fixes

Author: sanderegg

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

@@ -101,22 +101,30 @@ class DistributedSemaphore(BaseModel):
     ] = DEFAULT_FACTORY
 
     # Private state attributes (not part of the model)
-    _acquire_script: ClassVar[AsyncScript]
-    _count_script: ClassVar[AsyncScript]
-    _release_script: ClassVar[AsyncScript]
-    _renew_script: ClassVar[AsyncScript]
+    acquire_script: ClassVar[AsyncScript | None] = None
+    count_script: ClassVar[AsyncScript | None] = None
+    release_script: ClassVar[AsyncScript | None] = None
+    renew_script: ClassVar[AsyncScript | None] = None
 
     @classmethod
-    def _register_scripts(cls, redis_client) -> None:
-        if cls._acquire_script is None:
-            cls._acquire_script = redis_client.register_script(ACQUIRE_SEMAPHORE_SCRIPT)
-            cls._count_script = redis_client.register_script(COUNT_SEMAPHORE_SCRIPT)
-            cls._release_script = redis_client.register_script(RELEASE_SEMAPHORE_SCRIPT)
-            cls._renew_script = redis_client.register_script(RENEW_SEMAPHORE_SCRIPT)
+    def _register_scripts(cls, redis_client: RedisClientSDK) -> None:
+        if cls.acquire_script is None:
+            cls.acquire_script = redis_client.redis.register_script(
+                ACQUIRE_SEMAPHORE_SCRIPT
+            )
+            cls.count_script = redis_client.redis.register_script(
+                COUNT_SEMAPHORE_SCRIPT
+            )
+            cls.release_script = redis_client.redis.register_script(
+                RELEASE_SEMAPHORE_SCRIPT
+            )
+            cls.renew_script = redis_client.redis.register_script(
+                RENEW_SEMAPHORE_SCRIPT
+            )
 
     def __init__(self, **data) -> None:
         super().__init__(**data)
-        type(self)._register_scripts(self.redis_client)  # noqa: SLF001
+        self.__class__._register_scripts(self.redis_client)  # noqa: SLF001
 
     @computed_field  # type: ignore[prop-decorator]
     @property
@@ -196,7 +204,7 @@ async def release(self) -> None:
         ttl_seconds = int(self.ttl.total_seconds())
 
         # Execute the release Lua script atomically
