deerflow2/backend/packages/harness/deerflow/sandbox/file_operation_lock.py

28 lines
997 B
Python

import threading
import weakref
from deerflow.sandbox.sandbox import Sandbox
# Use WeakValueDictionary to prevent memory leak in long-running processes.
# Locks are automatically removed when no longer referenced by any thread.
_LockKey = tuple[str, str]
_FILE_OPERATION_LOCKS: weakref.WeakValueDictionary[_LockKey, threading.Lock] = weakref.WeakValueDictionary()
_FILE_OPERATION_LOCKS_GUARD = threading.Lock()
def get_file_operation_lock_key(sandbox: Sandbox, path: str) -> tuple[str, str]:
sandbox_id = getattr(sandbox, "id", None)
if not sandbox_id:
sandbox_id = f"instance:{id(sandbox)}"
return sandbox_id, path
def get_file_operation_lock(sandbox: Sandbox, path: str) -> threading.Lock:
lock_key = get_file_operation_lock_key(sandbox, path)
with _FILE_OPERATION_LOCKS_GUARD:
lock = _FILE_OPERATION_LOCKS.get(lock_key)
if lock is None:
lock = threading.Lock()
_FILE_OPERATION_LOCKS[lock_key] = lock
return lock