work with lost tokens

Author: sanderegg

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

@@ -107,6 +107,8 @@ class DistributedSemaphore(BaseModel):
     release_script: ClassVar[AsyncScript | None] = None
     renew_script: ClassVar[AsyncScript | None] = None
 
+    _token: str | None = None  # currently held token, if any
+
     @classmethod
     def _register_scripts(cls, redis_client: RedisClientSDK) -> None:
         """Register Lua scripts with Redis if not already done.
@@ -258,15 +260,15 @@ async def _try_acquire() -> list[str] | None:
         assert tokens_key_token is not None  # nosec
         assert len(tokens_key_token) == 2  # nosec  # noqa: PLR2004
         assert tokens_key_token[0] == self.tokens_key  # nosec
-        token = tokens_key_token[1]
+        self._token = tokens_key_token[1]
 
         # set up the semaphore holder with a TTL
         cls = type(self)
         assert cls.acquire_script is not None  # nosec
         result = await cls.acquire_script(  # pylint: disable=not-callable
             keys=[self.holders_set, self.holder_key],
             args=[
-                token,
+                self._token,
                 self.instance_id,
                 ttl_seconds,
             ],
@@ -300,9 +302,12 @@ async def release(self) -> None:
         # Execute the release Lua script atomically
         cls = type(self)
         assert cls.release_script is not None  # nosec
+        release_args = [self.instance_id]
+        if self._token is not None:
+            release_args.append(self._token)
         result = await cls.release_script(  # pylint: disable=not-callable
             keys=[self.tokens_key, self.holders_set, self.holder_key],
-            args=[self.instance_id],
+            args=release_args,
             client=self.redis_client.redis,
         )
 

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

@@ -4,6 +4,7 @@
 -- KEYS[3]: holder_key (individual holder TTL key for this instance)
 
 -- ARGV[1]: instance_id
+-- ARGV[2]: passed_token (the token held by this instance or nil if unknown)
 --
 -- Returns: {exit_code, status, current_count}
 -- exit_code: 0 if released, 255 if failed
@@ -14,6 +15,7 @@ local holders_key = KEYS[2]
 local holder_key = KEYS[3]
 
 local instance_id = ARGV[1]
+local passed_token = ARGV[2]
 
 -- Step 1: Check if this instance is currently a holder
 local is_holder = redis.call('SISMEMBER', holders_key, instance_id)
@@ -29,6 +31,11 @@ if not token then
     -- this indicates a lost semaphore (e.g. due to TTL expiry)
     -- remove from holders set and return error
     redis.call('SREM', holders_key, instance_id)
+    -- if the token was passed return it to the pool
+    if passed_token then
+        redis.call('LPUSH', tokens_key, passed_token)
+    end
+    -- Note: we do NOT push a recovered token since we don't know its state
     return {255, 'expired', redis.call('SCARD', holders_key)}
 end
 

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

@@ -58,9 +58,9 @@ async def test_semaphore_initialization(
     assert semaphore.blocking is True
     assert semaphore.instance_id is not None
     assert semaphore.semaphore_key == f"{SEMAPHORE_KEY_PREFIX}{semaphore_name}"
-    assert semaphore.holder_key.startswith(
-        f"{SEMAPHORE_KEY_PREFIX}{semaphore_name}:holders:"
-    )
+    assert semaphore.tokens_key.startswith(f"{semaphore.semaphore_key}:")
+    assert semaphore.holders_set.startswith(f"{semaphore.semaphore_key}:")
+    assert semaphore.holder_key.startswith(f"{semaphore.semaphore_key}:")
 
 
 async def test_invalid_semaphore_initialization(
@@ -102,6 +102,28 @@ async def test_invalid_semaphore_initialization(
         )
 
 
+async def _assert_semaphore_redis_state(
+    redis_client_sdk: RedisClientSDK,
+    semaphore: DistributedSemaphore,
+    *,
+    expected_count: int,
+    expected_free_tokens: int,
+    expected_expired: bool = False,
+):
+    """Helper to assert the internal Redis state of the semaphore"""
+    holders = await redis_client_sdk.redis.smembers(semaphore.holders_set)
+    assert len(holders) == expected_count
+    if expected_count > 0:
+        assert semaphore.instance_id in holders
+        holder_key_exists = await redis_client_sdk.redis.exists(semaphore.holder_key)
+        if expected_expired:
+            assert holder_key_exists == 0
+        else:
+            assert holder_key_exists == 1
+    tokens = await redis_client_sdk.redis.lrange(semaphore.tokens_key, 0, -1)
+    assert len(tokens) == expected_free_tokens
+
+
 async def test_semaphore_acquire_release_basic(
     redis_client_sdk: RedisClientSDK,
     semaphore_name: str,
@@ -119,46 +141,88 @@ async def test_semaphore_acquire_release_basic(
     assert await semaphore.current_count() == 0
     assert await semaphore.size() == semaphore_capacity
     assert await semaphore.is_acquired() is False
+    await _assert_semaphore_redis_state(
+        redis_client_sdk,
+        semaphore,
+        expected_count=0,
+        expected_free_tokens=semaphore_capacity,
+    )
 
     # Acquire
     result = await semaphore.acquire()
     assert result is True
     assert await semaphore.current_count() == 1
     assert await semaphore.size() == semaphore_capacity - 1
     assert await semaphore.is_acquired() is True
+    await _assert_semaphore_redis_state(
+        redis_client_sdk,
+        semaphore,
+        expected_count=1,
+        expected_free_tokens=semaphore_capacity - 1,
+    )
 
     # Acquire again on same instance should return True immediately and keep the same count (reentrant)
     result = await semaphore.acquire()
     assert result is True
     assert await semaphore.current_count() == 1
     assert await semaphore.size() == semaphore_capacity - 1
     assert await semaphore.is_acquired() is True
+    await _assert_semaphore_redis_state(
+        redis_client_sdk,
+        semaphore,
+        expected_count=1,
+        expected_free_tokens=semaphore_capacity - 1,
+    )
 
     # reacquire should just work
     await semaphore.reacquire()
     assert await semaphore.current_count() == 1
     assert await semaphore.size() == semaphore_capacity - 1
     assert await semaphore.is_acquired() is True
+    await _assert_semaphore_redis_state(
+        redis_client_sdk,
+        semaphore,
+        expected_count=1,
+        expected_free_tokens=semaphore_capacity - 1,
+    )
 
     # Release
     await semaphore.release()
     assert await semaphore.current_count() == 0
     assert await semaphore.size() == semaphore_capacity
     assert await semaphore.is_acquired() is False
+    await _assert_semaphore_redis_state(
+        redis_client_sdk,
+        semaphore,
+        expected_count=0,
+        expected_free_tokens=semaphore_capacity,
+    )
 
     # reacquire after release should fail
     with pytest.raises(
         SemaphoreNotAcquiredError,
         match=f"Semaphore '{semaphore_name}' was not acquired by this instance",
     ):
         await semaphore.reacquire()
+    await _assert_semaphore_redis_state(
+        redis_client_sdk,
+        semaphore,
+        expected_count=0,
+        expected_free_tokens=semaphore_capacity,
+    )
 
     # so does release again
     with pytest.raises(
         SemaphoreNotAcquiredError,
         match=f"Semaphore '{semaphore_name}' was not acquired by this instance",
     ):
         await semaphore.release()
+    await _assert_semaphore_redis_state(
+        redis_client_sdk,
+        semaphore,
+        expected_count=0,
+        expected_free_tokens=semaphore_capacity,
+    )
 
 
 async def test_semaphore_acquire_release_with_ttl_expiry(
@@ -176,21 +240,48 @@ async def test_semaphore_acquire_release_with_ttl_expiry(
     await semaphore.acquire()
     assert await semaphore.current_count() == 1
     assert await semaphore.size() == semaphore_capacity - 1
+    await _assert_semaphore_redis_state(
+        redis_client_sdk,
+        semaphore,
+        expected_count=1,
+        expected_free_tokens=semaphore_capacity - 1,
+    )
+
     # wait for TTL to expire
     await asyncio.sleep(with_short_default_semaphore_ttl.total_seconds() + 0.1)
+    await _assert_semaphore_redis_state(
+        redis_client_sdk,
+        semaphore,
+        expected_count=1,
+        expected_free_tokens=semaphore_capacity - 1,
+        expected_expired=True,
+    )
 
     # TTL expired, reacquire should fail
     with pytest.raises(
         SemaphoreLostError,
         match=f"Semaphore '{semaphore_name}' was lost by this instance",
     ):
         await semaphore.reacquire()
+    await _assert_semaphore_redis_state(
+        redis_client_sdk,
+        semaphore,
+        expected_count=1,
+        expected_free_tokens=semaphore_capacity - 1,
+        expected_expired=True,
+    )
     # and release should also fail
     with pytest.raises(
         SemaphoreLostError,
         match=f"Semaphore '{semaphore_name}' was lost by this instance",
     ):
         await semaphore.release()
+    await _assert_semaphore_redis_state(
+        redis_client_sdk,
+        semaphore,
+        expected_count=0,
+        expected_free_tokens=semaphore_capacity,
+    )
 
 
 async def test_semaphore_context_manager(