cleanup

sanderegg · sanderegg · commit 0b59b6ff03af · 2025-01-09T17:33:50.000+01:00
diff --git a/packages/service-library/src/servicelib/background_task.py b/packages/service-library/src/servicelib/background_task.py
@@ -6,7 +6,7 @@
 from typing import Final
 
 from common_library.errors_classes import OsparcErrorMixin
-from tenacity import TryAgain
+from tenacity import TryAgain, retry_always, retry_if_exception_type
 from tenacity.asyncio import AsyncRetrying
 from tenacity.stop import stop_after_attempt
 from tenacity.wait import wait_fixed
@@ -26,14 +26,14 @@ class PeriodicTaskCancellationError(OsparcErrorMixin, Exception):
 
 
 class SleepUsingAsyncioEvent:
-    """Sleep strategy that waits on an event to be set."""
+    """Sleep strategy that waits on an event to be set or sleeps."""
 
     def __init__(self, event: "asyncio.Event") -> None:
         self.event = event
 
-    async def __call__(self, timeout: float | None) -> None:
+    async def __call__(self, delay: float | None) -> None:
         with contextlib.suppress(TimeoutError):
-            await asyncio.wait_for(self.event.wait(), timeout=timeout)
+            await asyncio.wait_for(self.event.wait(), timeout=delay)
             self.event.clear()
 
 
@@ -42,10 +42,18 @@ async def _periodic_scheduled_task(
     *,
     interval: datetime.timedelta,
     task_name: str,
+    raise_on_error: bool,
     early_wake_up_event: asyncio.Event | None,
     **task_kwargs,
 ) -> None:
-    # NOTE: This retries forever unless cancelled
+    """periodically runs task with a given interval.
+    If raise_on_error is False, the task will be retried indefinitely unless cancelled.
+    If raise_on_error is True, the task will be retried indefinitely unless cancelled or an exception is raised.
+    If early_wake_up_event is set, the task might be woken up earlier than interval when the event is set.
+
+    Raises:
+        task exception if raise_on_error is True
+    """
     nap = (
         asyncio.sleep
         if early_wake_up_event is None
@@ -54,6 +62,8 @@ async def _periodic_scheduled_task(
     async for attempt in AsyncRetrying(
         sleep=nap,
         wait=wait_fixed(interval.total_seconds()),
+        reraise=True,
+        retry=retry_if_exception_type(TryAgain) if raise_on_error else retry_always,
     ):
         with attempt:
             with (
@@ -74,6 +84,7 @@ def start_periodic_task(
     *,
     interval: datetime.timedelta,
     task_name: str,
+    raise_on_error: bool = False,
     wait_before_running: datetime.timedelta = datetime.timedelta(0),
     early_wake_up_event: asyncio.Event | None = None,
     **kwargs,
@@ -89,6 +100,7 @@ def start_periodic_task(
                 task,
                 interval=interval,
                 task_name=task_name,
+                raise_on_error=raise_on_error,
                 early_wake_up_event=early_wake_up_event,
                 **kwargs,
             ),
@@ -149,7 +161,11 @@ async def periodic_task(
     asyncio_task: asyncio.Task | None = None
     try:
         asyncio_task = start_periodic_task(
-            task, interval=interval, task_name=task_name, **kwargs
+            task,
+            interval=interval,
+            task_name=task_name,
+            raise_on_error=raise_on_error,
+            **kwargs,
         )
         yield asyncio_task
     finally:
diff --git a/packages/service-library/src/servicelib/redis/_decorators.py b/packages/service-library/src/servicelib/redis/_decorators.py
@@ -2,14 +2,15 @@
 import contextlib
 import functools
 import logging
-from collections.abc import Awaitable, Callable
-from typing import ParamSpec, TypeVar
+from collections.abc import Callable
+from typing import Any, Coroutine, ParamSpec, TypeVar
 
 import redis.exceptions
+from servicelib.logging_utils import log_context
 
 from ..background_task import periodic_task
 from ._client import RedisClientSDK
-from ._constants import DEFAULT_LOCK_TTL
+from ._constants import DEFAULT_LOCK_TTL, SHUTDOWN_TIMEOUT_S
 from ._errors import CouldNotAcquireLockError
 from ._utils import auto_extend_lock
 
@@ -24,26 +25,30 @@ def exclusive(
     *,
     lock_key: str | Callable[..., str],
     lock_value: bytes | str | None = None,
-) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]:
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, R]]], Callable[P, Coroutine[Any, Any, R]]
+]:
     """
-    Define a method to run exclusively across
-    processes by leveraging a Redis Lock.
-
-    parameters:
-    redis: the redis client SDK
-    lock_key: a string as the name of the lock (good practice: app_name:lock_name)
-    lock_value: some additional data that can be retrieved by another client
-
-    Raises:
-        - ValueError if used incorrectly
-        - CouldNotAcquireLockError if the lock could not be acquired
+        Define a method to run exclusively across
+        processes by leveraging a Redis Lock.
+    a1f69fdefa14fae2fee03fac7e89f27e44b13aa9
+        parameters:
+        redis: the redis client SDK
+        lock_key: a string as the name of the lock (good practice: app_name:lock_name)
+        lock_value: some additional data that can be retrieved by another client
+
+        Raises:
+            - ValueError if used incorrectly
+            - CouldNotAcquireLockError if the lock could not be acquired
     """
 
     if not lock_key:
         msg = "lock_key cannot be empty string!"
         raise ValueError(msg)
 
