redis-allocator 0.0.1__py3-none-any.whl → 0.3.2__py3-none-any.whl

redis_allocator/lock.py

@@ -207,7 +207,7 @@ class BaseLock(ABC):
         """Deletes a key when the comparison value is not equal to the current value."""
         return self._conditional_setdel('!=', key, value, None, None, True)
 
-    def _to_seconds(self, timeout: Timeout):
+    def _to_seconds(self, timeout: Timeout) -> float:
         """Convert a timeout to seconds."""
         if timeout is None:
             timeout = datetime(2099, 1, 1).timestamp()
@@ -288,13 +288,14 @@ class BaseLockPool(BaseLock, metaclass=ABCMeta):
 class RedisLock(BaseLock):
     """Redis-based lock implementation.
 
-    Provides distributed locking capabilities using Redis as the backend storage.
+    Uses standard Redis commands (SET with NX, EX options) for basic locking
+    and Lua scripts for conditional operations (set/del based on value comparison).
 
     Attributes:
-        redis: The Redis client instance.
-        prefix: Prefix for Redis keys.
-        suffix: Suffix for Redis keys.
-        eps: Epsilon value for floating point comparison.
+        redis: StrictRedis client instance (must decode responses).
+        prefix: Prefix for all Redis keys managed by this lock instance.
+        suffix: Suffix for Redis keys to distinguish lock types (e.g., 'lock').
+        eps: Epsilon for float comparisons in conditional Lua scripts.
     """
 
     redis: Redis
@@ -319,6 +320,10 @@ class RedisLock(BaseLock):
 
     @property
     def _lua_required_string(self):
+        """Base Lua script providing the key_str function.
+
+        - key_str(key: str): Constructs the full Redis key using prefix and suffix.
+        """
         return f'''
         local function key_str(key)
             return '{self.prefix}|{self.suffix}:' .. key
@@ -361,6 +366,11 @@ class RedisLock(BaseLock):
 
     def _conditional_setdel(self, op: str, key: str, value: float, set_value: Optional[float] = None, ex: Optional[int] = None,
                             isdel: bool = False) -> bool:
+        """Executes the conditional set/delete Lua script.
+
+        Passes necessary arguments (key, compare_value, set_value, expiry, isdel flag)
+        to the cached Lua script corresponding to the comparison operator (`op`).
+        """
         # Convert None to a valid value for Redis (using -1 to indicate no expiration)
         key_value = self._key_str(key)
         ex_value = -1 if ex is None else ex
@@ -376,6 +386,21 @@ class RedisLock(BaseLock):
                 ('>', '<', '>=', '<=', '==', '!=')}
 
     def _conditional_setdel_lua_script(self, op: str, eps: float = 1e-6) -> str:
+        """Generates the Lua script for conditional set/delete operations.
+
+        Args:
+            op: The comparison operator ('>', '<', '>=', '<=', '==', '!=').
+            eps: Epsilon for floating-point comparisons ('==', '!=', '>=', '<=').
+
+        Returns:
+            A Lua script string that:
+            1. Gets the current numeric value of the target key (KEYS[1]).
+            2. Compares it with the provided compare_value (ARGV[1]) using the specified `op`.
+            3. If the key doesn't exist or the condition is true:
+               - If `isdel` (ARGV[4]) is true, deletes the key.
+               - Otherwise, sets the key to `new_value` (ARGV[2]) with optional expiry `ex` (ARGV[3]).
+            4. Returns true if the operation was performed, false otherwise.
+        """
         match op:
             case '>':
                 condition = 'compare_value > current_value'
@@ -425,9 +450,14 @@ class RedisLock(BaseLock):
 
 
 class RedisLockPool(RedisLock, BaseLockPool):
-    """Redis-based lock pool implementation.
+    """Manages a collection of RedisLock keys as a logical pool.
+
+    Uses a Redis Set (`<prefix>|<suffix>|pool`) to store the identifiers (keys)
+    belonging to the pool. Inherits locking logic from RedisLock.
 
-    Manages a pool of Redis locks, stored as a Redis set.
+    Provides methods to add (`extend`), remove (`shrink`), replace (`assign`),
+    and query (`keys`, `__contains__`) the members of the pool.
+    Also offers methods to check the lock status of pool members (`_get_key_lock_status`).
     """
 
     def __init__(self, redis: Redis, prefix: str, suffix='lock-pool', eps: float = 1e-6):
@@ -444,6 +474,11 @@ class RedisLockPool(RedisLock, BaseLockPool):
 
     @property
     def _lua_required_string(self):
+        """Base Lua script providing key_str and pool_str functions.
+
+        - key_str(key: str): Inherited from RedisLock.
+        - pool_str(): Returns the Redis key for the pool Set.
+        """
         return f'''
         {super()._lua_required_string}
         local function pool_str()
