[fix] lock del idle

repository/sandbox_lock.py

@@ -1,4 +1,5 @@
 import threading
+from dataclasses import dataclass
 from types import TracebackType
 from typing import Protocol
 from uuid import UUID
@@ -12,12 +13,25 @@ class _SyncLock(Protocol):
     def release(self) -> None: ...
 
 
+@dataclass(slots=True)
+class _LockEntry:
+    lock: _SyncLock
+    users: int = 0
+
+
 class _ChatLock(LockContext):
-    def __init__(self, lock: _SyncLock) -> None:
-        self._lock = lock
+    def __init__(
+        self,
+        locker: 'ProcessLocalSandboxLifecycleLocker',
+        chat_id: UUID,
+        entry: _LockEntry,
+    ) -> None:
+        self._locker = locker
+        self._chat_id = chat_id
+        self._entry = entry
 
     def __enter__(self) -> None:
-        self._lock.acquire()
+        self._entry.lock.acquire()
 
     def __exit__(
         self,
@@ -25,20 +39,32 @@ class _ChatLock(LockContext):
         exc: BaseException | None,
         traceback: TracebackType | None,
     ) -> bool | None:
-        self._lock.release()
+        self._entry.lock.release()
+        self._locker._release(self._chat_id, self._entry)
         return None
 
 
 class ProcessLocalSandboxLifecycleLocker(SandboxLifecycleLocker):
     def __init__(self) -> None:
         self._registry_lock = threading.Lock()
-        self._locks_by_chat_id: dict[UUID, _SyncLock] = {}
+        self._locks_by_chat_id: dict[UUID, _LockEntry] = {}
 
     def lock(self, chat_id: UUID) -> LockContext:
         with self._registry_lock:
-            lock = self._locks_by_chat_id.get(chat_id)
-            if lock is None:
-                lock = threading.Lock()
-                self._locks_by_chat_id[chat_id] = lock
+            entry = self._locks_by_chat_id.get(chat_id)
+            if entry is None:
+                entry = _LockEntry(lock=threading.Lock())
+                self._locks_by_chat_id[chat_id] = entry
+            entry.users += 1
 
-        return _ChatLock(lock)
+        return _ChatLock(self, chat_id, entry)
+
+    def _release(self, chat_id: UUID, entry: _LockEntry) -> None:
+        with self._registry_lock:
+            entry.users -= 1
+            if entry.users != 0:
+                return
+
+            current_entry = self._locks_by_chat_id.get(chat_id)
+            if current_entry is entry:
+                del self._locks_by_chat_id[chat_id]