ongoing fair queuing semaphore

Author: sanderegg

packages/service-library/src/servicelib/redis/fair_semaphore.py

@@ -0,0 +1,272 @@
+"""Fair semaphore decorator with automatic renewal and crash recovery."""
+
+import asyncio
+import datetime
+import functools
+import logging
+from collections.abc import Callable, Coroutine
+from contextlib import asynccontextmanager
+from typing import Any, ParamSpec, TypeVar
+
+from common_library.logging.logging_errors import create_troubleshooting_log_kwargs
+
+from ._constants import (
+    DEFAULT_EXPECTED_LOCK_OVERALL_TIME,
+    DEFAULT_SEMAPHORE_TTL,
+    DEFAULT_SOCKET_TIMEOUT,
+)
+from ._errors import (
+    SemaphoreAcquisitionError,
+    SemaphoreLostError,
+    SemaphoreNotAcquiredError,
+)
+from .fair_semaphore import FairSemaphore
+
+_logger = logging.getLogger(__name__)
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+@asynccontextmanager
+async def _managed_fair_semaphore_execution(
+    semaphore: FairSemaphore,
+    semaphore_key: str,
+    ttl: datetime.timedelta,
+    execution_context: str,
+    enable_auto_renewal: bool = True,
+):
+    """Context manager for fair semaphore with auto-renewal."""
+
+    async def _auto_renewal():
+        """Background task to automatically renew semaphore."""
+        if not enable_auto_renewal:
+            return
+
+        renewal_interval = ttl.total_seconds() / 3  # Renew at 1/3 TTL
+
+        while semaphore.acquired:
+            try:
+                await asyncio.sleep(renewal_interval)
+                if semaphore.acquired:  # Check again after sleep
+                    await semaphore.renew()
+                    _logger.debug(f"Renewed fair semaphore {semaphore_key}")
+            except SemaphoreLostError:
+                _logger.error(
+                    f"Fair semaphore {semaphore_key} was lost during execution"
+                )
+                break
+            except Exception as e:
+                _logger.warning(f"Failed to renew fair semaphore {semaphore_key}: {e}")
+                break
+
+    renewal_task = None
+    try:
+        # Acquire the semaphore (blocks until available)
+        if not await semaphore.acquire():
+            raise SemaphoreAcquisitionError(
+                f"Failed to acquire fair semaphore {semaphore_key}"
+            )
+
+        _logger.info(f"Acquired fair semaphore {semaphore_key} for {execution_context}")
+
+        # Start auto-renewal task if enabled
+        if enable_auto_renewal:
+            renewal_task = asyncio.create_task(_auto_renewal())
+
+        yield
+
+    except Exception as e:
+        _logger.error(
+            f"Error in fair semaphore-protected execution: {e}",
+            extra=create_troubleshooting_log_kwargs(
+                context=execution_context,
+                semaphore_key=semaphore_key,
+            ),
+        )
+        raise
+    finally:
+        # Cancel renewal task
+        if renewal_task and not renewal_task.done():
+            renewal_task.cancel()
+            try:
+                await renewal_task
+            except asyncio.CancelledError:
+                pass
+
+        # Release semaphore
+        if semaphore.acquired:
+            try:
+                await semaphore.release()
+                _logger.info(f"Released fair semaphore {semaphore_key}")
+            except Exception as e:
+                _logger.error(f"Failed to release fair semaphore {semaphore_key}: {e}")
+
+
+def fair_semaphore(
+    *,
+    key: str,
+    capacity: int,
+    ttl: datetime.timedelta = DEFAULT_SEMAPHORE_TTL,
+    timeout: datetime.timedelta = DEFAULT_SOCKET_TIMEOUT,
+    expected_execution_time: datetime.timedelta = DEFAULT_EXPECTED_LOCK_OVERALL_TIME,
+    cleanup_interval: datetime.timedelta = datetime.timedelta(seconds=30),
+    enable_auto_cleanup: bool = True,
+    enable_auto_renewal: bool = True,
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, R]]], Callable[P, Coroutine[Any, Any, R]]
+]:
+    """
+    Decorator that protects async functions with a fair distributed semaphore.
+
+    Uses Redis BRPOP for true FIFO fairness - first requester gets first slot.
+    No starvation possible, automatic crash recovery.
+
+    Args:
+        key: Unique semaphore identifier
+        capacity: Maximum concurrent executions allowed
+        ttl: How long each holder can keep the semaphore
+        timeout: How long to wait for semaphore (0 = infinite wait)
+        expected_execution_time: Expected total execution time (unused, kept for compatibility)
+        cleanup_interval: How often to run cleanup for crashed clients
+        enable_auto_cleanup: Whether to run background cleanup
+        enable_auto_renewal: Whether to automatically renew TTL during execution
+
+    Example:
+        @fair_semaphore(
+            key="api_calls",
+            capacity=10,
+            ttl=datetime.timedelta(seconds=30),
+            timeout=datetime.timedelta(seconds=60)
+        )
+        async def call_external_api():
+            # This will block fairly until semaphore available
+            # Maximum 10 concurrent executions
+            # First-come-first-served ordering guaranteed
+            pass
+    """
+
+    def decorator(
+        func: Callable[P, Coroutine[Any, Any, R]],
+    ) -> Callable[P, Coroutine[Any, Any, R]]:
+        @functools.wraps(func)
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            semaphore = FairSemaphore(
+                key=key,
+                capacity=capacity,
+                ttl=ttl,
+                timeout=timeout,
+                cleanup_interval=cleanup_interval,
+                enable_auto_cleanup=enable_auto_cleanup,
+            )
+
+            execution_context = f"{func.__module__}.{func.__qualname__}"
+
+            async with _managed_fair_semaphore_execution(
+                semaphore=semaphore,
+                semaphore_key=key,
+                ttl=ttl,
+                execution_context=execution_context,
+                enable_auto_renewal=enable_auto_renewal,
+            ):
+                return await func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+class FairSemaphoreContext:
+    """Async context manager for manual fair semaphore control."""
+
+    def __init__(
+        self,
+        key: str,
+        capacity: int,
+        ttl: datetime.timedelta = DEFAULT_SEMAPHORE_TTL,
+        timeout: datetime.timedelta = DEFAULT_SOCKET_TIMEOUT,
+        cleanup_interval: datetime.timedelta = datetime.timedelta(seconds=30),
+        enable_auto_cleanup: bool = True,
+        enable_auto_renewal: bool = True,
+    ):
+        self.semaphore = FairSemaphore(
+            key=key,
+            capacity=capacity,
+            ttl=ttl,
+            timeout=timeout,
+            cleanup_interval=cleanup_interval,
+            enable_auto_cleanup=enable_auto_cleanup,
+        )
+        self.ttl = ttl
+        self.enable_auto_renewal = enable_auto_renewal
+        self._renewal_task: Optional[asyncio.Task] = None
+
+    async def __aenter__(self) -> FairSemaphore:
+        """Acquire semaphore and start auto-renewal."""
+        await self.semaphore.acquire()
+
+        # Start auto-renewal if enabled
+        if self.enable_auto_renewal:
+
+            async def _auto_renewal():
+                renewal_interval = self.ttl.total_seconds() / 3
+                while self.semaphore.acquired:
+                    try:
+                        await asyncio.sleep(renewal_interval)
+                        if self.semaphore.acquired:
+                            await self.semaphore.renew()
+                    except (SemaphoreLostError, SemaphoreNotAcquiredError):
+                        break
+                    except Exception as e:
+                        _logger.warning(f"Auto-renewal failed: {e}")
+
+            self._renewal_task = asyncio.create_task(_auto_renewal())
+
+        return self.semaphore
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Stop renewal and release semaphore."""
+        if self._renewal_task and not self._renewal_task.done():
+            self._renewal_task.cancel()
+            try:
+                await self._renewal_task
+            except asyncio.CancelledError:
+                pass
+
+        if self.semaphore.acquired:
+            await self.semaphore.release()
+
+
+# Convenience function for creating fair semaphore contexts
+def fair_semaphore_context(
+    key: str,
+    capacity: int,
+    ttl: datetime.timedelta = DEFAULT_SEMAPHORE_TTL,
+    timeout: datetime.timedelta = DEFAULT_SOCKET_TIMEOUT,
+    cleanup_interval: datetime.timedelta = datetime.timedelta(seconds=30),
+    enable_auto_cleanup: bool = True,
+    enable_auto_renewal: bool = True,
+) -> FairSemaphoreContext:
+    """
+    Create an async context manager for fair semaphore usage.
+
+    Example:
+        async with fair_semaphore_context(
+            "my_resource",
+            capacity=5,
+            timeout=datetime.timedelta(seconds=30)
+        ) as sem:
+            # Protected code here - guaranteed fair access
+            # sem is the FairSemaphore instance
+            stats = await sem.count()
+            print(f"Current holders: {stats['current_holders']}")
+    """
+    return FairSemaphoreContext(
+        key=key,
+        capacity=capacity,
+        ttl=ttl,
+        timeout=timeout,
+        cleanup_interval=cleanup_interval,
+        enable_auto_cleanup=enable_auto_cleanup,
+        enable_auto_renewal=enable_auto_renewal,
+    )

packages/service-library/src/servicelib/redis/fair_semaphore_decorator.py

@@ -0,0 +1,32 @@
+-- Fair distributed semaphore using token pool (BRPOP-based)
+-- KEYS[1]: tokens_key (LIST of available tokens)
+-- KEYS[2]: holders_key (SET of current holder instance IDs)
+-- KEYS[3]: holder_key (individual holder TTL key for this instance)
+-- ARGV[1]: instance_id
+-- ARGV[2]: capacity (max concurrent holders)
+-- ARGV[3]: ttl_seconds
+-- ARGV[4]: timeout_seconds (for BRPOP)
+--
+-- Returns: {exit_code, status, token, current_count}
+-- exit_code: 0 if acquired, 255 if timeout/failed
+-- status: 'acquired' or 'timeout'
+
+local holders_key = KEYS[1]
+local holder_key = KEYS[2]
+
+local token = ARGV[1]
+local instance_id = ARGV[2]
+local ttl_seconds = tonumber(ARGV[3])
+
+
+
+-- Step 1: Register as holder
+redis.call('SADD', holders_key, instance_id)
+redis.call('SETEX', holder_key, ttl_seconds, token)
+
+-- Step 2: Set expiry on holders set to prevent infinite growth
+redis.call('EXPIRE', holders_key, ttl_seconds * 10)
+
+local current_count = redis.call('SCARD', holders_key)
+
+return {0, 'acquired', token, current_count}

packages/service-library/src/servicelib/redis/lua/acquire_fair_semaphore_v2.lua

@@ -0,0 +1,66 @@
+-- Cleanup orphaned tokens from crashed clients
+-- KEYS[1]: tokens_key (LIST of available tokens)
+-- KEYS[2]: holders_key (SET of current holders)
+-- KEYS[3]: holder_prefix (prefix for holder keys, e.g. "semaphores:holders:key:")
+-- ARGV[1]: capacity (total semaphore capacity)
+--
+-- Returns: {recovered_tokens, current_holders, available_tokens, total_cleaned}
+-- This script should be run periodically to recover tokens from crashed clients
+
+local tokens_key = KEYS[1]
+local holders_key = KEYS[2]
+local holder_prefix = KEYS[3]
+
+local capacity = tonumber(ARGV[1])
+
+-- Step 1: Get all current holders
+local current_holders = redis.call('SMEMBERS', holders_key)
+local recovered_tokens = 0
+local cleaned_holders = {}
+
+-- Step 2: Check each holder to see if their TTL key still exists
+for i = 1, #current_holders do
+    local holder_id = current_holders[i]
+    local holder_key = holder_prefix .. holder_id
+    local exists = redis.call('EXISTS', holder_key)
+
+    if exists == 0 then
+        -- Holder key doesn't exist but holder is in SET
+        -- This indicates a crashed client - clean up and recover token
+        redis.call('SREM', holders_key, holder_id)
+        redis.call('LPUSH', tokens_key, 'token_recovered_' .. holder_id)
+        recovered_tokens = recovered_tokens + 1
+        table.insert(cleaned_holders, holder_id)
+    end
+end
+
+-- Step 3: Ensure we have the correct total number of tokens
+local remaining_holders = redis.call('SCARD', holders_key)
+local available_tokens_count = redis.call('LLEN', tokens_key)
+local total_tokens = remaining_holders + available_tokens_count
+
+-- If we're missing tokens (due to crashes or Redis issues), add them back
+local missing_tokens = capacity - total_tokens
+for i = 1, missing_tokens do
+    redis.call('LPUSH', tokens_key, 'token_missing_' .. i)
+    recovered_tokens = recovered_tokens + 1
+end
+
+-- If we somehow have too many tokens (shouldn't happen), remove extras
+local excess_tokens = total_tokens - capacity
+for i = 1, excess_tokens do
+    redis.call('RPOP', tokens_key)
+end
+
+-- Step 4: Refresh expiry on data structures to prevent cleanup
+local final_holders = redis.call('SCARD', holders_key)
+local final_available = redis.call('LLEN', tokens_key)
+
+if final_holders > 0 then
+    redis.call('EXPIRE', holders_key, 3600)  -- 1 hour expiry
+end
+if final_available > 0 then
+    redis.call('EXPIRE', tokens_key, 3600)   -- 1 hour expiry
+end
+
+return {recovered_tokens, final_holders, final_available, #cleaned_holders}

packages/service-library/src/servicelib/redis/lua/cleanup_fair_semaphore_v2.lua

@@ -0,0 +1,17 @@
+-- Count current semaphore holders (simplified for token pool design)
+-- KEYS[1]: holders_key (SET of current holders)
+-- KEYS[2]: tokens_key (LIST of available tokens)
+-- ARGV[1]: capacity (total semaphore capacity)
+--
+-- Returns: {current_holders, available_tokens, total_capacity}
+
+local holders_key = KEYS[1]
+local tokens_key = KEYS[2]
+
+local capacity = tonumber(ARGV[1])
+
+-- Count current holders and available tokens
+local current_holders = redis.call('SCARD', holders_key)
+local available_tokens = redis.call('LLEN', tokens_key)
+
+return {current_holders, available_tokens, capacity}