@@ -467,10 +502,16 @@ class RedisLockPool(RedisLock, BaseLockPool):
 
     @property
     def _assign_lua_string(self):
+        """Lua script to atomically replace the contents of the pool Set.
+
+        1. Deletes the existing pool Set key (KEYS[1]).
+        2. Adds all provided keys (ARGV) to the (now empty) pool Set using SADD.
+        """
         return f'''
         {self._lua_required_string}
-        redis.call('DEL', pool_str())
-        redis.call('SADD', pool_str(), unpack(ARGV))
+        local _pool_str = KEYS[1]
+        redis.call('DEL', _pool_str)
+        redis.call('SADD', _pool_str, unpack(ARGV))
         '''
 
     @cached_property
@@ -480,7 +521,7 @@ class RedisLockPool(RedisLock, BaseLockPool):
     def assign(self, keys: Optional[Sequence[str]] = None):
         """Assign keys to the pool, replacing any existing keys."""
         if keys is not None and len(keys) > 0:
-            self._assign_lua_script(args=keys)
+            self._assign_lua_script(args=keys, keys=[self._pool_str()])
         else:
             self.clear()
 
@@ -514,13 +555,16 @@ class LockData:
 
 
 class ThreadLock(BaseLock):
-    """Thread-safe lock implementation for local process use.
+    """In-memory, thread-safe lock implementation conforming to BaseLock.
 
-    Provides similar functionality to RedisLock but works locally
-    within a single Python process using threading mechanisms.
+    Simulates Redis lock behavior using Python's `threading.RLock` for concurrency
+    control and a `defaultdict` to store lock data (value and expiry timestamp).
+    Suitable for single-process scenarios or testing.
 
     Attributes:
-        eps: Epsilon value for floating point comparison.
+        eps: Epsilon for float comparisons.
+        _locks: defaultdict storing LockData(value, expiry) for each key.
+        _lock: threading.RLock protecting access to _locks.
     """
 
     def __init__(self, eps: float = 1e-6):
@@ -636,9 +680,14 @@ class ThreadLock(BaseLock):
 
 
 class ThreadLockPool(ThreadLock, BaseLockPool):
-    """Thread-safe lock pool implementation for local process use.
+    """In-memory, thread-safe lock pool implementation.
 
-    Maintains a set of keys as the pool and provides operations to manage them.
+    Manages a collection of lock keys using a Python `set` for the pool members
+    and inherits the locking logic from `ThreadLock`.
+
+    Attributes:
+        _pool: Set containing the keys belonging to this pool.
+        _lock: threading.RLock protecting access to _locks and _pool.
     """
 
     def __init__(self, eps: float = 1e-6):

redis_allocator/task_queue.py

@@ -22,7 +22,7 @@ logger = logging.getLogger(__name__)
 
 class TaskExecutePolicy(Enum):
     """Defines different policies for task execution.
-    
+
     Attributes:
         Local: Execute task only locally
         Remote: Execute task only remotely
@@ -39,22 +39,24 @@ class TaskExecutePolicy(Enum):
 
 @dataclass
 class RedisTask:
-    """Represents a task in the Redis task queue system.
-    
-    This class encapsulates all information related to a task, including
-    its unique ID, name, parameters, and execution status.
-    
+    """Represents a task to be processed via the RedisTaskQueue.
+
+    Encapsulates task details like ID, category (name), parameters, and its current state
+    (progress, result, error). It includes methods for saving state and updating progress.
+
     Attributes:
-        id: Unique identifier for the task
-        name: Name of the task category
-        params: Dictionary of parameters for the task
-        expiry: Unix timestamp when this task expires
-        result: The result of the task, initially None
-        error: Any error that occurred during task execution
-        update_progress_time: Last time the progress was updated
-        current_progress: Current progress value
-        total_progress: Total progress value for completion
-        _save: Internal callable to save task state to Redis
+        id: Unique identifier for this specific task instance.
+        name: Categorical name for the task (used for routing in the queue).
+        params: Dictionary containing task-specific input parameters.
+        expiry: Absolute Unix timestamp when the task should be considered expired.
+               Used both locally and remotely to timeout waiting operations.
+        result: Stores the successful return value of the task execution.
+        error: Stores any exception raised during task execution.
+        update_progress_time: Timestamp of the last progress update.
+        current_progress: Current progress value (e.g., items processed).
+        total_progress: Total expected value for completion (e.g., total items).
+        _save: Internal callback function provided by RedisTaskQueue to persist
+               the task's state (result, error, progress) back to Redis.
     """
     id: str
     name: str