-    def decorator(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]:
+    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:
             redis_lock_key = (
@@ -58,51 +63,52 @@ async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
             )
             assert isinstance(redis_client, RedisClientSDK)  # nosec
 
-            lock_ttl = DEFAULT_LOCK_TTL
-
-            lock = client.create_lock(redis_lock_key, ttl=lock_ttl)
+            lock = client.create_lock(redis_lock_key, ttl=DEFAULT_LOCK_TTL)
             if not await lock.acquire(token=lock_value):
                 raise CouldNotAcquireLockError(lock=lock)
 
             try:
                 async with periodic_task(
                     auto_extend_lock,
-                    interval=lock_ttl / 2,
+                    interval=DEFAULT_LOCK_TTL / 2,
                     task_name=f"autoextend_exclusive_lock_{redis_lock_key}",
                     raise_on_error=True,
                     lock=lock,
                 ) as auto_extend_task:
                     work_task = asyncio.create_task(
                         func(*args, **kwargs), name=f"exclusive_{func.__name__}"
                     )
-                    done, pending = await asyncio.wait(
+                    done, _pending = await asyncio.wait(
                         [work_task, auto_extend_task],
                         return_when=asyncio.FIRST_COMPLETED,
                     )
-                    # the task finished first, let's return it
+                    # the task finished, let's return its result whatever it is
                     if work_task in done:
                         return await work_task
 
-                    # the auto extend tasks finished first, meaning it could not extend the lock!
-                    # let's cancel the work task and raise an error
-                    work_task.cancel()
-                    with contextlib.suppress(asyncio.CancelledError, TimeoutError):
-                        # TODO: shall we raise the other errors?
-                        await asyncio.wait_for(work_task, timeout=3)
-
+                    # the auto extend task can only finish if it raised an error, so it's bad
+                    _logger.error(
+                        "lock %s could not be auto-extended, cancelling work task! "
+                        "TIP: check connection to Redis DBs or look for Synchronous "
+                        "code that might block the auto-extender task.",
+                        lock.name,
+                    )
+                    with log_context(_logger, logging.DEBUG, msg="cancel work task"):
+                        work_task.cancel()
+                        with contextlib.suppress(asyncio.CancelledError, TimeoutError):
+                            # this will raise any other error that could have happened in the work task
+                            await asyncio.wait_for(
+                                work_task, timeout=SHUTDOWN_TIMEOUT_S
+                            )
+                    # return the extend task raised error
                     return await auto_extend_task
 
-                return result
             finally:
                 with contextlib.suppress(redis.exceptions.LockNotOwnedError):
-                    # in the case where the lock would have been lost, this would raise
+                    # in the case where the lock would have been lost,
+                    # this would raise again and is not necessary
                     await lock.release()
 
-            async with redis_client.lock_context(
-                lock_key=redis_lock_key, lock_value=lock_value
-            ):
-                return await func(*args, **kwargs)
-
         return wrapper
 
     return decorator
