add task notification utility

Author: Andrei Neagu

packages/service-library/src/servicelib/async_utils.py

@@ -1,14 +1,18 @@
 import asyncio
 import logging
-from collections import deque
+from collections.abc import Awaitable, Callable
 from contextlib import suppress
 from dataclasses import dataclass
+from datetime import timedelta
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Awaitable, Callable, Deque
+from typing import TYPE_CHECKING, Any, Coroutine, Final, TypeVar
+
+from pydantic import NonNegativeFloat
+from servicelib.background_task import cancel_task
 
 from .utils_profiling_middleware import dont_profile, is_profiling, profile_context
 
-logger = logging.getLogger(__name__)
+_logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
     Queue = asyncio.Queue
@@ -51,7 +55,7 @@ async def _safe_cancel(context: Context) -> None:
                 await context.task
     except RuntimeError as e:
         if "Event loop is closed" in f"{e}":
-            logger.warning("event loop is closed and could not cancel %s", context)
+            _logger.warning("event loop is closed and could not cancel %s", context)
         else:
             raise
 
@@ -62,7 +66,7 @@ async def cancel_sequential_workers() -> None:
         await _safe_cancel(context)
 
     _sequential_jobs_contexts.clear()
-    logger.info("All run_sequentially_in_context pending workers stopped")
+    _logger.info("All run_sequentially_in_context pending workers stopped")
 
 
 # NOTE: If you get funny mismatches with mypy in returned values it might be due to this decorator.
@@ -118,22 +122,22 @@ def _get_context(args: Any, kwargs: dict) -> Context:
             search_args = dict(zip(arg_names, args))
             search_args.update(kwargs)
 
-            key_parts: Deque[str] = deque()
+            key_parts: list[str] = []
             for arg in target_args:
                 sub_args = arg.split(".")
                 main_arg = sub_args[0]
                 if main_arg not in search_args:
-                    raise ValueError(
+                    msg = (
                         f"Expected '{main_arg}' in '{decorated_function.__name__}'"
                         f" arguments. Got '{search_args}'"
                     )
+                    raise ValueError(msg)
                 context_key = search_args[main_arg]
                 for attribute in sub_args[1:]:
                     potential_key = getattr(context_key, attribute)
                     if not potential_key:
-                        raise ValueError(
-                            f"Expected '{attribute}' attribute in '{context_key.__name__}' arguments."
-                        )
+                        msg = f"Expected '{attribute}' attribute in '{context_key.__name__}' arguments."
+                        raise ValueError(msg)
                     context_key = potential_key
 
                 key_parts.append(f"{decorated_function.__name__}_{context_key}")
@@ -200,3 +204,35 @@ async def worker(in_q: Queue[QueueElement], out_q: Queue) -> None:
         return wrapper
 
     return decorator
+
+
+T = TypeVar("T")
+_CANCELLATION_TIMEOUT: Final[NonNegativeFloat] = 0.1
+
+
+async def _monitor_task(
+    notification_hook: Callable[[], Awaitable[None]], notify_after: timedelta
+) -> None:
+    await asyncio.sleep(notify_after.total_seconds())
+    await notification_hook()
+
+
+async def notify_when_over_threshold(
+    task: Coroutine[Any, Any, T],
+    *,
+    notification_hook: Callable[[], Awaitable[None]],
+    notify_after: timedelta,
+) -> T:
+    monitor_task = asyncio.create_task(_monitor_task(notification_hook, notify_after))
+
+    try:
+        result = await task
+        await cancel_task(monitor_task, timeout=_CANCELLATION_TIMEOUT)
+    except asyncio.CancelledError:
+        await cancel_task(monitor_task, timeout=_CANCELLATION_TIMEOUT)
+        raise
+    except Exception:
+        await cancel_task(monitor_task, timeout=_CANCELLATION_TIMEOUT)
+        raise
+
+    return result

packages/service-library/tests/test_async_utils.py

@@ -7,13 +7,16 @@
 import random
 from collections import deque
 from dataclasses import dataclass
+from datetime import timedelta
 from time import time
 from typing import Any
+from unittest.mock import Mock
 
 import pytest
 from faker import Faker
 from servicelib.async_utils import (
     _sequential_jobs_contexts,
+    notify_when_over_threshold,
     run_sequentially_in_context,
 )
 
@@ -224,3 +227,61 @@ async def test_multiple_context_calls(context_param: int) -> int:
             assert i == await test_multiple_context_calls(i)
 
     assert len(_sequential_jobs_contexts) == RETRIES
+
+
+async def test_notify_when_over_threshold():
+
+    notification_spy = Mock()
+
+    async def notification() -> None:
+        notification_spy()
+        print("notified")
+
+    async def _worker(*, sleep_for: float, raise_error: bool = False) -> float:
+        await asyncio.sleep(sleep_for)
+
+        if raise_error:
+            msg = "raise error as planned"
+            raise RuntimeError(msg)
+
+        return sleep_for
+
+    # 1. finish after
+    result = await notify_when_over_threshold(
+        _worker(sleep_for=0.5),
+        notification_hook=notification,
+        notify_after=timedelta(seconds=0.1),
+    )
+    assert isinstance(result, float)
+    assert result == 0.5
+    assert notification_spy.call_count == 1
+
+    # 2. finish before
+    notification_spy.reset_mock()
+    await notify_when_over_threshold(
+        _worker(sleep_for=0.1),
+        notification_hook=notification,
+        notify_after=timedelta(seconds=0.2),
+    )
+    await asyncio.sleep(0.2)
+    assert notification_spy.call_count == 0
+
+    # 3. raise error before notification
+    notification_spy.reset_mock()
+    with pytest.raises(RuntimeError):
+        await notify_when_over_threshold(
+            _worker(sleep_for=0, raise_error=True),
+            notification_hook=notification,
+            notify_after=timedelta(seconds=0.2),
+        )
+    assert notification_spy.call_count == 0
+
+    # 4. raise after notification
+    notification_spy.reset_mock()
+    with pytest.raises(RuntimeError):
+        await notify_when_over_threshold(
+            _worker(sleep_for=0.2, raise_error=True),
+            notification_hook=notification,
+            notify_after=timedelta(seconds=0.1),
+        )
+    assert notification_spy.call_count == 1