@@ -74,11 +76,11 @@ class RedisTask:
 
     def update(self, current_progress: float, total_progress: float):
         """Update the progress of the task.
-        
+
         Args:
             current_progress: The current progress value
             total_progress: The total progress value for completion
-            
+
         Raises:
             TimeoutError: If the task has expired
         """
@@ -91,14 +93,36 @@ class RedisTask:
 
 
 class RedisTaskQueue:
-    """A class that provides a simple interface for managing a task queue in Redis.
+    """Provides a distributed task queue using Redis lists and key/value storage.
+
+    Enables submitting tasks (represented by `RedisTask` objects) to named queues
+    and having them processed either locally (if `task_fn` is provided) or by
+    remote listeners polling the corresponding Redis list.
 
-    This class enables distributed task processing through Redis, with support for
-    asynchronous processing, task listening, and result retrieval. Tasks can be 
-    executed locally or remotely based on configurable policies.
+    Key Concepts:
+    - Task Queues: Redis lists (`<prefix>|<suffix>|task-queue|<name>`) where task IDs
+                   are pushed for remote workers.
+    - Task Data: Serialized `RedisTask` objects stored in Redis keys
+                 (`<prefix>|<suffix>|task-result:<id>`) with an expiry time.
+                 This stores parameters, progress, results, and errors.
+    - Listeners: Remote workers use BLPOP on queue lists. To signal their presence,
+                 they periodically set a listener key (`<prefix>|<suffix>|task-listen|<name>`).
+    - Execution Policies (`TaskExecutePolicy`): Control whether a task is executed
+                                            locally, remotely, or attempts one then the other.
+                                            `Auto` mode checks for a listener key.
+    - Task Function (`task_fn`): A user-provided function that takes a `RedisTask`
+                                 and performs the actual work locally.
+
+    Attributes:
+        redis: StrictRedis client instance.
+        prefix: Prefix for all Redis keys.
+        suffix: Suffix for Redis keys (default: 'task-queue').
+        timeout: Default expiry/timeout for tasks and listener keys (seconds).
+        interval: Polling interval for remote task fetching (seconds).
+        task_fn: Callable[[RedisTask], Any] to execute tasks locally.
     """
 
-    def __init__(self, redis: Redis, prefix: str, suffix='task-queue', timeout=300, interval=5, 
+    def __init__(self, redis: Redis, prefix: str, suffix='task-queue', timeout=300, interval=5,
                  task_fn: Callable[[RedisTask], Any] = None):
         """Initialize a RedisTaskQueue instance.
 
@@ -119,12 +143,12 @@ class RedisTaskQueue:
 
     def build_task(self, id: str, name: str, params: dict) -> RedisTask:
         """Create a new RedisTask instance with the given parameters.
-        
+
         Args:
             id: Unique identifier for the task
             name: Name of the task category
             params: Dictionary of parameters for the task
-            
+
         Returns:
             A new RedisTask instance with a save function configured
         """
@@ -134,15 +158,15 @@ class RedisTaskQueue:
 
     def execute_task_remotely(self, task: RedisTask, timeout: Optional[float] = None, once: bool = False) -> Any:
         """Execute a task remotely by pushing it to the queue.
-        
+
         Args:
             task: The RedisTask to execute
             timeout: Optional timeout in seconds, defaults to self.timeout
             once: Whether to delete the result after getting it
-            
+
         Returns:
             The result of the task
-            
+
         Raises:
             TimeoutError: If the task times out
             Exception: Any exception raised during task execution
@@ -164,14 +188,14 @@ class RedisTaskQueue:
 
     def execute_task_locally(self, task: RedisTask, timeout: Optional[float] = None) -> Any:
         """Execute a task locally using the task_fn.
