refactor

Author: sanderegg

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

@@ -1,12 +1,15 @@
 import asyncio
+import contextlib
 import datetime
 import logging
 from collections import deque
 from collections.abc import Awaitable, Callable, Coroutine
 from contextlib import suppress
 from dataclasses import dataclass
 from functools import wraps
-from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar
+from typing import TYPE_CHECKING, Any, Final, ParamSpec, TypeVar
+
+from tenacity import before_sleep_log, retry, stop_after_attempt
 
 from . import tracing
 from .utils_profiling_middleware import dont_profile, is_profiling, profile_context
@@ -229,3 +232,32 @@ async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
         return wrapper
 
     return decorator
+
+
+_MAX_TASK_CANCELLATION_ATTEMPTS: Final[int] = 3
+
+
+@retry(
+    reraise=True,
+    stop=stop_after_attempt(_MAX_TASK_CANCELLATION_ATTEMPTS),
+    before_sleep=before_sleep_log(_logger, logging.WARNING, exc_info=True),
+)
+async def retried_cancel_task(
+    task: asyncio.Task,
+    *,
+    timeout: float | None,  # noqa: ASYNC109
+) -> None:
+    """Reliable task cancellation. Some libraries will just hang without
+    cancelling the task. The operation is retried _MAX_TASK_CANCELLATION_ATTEMPTS times
+    in case of timeout.
+
+    :param task: task to be canceled
+    :param timeout_delay: duration (in seconds) to wait before giving
+        up the cancellation. If None it waits forever. Will be repeated _MAX_TASK_CANCELLATION_ATTEMPTS times.
+    :raises TimeoutError: raised if cannot cancel the task. Errors during cancellations are suppressed (backwards compatibility).
+    """
+
+    task.cancel()
+    async with asyncio.timeout(timeout):
+        with contextlib.suppress(BaseException):
+            await task

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

@@ -5,24 +5,17 @@
 from collections.abc import AsyncIterator, Awaitable, Callable
 from typing import Final
 
-from common_library.errors_classes import OsparcErrorMixin
 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
 
-from .async_utils import with_delay
+from .async_utils import retried_cancel_task, with_delay
 from .logging_utils import log_catch, log_context
 
 _logger = logging.getLogger(__name__)
 
 
 _DEFAULT_STOP_TIMEOUT_S: Final[int] = 5
-_MAX_TASK_CANCELLATION_ATTEMPTS: Final[int] = 3
-
-
-class PeriodicTaskCancellationError(OsparcErrorMixin, Exception):
-    msg_template: str = "Could not cancel task '{task_name}'"
 
 
 class SleepUsingAsyncioEvent:
@@ -108,35 +101,6 @@ def start_periodic_task(
         )
 
 
-async def cancel_task(
-    task: asyncio.Task,
-    *,
-    timeout: float | None,
-    cancellation_attempts: int = _MAX_TASK_CANCELLATION_ATTEMPTS,
-) -> None:
-    """Reliable task cancellation. Some libraries will just hang without
-    cancelling the task. It is important to retry the operation to provide
-    a timeout in that situation to avoid forever pending tasks.
-
-    :param task: task to be canceled
-    :param timeout: total duration (in seconds) to wait before giving
-        up the cancellation. If None it waits forever.
-    :raises TryAgain: raised if cannot cancel the task.
-    """
-    async for attempt in AsyncRetrying(
-        stop=stop_after_attempt(cancellation_attempts), reraise=True
-    ):
-        with attempt:
-            task.cancel()
-            _, pending = await asyncio.wait((task,), timeout=timeout)
-            if pending:
-                task_name = task.get_name()
-                _logger.info(
-                    "tried to cancel '%s' but timed-out! %s", task_name, pending
-                )
-                raise PeriodicTaskCancellationError(task_name=task_name)
-
-
 async def stop_periodic_task(
     asyncio_task: asyncio.Task, *, timeout: float | None = None
 ) -> None:
@@ -145,7 +109,7 @@ async def stop_periodic_task(
         logging.DEBUG,
         msg=f"cancel periodic background task '{asyncio_task.get_name()}'",
     ):
-        await cancel_task(asyncio_task, timeout=timeout)
+        await retried_cancel_task(asyncio_task, timeout=timeout)
 
 
 @contextlib.asynccontextmanager

services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_scheduler.py

@@ -41,11 +41,8 @@
 from models_library.users import UserID
 from models_library.wallets import WalletID
 from pydantic import NonNegativeFloat
-from servicelib.background_task import (
-    cancel_task,
-    start_periodic_task,
-    stop_periodic_task,
-)
+from servicelib.async_utils import retried_cancel_task
+from servicelib.background_task import start_periodic_task, stop_periodic_task
 from servicelib.fastapi.long_running_tasks.client import ProgressCallback
 from servicelib.fastapi.long_running_tasks.server import TaskProgress
 from servicelib.redis import RedisClientsManager, exclusive
@@ -368,7 +365,7 @@ async def mark_service_for_removal(
                     self._service_observation_task[service_name]
                 )
                 if isinstance(service_task, asyncio.Task):
-                    await cancel_task(service_task, timeout=10)
+                    await retried_cancel_task(service_task, timeout=10)
 
             if skip_observation_recreation:
                 return

services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/prometheus_metrics.py

@@ -9,7 +9,7 @@
 from fastapi import FastAPI, status
 from models_library.callbacks_mapping import CallbacksMapping, UserServiceCommand
 from pydantic import BaseModel, NonNegativeFloat, NonNegativeInt
-from servicelib.background_task import cancel_task
+from servicelib.async_utils import retried_cancel_task
 from servicelib.logging_utils import log_context
 from servicelib.sequences_utils import pairwise
 from simcore_service_dynamic_sidecar.core.errors import (
@@ -143,7 +143,7 @@ async def start(self) -> None:
     async def stop(self) -> None:
         with log_context(_logger, logging.INFO, "shutdown service metrics recovery"):
             if self._metrics_recovery_task:
-                await cancel_task(
+                await retried_cancel_task(
                     self._metrics_recovery_task, timeout=_TASK_CANCELLATION_TIMEOUT_S
                 )