feat: implement advisory locking in Channel base class and enhance serialization support

Author: myui

graflow/channels/base.py

@@ -2,16 +2,32 @@
 
 from __future__ import annotations
 
+import threading
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
-from typing import Any, Iterator, List, Optional, Union
+from typing import Any, Dict, Iterator, List, Optional, Union
 
 
 class Channel(ABC):
     """Abstract base class for all channel implementations."""
 
     def __init__(self, name: str):
         self.name = name
+        self._key_locks: Dict[str, threading.RLock] = {}
+        self._key_locks_guard = threading.Lock()  # protects _key_locks dict itself
+
+    def __getstate__(self) -> Dict[str, Any]:
+        """Exclude unpicklable lock objects during serialization."""
+        state = self.__dict__.copy()
+        state.pop("_key_locks", None)
+        state.pop("_key_locks_guard", None)
+        return state
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        """Recreate lock objects after deserialization."""
+        self.__dict__.update(state)  # type: ignore[arg-type]
+        self._key_locks = {}
+        self._key_locks_guard = threading.Lock()
 
     @abstractmethod
     def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
@@ -90,9 +106,20 @@ def atomic_add(self, key: str, amount: Union[int, float] = 1) -> Union[int, floa
         """
         pass
 
+    # -- advisory locking for compound operations --
+
+    def _get_key_lock(self, key: str) -> threading.RLock:
+        """Return (or create) the ``RLock`` associated with *key*."""
+        if key not in self._key_locks:
+            with self._key_locks_guard:
+                # Double-checked locking
+                if key not in self._key_locks:
+                    self._key_locks[key] = threading.RLock()
+        return self._key_locks[key]
+
     @contextmanager
     def lock(self, key: str, timeout: float = 10.0) -> Iterator[None]:
-        """Acquire an advisory lock scoped to *key* for compound operations.
+        """Acquire an advisory per-key lock for compound read-modify-write.
 
         Usage::
 
@@ -105,22 +132,22 @@ def lock(self, key: str, timeout: float = 10.0) -> Iterator[None]:
         protect read-modify-write sequences that cannot be expressed with
         ``atomic_add()``.
 
-        .. warning::
-
-            The default implementation is a **no-op** (yields immediately)
-            and provides **no mutual exclusion**.  Subclasses that need
-            compound-operation safety **must** override this method — this
-            includes both in-process backends under multi-threading
-            (e.g. ``MemoryChannel``) and distributed backends where
-            multi-client read-modify-write sequences are racy
-            (e.g. Redis without a distributed lock).
-
         Args:
             key: Logical key to lock on (does not need to correspond to a
                  stored key).
             timeout: Maximum seconds to wait for the lock.
 
         Yields:
             None — the lock is held for the duration of the ``with`` block.
+
+        Raises:
+            TimeoutError: If the lock cannot be acquired within *timeout*.
         """
-        yield
+        rlock = self._get_key_lock(key)
+        acquired = rlock.acquire(timeout=timeout)
+        if not acquired:
+            raise TimeoutError(f"Could not acquire lock for key '{key}' within {timeout}s")
+        try:
+            yield
+        finally:
+            rlock.release()

graflow/channels/memory_channel.py

@@ -2,10 +2,8 @@
 
 from __future__ import annotations
 
-import threading
 import time
-from contextlib import contextmanager
-from typing import Any, Dict, Iterator, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from graflow.channels.base import Channel
 
@@ -18,21 +16,6 @@ def __init__(self, name: str, **kwargs):
         super().__init__(name)
         self.data: Dict[str, Any] = {}
         self.ttl_data: Dict[str, float] = {}
-        self._key_locks: Dict[str, threading.RLock] = {}
-        self._key_locks_guard = threading.Lock()  # protects _key_locks dict itself
-
-    def __getstate__(self) -> Dict[str, Any]:
-        """Exclude unpicklable lock objects during serialization."""
-        state = self.__dict__.copy()
-        state.pop("_key_locks", None)
-        state.pop("_key_locks_guard", None)
-        return state
-
-    def __setstate__(self, state: Dict[str, Any]) -> None:
-        """Recreate lock objects after deserialization."""
-        self.__dict__.update(state)  # type: ignore[assignment]
-        self._key_locks = {}
-        self._key_locks_guard = threading.Lock()
 
     def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
         """Store data in the channel."""
@@ -164,34 +147,3 @@ def atomic_add(self, key: str, amount: Union[int, float] = 1) -> Union[int, floa
             new_value = current + amount
             self.data[key] = new_value
             return new_value
-
-    # -- advisory locking for compound operations --
-
-    def _get_key_lock(self, key: str) -> threading.RLock:
-        """Return (or create) the RLock associated with *key*."""
-        if key not in self._key_locks:
-            with self._key_locks_guard:
-                # Double-checked locking
-                if key not in self._key_locks:
-                    self._key_locks[key] = threading.RLock()
-        return self._key_locks[key]
-
-    @contextmanager
-    def lock(self, key: str, timeout: float = 10.0) -> Iterator[None]:
-        """Acquire an advisory per-key lock for compound read-modify-write.
-
-        Args:
-            key: Logical key to lock on.
-            timeout: Maximum seconds to wait for the lock.
-
-        Raises:
-            TimeoutError: If the lock cannot be acquired within *timeout*.
-        """
-        rlock = self._get_key_lock(key)
-        acquired = rlock.acquire(timeout=timeout)
-        if not acquired:
-            raise TimeoutError(f"Could not acquire lock for key '{key}' within {timeout}s")
-        try:
-            yield
-        finally:
-            rlock.release()

graflow/channels/redis_channel.py

@@ -240,15 +240,13 @@ def atomic_add(self, key: str, amount: Union[int, float] = 1) -> Union[int, floa
 
     def __getstate__(self):
         """Support for pickle serialization."""
-        state = self.__dict__.copy()
-        # Remove the unpicklable Redis client
+        state = super().__getstate__()
         del state["redis_client"]
         return state
 
     def __setstate__(self, state):
         """Support for pickle deserialization."""
-        self.__dict__.update(state)
-        # Recreate the Redis client
+        super().__setstate__(state)
         assert redis is not None, "redis package is required for RedisChannel"
         self.redis_client = redis.Redis(
             host=self._host, port=self._port, db=self._db, decode_responses=True, **self._kwargs