passing tests

Author: sanderegg

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

@@ -4,6 +4,7 @@
 from types import TracebackType
 from typing import Annotated, ClassVar
 
+import redis.exceptions
 from common_library.basic_types import DEFAULT_FACTORY
 from pydantic import (
     BaseModel,
@@ -13,6 +14,7 @@
     field_validator,
 )
 from redis.commands.core import AsyncScript
+from servicelib.redis._utils import handle_redis_returns_union_types
 
 from ._client import RedisClientSDK
 from ._constants import (
@@ -29,6 +31,7 @@
 from ._semaphore_lua import (
     ACQUIRE_FAIR_SEMAPHORE_V2_SCRIPT,
     COUNT_FAIR_SEMAPHORE_V2_SCRIPT,
+    REGISTER_FAIR_SEMAPHORE_SCRIPT,
     RELEASE_FAIR_SEMAPHORE_V2_SCRIPT,
     RENEW_FAIR_SEMAPHORE_V2_SCRIPT,
     SCRIPT_BAD_EXIT_CODE,
@@ -92,6 +95,7 @@ class DistributedSemaphore(BaseModel):
     ] = DEFAULT_FACTORY
 
     # Class and/or Private state attributes (not part of the model)
+    register_semaphore: ClassVar[AsyncScript | None] = None
     acquire_script: ClassVar[AsyncScript | None] = None
     count_script: ClassVar[AsyncScript | None] = None
     release_script: ClassVar[AsyncScript | None] = None
@@ -104,6 +108,9 @@ def _register_scripts(cls, redis_client: RedisClientSDK) -> None:
         caches the script SHA, so this is efficient. Even if called multiple times,
         the script is only registered once."""
         if cls.acquire_script is None:
+            cls.register_semaphore = redis_client.redis.register_script(
+                REGISTER_FAIR_SEMAPHORE_SCRIPT
+            )
             cls.acquire_script = redis_client.redis.register_script(
                 ACQUIRE_FAIR_SEMAPHORE_V2_SCRIPT
             )
@@ -145,6 +152,12 @@ def holder_key(self) -> str:
         """Redis key for this instance's holder entry."""
         return f"{SEMAPHORE_HOLDER_KEY_PREFIX}{self.key}:{self.instance_id}"
 
+    @computed_field
+    @property
+    def holder_prefix(self) -> str:
+        """Prefix for holder keys (used in cleanup)."""
+        return f"{SEMAPHORE_HOLDER_KEY_PREFIX}{self.key}:"
+
     # Additional validation
     @field_validator("ttl")
     @classmethod
@@ -176,22 +189,51 @@ async def acquire(self) -> bool:
         """
 
         ttl_seconds = int(self.ttl.total_seconds())
-        blocking_timeout_seconds = 0.1
+        blocking_timeout_seconds = 1
         if self.blocking:
-            blocking_timeout_seconds = (
+            blocking_timeout_seconds = int(
                 self.blocking_timeout.total_seconds() if self.blocking_timeout else 0
             )
 
-        # Execute the Lua script atomically
+        # Execute the Lua scripts atomically
         cls = type(self)
+        assert cls.register_semaphore is not None  # nosec
+        await cls.register_semaphore(
+            keys=[self.tokens_key, self.holders_key],
+            args=[self.capacity, ttl_seconds],
+            client=self.redis_client.redis,
+        )  # pylint: disable=not-callable
+
+        try:
+            # this is blocking pop with timeout
+            tokens_key_token: list[str] = await handle_redis_returns_union_types(
+                self.redis_client.redis.brpop(
+                    [self.tokens_key], timeout=blocking_timeout_seconds
+                )
+            )
+        except redis.exceptions.TimeoutError as e:
+            _logger.debug(
+                "Timeout acquiring semaphore '%s' (instance: %s)",
+                self.key,
+                self.instance_id,
+            )
+            if self.blocking:
+                raise SemaphoreAcquisitionError(
+                    name=self.key, capacity=self.capacity
+                ) from e
+            return False
+
+        assert len(tokens_key_token) == 2  # nosec
+        assert tokens_key_token[0] == self.tokens_key  # nosec
+        token = tokens_key_token[1]
+
         assert cls.acquire_script is not None  # nosec
         result = await cls.acquire_script(  # pylint: disable=not-callable
-            keys=[self.tokens_key, self.holders_key, self.holder_key],
+            keys=[self.holders_key, self.holder_key],
             args=[
+                token[0],
                 self.instance_id,
-                self.capacity,
                 ttl_seconds,
-                blocking_timeout_seconds,
             ],
             client=self.redis_client.redis,
         )
@@ -209,6 +251,7 @@ async def acquire(self) -> bool:
                 current_count,
             )
             return True
+
         if status == "timeout":
             if self.blocking:
                 _logger.debug(
@@ -261,6 +304,7 @@ async def release(self) -> None:
                 self.instance_id,
                 current_count,
             )
+            return
 
         # Instance was already expired or not acquired
         assert exit_code == SCRIPT_BAD_EXIT_CODE  # nosec

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

@@ -33,6 +33,7 @@ def _load_script(script_name: str) -> str:
 COUNT_SEMAPHORE_SCRIPT: Final[str] = _load_script("count_semaphore")
 
 # fair semaphore scripts (token pool based)
+REGISTER_FAIR_SEMAPHORE_SCRIPT: Final[str] = _load_script("register_semaphore_holder")
 ACQUIRE_FAIR_SEMAPHORE_V2_SCRIPT: Final[str] = _load_script("acquire_fair_semaphore_v2")
 RELEASE_FAIR_SEMAPHORE_V2_SCRIPT: Final[str] = _load_script("release_fair_semaphore_v2")
 CLEANUP_FAIR_SEMAPHORE_V2_SCRIPT: Final[str] = _load_script("cleanup_fair_semaphore_v2")

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

@@ -11,43 +11,20 @@
 -- exit_code: 0 if acquired, 255 if timeout/failed
 -- status: 'acquired' or 'timeout'
 
-local tokens_key = KEYS[1]
-local holders_key = KEYS[2]
-local holder_key = KEYS[3]
+local holders_key = KEYS[1]
+local holder_key = KEYS[2]
 
-local instance_id = ARGV[1]
-local capacity = tonumber(ARGV[2])
+local token = ARGV[1]
+local instance_id = ARGV[2]
 local ttl_seconds = tonumber(ARGV[3])
-local timeout_seconds = tonumber(ARGV[4])
-
--- Step 1: Initialize token pool if needed (first time setup)
-local tokens_exist = redis.call('EXISTS', tokens_key)
-if tokens_exist == 0 then
-    -- Initialize with capacity number of tokens
-    for i = 1, capacity do
-        redis.call('LPUSH', tokens_key, 'token_' .. i)
-    end
-    -- Set expiry on tokens list to prevent infinite growth
-    redis.call('EXPIRE', tokens_key, ttl_seconds * 10)
-end
-
--- Step 2: Try to get a token using blocking pop
--- timeout_seconds = 0 means block indefinitely
-local token_result = redis.call('BRPOP', tokens_key, timeout_seconds)
-
-if token_result == false or token_result == nil then
-    -- Timeout occurred
-    local current_count = redis.call('SCARD', holders_key)
-    return {255, 'timeout', '', current_count}
-end
-
-local token = token_result[2]  -- BRPOP returns {key, value}
-
--- Step 3: Register as holder
+
+
+
+-- Step 1: Register as holder
 redis.call('SADD', holders_key, instance_id)
 redis.call('SETEX', holder_key, ttl_seconds, token)
 
--- Step 4: Set expiry on holders set to prevent infinite growth
+-- 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)

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

@@ -12,31 +12,20 @@
 
 local tokens_key = KEYS[1]
 local holders_key = KEYS[2]
-local holder_key = KEYS[3]
 
-local instance_id = ARGV[1]
-local capacity = tonumber(ARGV[2])
-local ttl_seconds = tonumber(ARGV[3])
-local token = ARGV[4]
+local capacity = tonumber(ARGV[1])
+local ttl_seconds = tonumber(ARGV[2])
 
 -- Step 1: Initialize token pool if needed (first time setup)
 local tokens_exist = redis.call('EXISTS', tokens_key)
-if tokens_exist == 0 then
+local holders_exist = redis.call('EXISTS', holders_key)
+if tokens_exist == 0 and holders_exist == 0 then
     -- Initialize with capacity number of tokens
     for i = 1, capacity do
         redis.call('LPUSH', tokens_key, 'token_' .. i)
     end
     -- Set expiry on tokens list to prevent infinite growth
-    redis.call('EXPIRE', tokens_key, ttl_seconds * 10)
+    -- redis.call('EXPIRE', tokens_key, ttl_seconds)
 end
 
--- Step 2: Register as holder (token was already popped by Python BRPOP)
-redis.call('SADD', holders_key, instance_id)
-redis.call('SETEX', holder_key, ttl_seconds, token)
-
--- Step 3: 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, 'registered', current_count}
+return 0

packages/service-library/tests/redis/test_semaphore.py

@@ -290,43 +290,6 @@ async def _raising_context():
     assert await captured_semaphore.get_current_count() == 0
 
 
-async def test_semaphore_ttl_cleanup(
-    redis_client_sdk: RedisClientSDK,
-    semaphore_name: str,
-    semaphore_capacity: int,
-    short_ttl: datetime.timedelta,
-):
-    # Create semaphore with explicit short TTL
-    semaphore = DistributedSemaphore(
-        redis_client=redis_client_sdk,
-        key=semaphore_name,
-        capacity=semaphore_capacity,
-        ttl=short_ttl,
-    )
-
-    # Manually add an expired entry
-    expired_instance_id = "expired-instance"
-    current_time = asyncio.get_event_loop().time()
-    # Make sure it's definitely expired by using the short TTL
-    expired_time = current_time - short_ttl.total_seconds() - 1
-
-    await redis_client_sdk.redis.zadd(
-        semaphore.semaphore_key, {expired_instance_id: expired_time}
-    )
-
-    # Verify the entry was added
-    initial_count = await redis_client_sdk.redis.zcard(semaphore.semaphore_key)
-    assert initial_count == 1
-
-    # Current count should clean up expired entries
-    count = await semaphore.get_current_count()
-    assert count == 0
-
-    # Verify expired entry was removed
-    remaining = await redis_client_sdk.redis.zcard(semaphore.semaphore_key)
-    assert remaining == 0
-
-
 async def test_multiple_semaphores_different_keys(
     redis_client_sdk: RedisClientSDK,
     faker: Faker,