-        
+
         Args:
             task: The RedisTask to execute
             timeout: Optional timeout in seconds, updates task.expiry if provided
-            
+
         Returns:
             The result of the task
-            
+
         Raises:
             Exception: Any exception raised during task execution
         """
@@ -189,7 +213,7 @@ class RedisTaskQueue:
     @cached_property
     def _queue_prefix(self) -> str:
         """Get the prefix for queue keys.
-        
+
         Returns:
             The queue prefix
         """
@@ -200,7 +224,7 @@ class RedisTaskQueue:
 
         Args:
             name: The task name.
-            
+
         Returns:
             The formatted queue key.
         """
@@ -208,13 +232,13 @@ class RedisTaskQueue:
 
     def _queue_name(self, key: str) -> str:
         """Extract the queue name from a queue key.
-        
+
         Args:
             key: The queue key.
-            
+
         Returns:
             The queue name.
-            
+
         Raises:
             AssertionError: If the key doesn't start with the queue prefix.
         """
@@ -223,10 +247,10 @@ class RedisTaskQueue:
 
     def _queue_listen_name(self, name: str) -> str:
         """Generate a listen name for the given task name.
-        
+
         Args:
             name: The task name.
-            
+
         Returns:
             The formatted listen name.
         """
@@ -242,10 +266,10 @@ class RedisTaskQueue:
 
     def _result_key(self, task_id: str) -> str:
         """Generate a result key for the given task ID.
-        
+
         Args:
             task_id: The task ID.
-            
+
         Returns:
             The formatted result key.
         """
@@ -253,10 +277,10 @@ class RedisTaskQueue:
 
     def set_task(self, task: RedisTask) -> str:
         """Save a task to Redis.
-        
+
         Args:
             task: The RedisTask to save.
-            
+
         Returns:
             The task ID.
         """
@@ -268,11 +292,11 @@ class RedisTaskQueue:
 
     def get_task(self, task_id: str, once: bool = False) -> Optional[RedisTask]:
         """Get a task from Redis.
-        
+
         Args:
             task_id: The task ID.
             once: Whether to delete the task after getting it.
-            
+
         Returns:
             The RedisTask, or None if no task is available.
         """
@@ -286,10 +310,10 @@ class RedisTaskQueue:
 
     def has_task(self, task_id: str) -> bool:
         """Check if a task exists.
-        
+
         Args:
             task_id: The task ID.
-            
+
         Returns:
             True if the task exists, False otherwise.
         """
@@ -298,12 +322,12 @@ class RedisTaskQueue:
     @contextmanager
     def _executor_context(self, max_workers: int = 128):
         """Create a ThreadPoolExecutor context manager.
-        
+
         This is a helper method for testing and internal use.
-        
+
         Args:
             max_workers: The maximum number of worker threads.
-            
+
         Yields:
             The ThreadPoolExecutor instance.
         """
@@ -312,10 +336,10 @@ class RedisTaskQueue:
 
     def listen(self, names: List[str], workers: int = 128, event: Optional[Event] = None) -> None:
         """Listen for tasks on the specified queues.
-        
-        This method continuously polls the specified queues for tasks, 
+
+        This method continuously polls the specified queues for tasks,
         and executes tasks locally when they are received.
-        
+
         Args:
             names: The names of the queues to listen to.
             workers: The number of worker threads to use. Default is 128.
@@ -339,10 +363,10 @@ class RedisTaskQueue:
     def query(self, id: str, name: str, params: dict, timeout: Optional[float] = None,
               policy: TaskExecutePolicy = TaskExecutePolicy.Auto, once: bool = False) -> Any:
         """Execute a task according to the specified policy.
-        
+
         This method provides a flexible way to execute tasks with different
         strategies based on the specified policy.
-        
+
         Args:
             id: The task ID.
             name: The task name.
@@ -350,10 +374,10 @@ class RedisTaskQueue:
             timeout: Optional timeout override.
             policy: The execution policy to use.
             once: Whether to delete the result after getting it.
-            
+
         Returns:
             The result of the task.
-            
+
         Raises:
             Exception: Any exception raised during task execution.
         """
@@ -379,4 +403,4 @@ class RedisTaskQueue:
                 if self.redis.exists(self._queue_listen_name(name)):
                     return self.execute_task_remotely(t, timeout)
                 else:
-                    return self.execute_task_locally(t, timeout)
+                    return self.execute_task_locally(t, timeout)