-        result = await type(self)._release_script(  # noqa: SLF001
+        result = await type(self).release_script(  # noqa: SLF001
             keys=(
                 self.semaphore_key,
                 self.holder_key,
@@ -231,7 +239,7 @@ async def _try_acquire(self) -> bool:
         ttl_seconds = int(self.ttl.total_seconds())
 
         # Execute the Lua script atomically
-        result = await type(self)._acquire_script(  # noqa: SLF001
+        result = await type(self).acquire_script(  # noqa: SLF001
             keys=(self.semaphore_key, self.holder_key),
             args=(self.instance_id, str(self.capacity), str(ttl_seconds)),
             client=self.redis_client.redis,
@@ -274,7 +282,7 @@ async def reacquire(self) -> None:
         ttl_seconds = int(self.ttl.total_seconds())
 
         # Execute the renewal Lua script atomically
-        result = await type(self)._renew_script(  # noqa: SLF001
+        result = await type(self).renew_script(  # noqa: SLF001
             keys=(self.semaphore_key, self.holder_key),
             args=(
                 self.instance_id,
@@ -322,7 +330,7 @@ async def get_current_count(self) -> int:
         ttl_seconds = int(self.ttl.total_seconds())
 
         # Execute the count Lua script atomically
-        result = await type(self)._count_script(  # noqa: SLF001
+        result = await type(self).count_script(  # noqa: SLF001
             keys=(self.semaphore_key,),
             args=(str(ttl_seconds),),
             client=self.redis_client.redis,

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

@@ -5,7 +5,7 @@
 import socket
 from collections.abc import AsyncIterator, Callable, Coroutine
 from contextlib import asynccontextmanager
-from typing import Any, ParamSpec, TypeVar
+from typing import Any, AsyncContextManager, ParamSpec, TypeVar
 
 from common_library.async_tools import cancel_wait_task
 
@@ -210,7 +210,10 @@ def with_limited_concurrency_cm(
     ttl: datetime.timedelta = DEFAULT_SEMAPHORE_TTL,
     blocking: bool = True,
     blocking_timeout: datetime.timedelta | None = DEFAULT_SOCKET_TIMEOUT,
-) -> Callable[[Callable[P, AsyncIterator[R]]], Callable[P, AsyncIterator[R]]]:
+) -> Callable[
+    [Callable[P, AsyncContextManager[R]]],
+    Callable[P, AsyncContextManager[R]],
+]:
     """
     Decorator to limit concurrent execution of async context managers using a distributed semaphore.
 
@@ -244,9 +247,10 @@ async def get_cluster_client():
     """
 
     def _decorator(
-        cm_func: Callable[P, AsyncIterator[R]],
-    ) -> Callable[P, AsyncIterator[R]]:
+        cm_func: Callable[P, AsyncContextManager[R]],
+    ) -> Callable[P, AsyncContextManager[R]]:
         @functools.wraps(cm_func)
+        @asynccontextmanager
         async def _wrapper(*args: P.args, **kwargs: P.kwargs) -> AsyncIterator[R]:
             semaphore, semaphore_key = _create_semaphore(
                 redis_client,
@@ -259,11 +263,13 @@ async def _wrapper(*args: P.args, **kwargs: P.kwargs) -> AsyncIterator[R]:
                 kwargs=kwargs,
             )
 
-            async with _managed_semaphore_execution(
-                semaphore, semaphore_key, ttl, f"context_manager_{cm_func.__name__}"
+            async with (
+                _managed_semaphore_execution(
+                    semaphore, semaphore_key, ttl, f"context_manager_{cm_func.__name__}"
+                ),
+                cm_func(*args, **kwargs) as value,
             ):
-                async for value in cm_func(*args, **kwargs):
-                    yield value
+                yield value
 
         return _wrapper
 

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

@@ -401,19 +401,19 @@ async def test_context_manager_basic_functionality(
 ):
     call_count = 0
 
-    @asynccontextmanager
     @with_limited_concurrency_cm(
         redis_client_sdk,
         key=semaphore_name,
         capacity=1,
     )
+    @asynccontextmanager
     async def limited_context_manager():
         nonlocal call_count
         call_count += 1
         yield call_count
 
     # Multiple concurrent context managers
-    async def use_context_manager():
+    async def use_context_manager() -> int:
         async with limited_context_manager() as value:
             await asyncio.sleep(0.1)
             return value
@@ -433,12 +433,12 @@ async def test_context_manager_capacity_enforcement(
     concurrent_count = 0
     max_concurrent = 0
 
-    @asynccontextmanager
     @with_limited_concurrency_cm(
         redis_client_sdk,
         key=semaphore_name,
         capacity=2,
     )
+    @asynccontextmanager
     async def limited_context_manager():
         nonlocal concurrent_count, max_concurrent
         concurrent_count += 1
@@ -449,12 +449,12 @@ async def limited_context_manager():
         finally:
             concurrent_count -= 1
 
-    async def use_context_manager():
+    async def use_context_manager() -> None:
         async with limited_context_manager():
             await asyncio.sleep(0.1)
 
-    # Start 5 concurrent context managers
-    tasks = [asyncio.create_task(use_context_manager()) for _ in range(5)]
+    # Start concurrent context managers
+    tasks = [asyncio.create_task(use_context_manager()) for _ in range(20)]
     await asyncio.gather(*tasks)
 
     # Should never exceed capacity of 2
@@ -465,12 +465,12 @@ async def test_context_manager_exception_handling(
     redis_client_sdk: RedisClientSDK,
     semaphore_name: str,
 ):
-    @asynccontextmanager
     @with_limited_concurrency_cm(
         redis_client_sdk,
         key=semaphore_name,
         capacity=1,
     )
+    @asynccontextmanager
     async def failing_context_manager():
         yield
         raise RuntimeError("Test exception")
@@ -480,12 +480,13 @@ async def failing_context_manager():
             pass
 
     # Semaphore should be released even after exception
-    @asynccontextmanager
+
     @with_limited_concurrency_cm(
         redis_client_sdk,
         key=semaphore_name,
         capacity=1,
     )
+    @asynccontextmanager
     async def success_context_manager():
         yield "success"
 
@@ -502,13 +503,13 @@ async def test_context_manager_auto_renewal(
     work_started = asyncio.Event()
     work_completed = asyncio.Event()
 
-    @asynccontextmanager
     @with_limited_concurrency_cm(
         redis_client_sdk,
         key=semaphore_name,
         capacity=semaphore_capacity,
         ttl=short_ttl,
     )
+    @asynccontextmanager
     async def long_running_context_manager():
         work_started.set()
         yield "data"
@@ -558,12 +559,12 @@ def get_key(user_id: str, resource: str) -> str:
     def get_capacity(user_id: str, resource: str) -> int:
         return 2
 
-    @asynccontextmanager
     @with_limited_concurrency_cm(
         get_redis_client,
         key=get_key,
         capacity=get_capacity,
     )
+    @asynccontextmanager
     async def process_user_resource_cm(user_id: str, resource: str):
         executed_keys.append(f"{user_id}-{resource}")
         yield f"processed-{user_id}-{resource}"
@@ -597,14 +598,14 @@ async def test_context_manager_non_blocking_behavior(
 ):
     started_event = asyncio.Event()
 
-    @asynccontextmanager
     @with_limited_concurrency_cm(
         redis_client_sdk,
         key=semaphore_name,
         capacity=1,
         blocking=True,
         blocking_timeout=datetime.timedelta(seconds=0.1),
     )
+    @asynccontextmanager
     async def limited_context_manager():
         started_event.set()
         yield
@@ -619,14 +620,15 @@ async def long_running_cm():
     await started_event.wait()  # Wait until semaphore is actually acquired
 
     # Second context manager should timeout and raise an exception
-    @asynccontextmanager
+
     @with_limited_concurrency_cm(
         redis_client_sdk,
         key=semaphore_name,
         capacity=1,
         blocking=True,
         blocking_timeout=datetime.timedelta(seconds=0.1),
     )
+    @asynccontextmanager
     async def timeout_context_manager():
         yield
 
@@ -645,13 +647,13 @@ async def test_context_manager_lose_semaphore_raises(
 ):
     work_started = asyncio.Event()
 
-    @asynccontextmanager
     @with_limited_concurrency_cm(
         redis_client_sdk,
         key=semaphore_name,
         capacity=semaphore_capacity,
         ttl=short_ttl,
     )
+    @asynccontextmanager
     async def context_manager_that_should_fail():
         work_started.set()
         yield "data"