@sanderegg review: replace cancel_wait_task by cancel_and_wait

Author: pcrespov

packages/celery-library/src/celery_library/task.py

@@ -12,8 +12,8 @@
     AbortableTask,
 )
 from celery.exceptions import Ignore  # type: ignore[import-untyped]
+from common_library.async_tools import cancel_and_wait
 from pydantic import NonNegativeInt
-from servicelib.async_utils import cancel_wait_task
 from servicelib.celery.models import TaskID
 
 from .errors import encode_celery_transferrable_error
@@ -62,7 +62,7 @@ async def abort_monitor():
                             abortable_result = AbortableAsyncResult(task_id, app=app)
                             while not main_task.done():
                                 if abortable_result.is_aborted():
-                                    await cancel_wait_task(
+                                    await cancel_and_wait(
                                         main_task,
                                         max_delay=_DEFAULT_CANCEL_TASK_TIMEOUT.total_seconds(),
                                     )

packages/common-library/src/common_library/async_tools.py

@@ -67,17 +67,29 @@ async def maybe_await(
     return obj
 
 
-async def cancel_and_wait(task: asyncio.Task) -> None:
+async def cancel_and_wait(
+    task: asyncio.Task, *, max_delay: float | None = None
+) -> None:
     """Cancels the given task and waits for it to complete.
 
+
+    :param task: task to be canceled
+    :param max_delay: duration (in seconds) to wait before giving
+        up the cancellation. If None it waits forever.
+    :raises TimeoutError: raised if cannot cancel the task.
+
     Accounts for the case where the tasks's owner function is being cancelled.
     In that case, it propagates the cancellation exception upstream.
+
     """
+
     task.cancel()
+
     try:
-        # NOTE shield ensures that cancellation of the caller function won’t stop you
+
+        # NOTE shield ensures that cancellation of the caller function won't stop you
         # from observing the cancellation/finalization of task.
-        await asyncio.shield(task)
+        await asyncio.shield(asyncio.wait_for(task, timeout=max_delay))
 
     except asyncio.CancelledError:
         if not task.cancelled():

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

@@ -8,8 +8,6 @@
 from functools import wraps
 from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar
 
-from common_library.async_tools import cancel_and_wait
-
 from . import tracing
 from .utils_profiling_middleware import dont_profile, is_profiling, profile_context
 
@@ -231,19 +229,3 @@ async def _wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
         return _wrapper
 
     return _decorator
-
-
-async def cancel_wait_task(
-    task: asyncio.Task,
-    *,
-    max_delay: float | None = None,
-) -> None:
-    """Cancel a asyncio.Task and waits for it to finish.
-
-    :param task: task to be canceled
-    :param max_delay: duration (in seconds) to wait before giving
-        up the cancellation. If None it waits forever.
-    :raises TimeoutError: raised if cannot cancel the task.
-    """
-    async with asyncio.timeout(max_delay):
-        await cancel_and_wait(task)

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

@@ -9,7 +9,7 @@
 from tenacity import TryAgain, before_sleep_log, retry, retry_if_exception_type
 from tenacity.wait import wait_fixed
 
-from .async_utils import cancel_wait_task, delayed_start
+from .async_utils import cancel_and_wait, delayed_start
 from .logging_utils import log_catch, log_context
 
 _logger = logging.getLogger(__name__)
@@ -142,4 +142,4 @@ async def periodic_task(
         if asyncio_task is not None:
             # NOTE: this stopping is shielded to prevent the cancellation to propagate
             # into the stopping procedure
-            await asyncio.shield(cancel_wait_task(asyncio_task, max_delay=stop_timeout))
+            await asyncio.shield(cancel_and_wait(asyncio_task, max_delay=stop_timeout))

packages/service-library/src/servicelib/long_running_tasks/task.py

@@ -9,9 +9,9 @@
 from typing import Any, Final, Protocol, TypeAlias
 from uuid import uuid4
 
+from common_library.async_tools import cancel_and_wait
 from models_library.api_schemas_long_running_tasks.base import TaskProgress
 from pydantic import PositiveFloat
-from servicelib.async_utils import cancel_wait_task
 from servicelib.background_task import create_periodic_task
 from servicelib.logging_utils import log_catch
 
@@ -105,7 +105,7 @@ async def teardown(self) -> None:
 
         if self._stale_tasks_monitor_task:
             with log_catch(_logger, reraise=False):
-                await cancel_wait_task(
+                await cancel_and_wait(
                     self._stale_tasks_monitor_task, max_delay=_CANCEL_TASK_TIMEOUT
                 )
 

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

@@ -12,7 +12,7 @@
 from redis.asyncio.retry import Retry
 from redis.backoff import ExponentialBackoff
 
-from ..async_utils import cancel_wait_task
+from ..async_utils import cancel_and_wait
 from ..background_task import periodic
 from ..logging_utils import log_catch, log_context
 from ._constants import (
@@ -88,7 +88,7 @@ async def shutdown(self) -> None:
                 assert self._health_check_task_started_event  # nosec
                 # NOTE: wait for the health check task to have started once before we can cancel it
                 await self._health_check_task_started_event.wait()
-                await cancel_wait_task(
+                await cancel_and_wait(
                     self._health_check_task, max_delay=_HEALTHCHECK_TASK_TIMEOUT_S
                 )
 

packages/service-library/tests/rabbitmq/test_rabbitmq_rpc_interfaces_async_jobs.py

@@ -4,6 +4,7 @@
 from dataclasses import dataclass, field
 
 import pytest
+from common_library.async_tools import cancel_and_wait
 from faker import Faker
 from models_library.api_schemas_rpc_async_jobs.async_jobs import (
     AsyncJobGet,
@@ -16,7 +17,6 @@
 from models_library.progress_bar import ProgressReport
 from models_library.rabbitmq_basic_types import RPCMethodName, RPCNamespace
 from pydantic import TypeAdapter
-from servicelib.async_utils import cancel_wait_task
 from servicelib.rabbitmq import RabbitMQRPCClient, RemoteMethodNotRegisteredError
 from servicelib.rabbitmq.rpc_interfaces.async_jobs.async_jobs import (
     list_jobs,
@@ -137,7 +137,7 @@ async def setup(self) -> None:
     yield
 
     for task in fake_server.tasks:
-        await cancel_wait_task(task)
+        await cancel_and_wait(task)
 
 
 @pytest.mark.parametrize("method", ["result", "status", "cancel"])

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

@@ -10,11 +10,11 @@
 from uuid import UUID
 
 import pytest
+from common_library.async_tools import cancel_and_wait
 from faker import Faker
 from models_library.projects import ProjectID
 from models_library.projects_access import Owner
 from models_library.projects_state import ProjectLocked, ProjectStatus
-from servicelib.async_utils import cancel_wait_task
 from servicelib.redis import (
     ProjectLockError,
     RedisClientSDK,
@@ -141,4 +141,4 @@ async def _locked_fct() -> None:
     with pytest.raises(ProjectLockError):
         await _locked_fct()
 
-    await cancel_wait_task(task1)
+    await cancel_and_wait(task1)

packages/service-library/tests/test_background_task.py

@@ -13,9 +13,9 @@
 from unittest.mock import AsyncMock
 
 import pytest
+from common_library.async_tools import cancel_and_wait
 from faker import Faker
 from pytest_mock.plugin import MockerFixture
-from servicelib.async_utils import cancel_wait_task
 from servicelib.background_task import create_periodic_task, periodic, periodic_task
 
 pytest_simcore_core_services_selection = [
@@ -78,7 +78,7 @@ async def _creator(
     yield _creator
     # cleanup
     await asyncio.gather(
-        *(cancel_wait_task(t, max_delay=stop_task_timeout) for t in created_tasks)
+        *(cancel_and_wait(t, max_delay=stop_task_timeout) for t in created_tasks)
     )
 
 

packages/service-library/tests/test_background_task_utils.py

@@ -13,7 +13,7 @@
 
 import arrow
 import pytest
-from servicelib.async_utils import cancel_wait_task
+from common_library.async_tools import cancel_and_wait
 from servicelib.background_task_utils import exclusive_periodic
 from servicelib.redis import RedisClientSDK
 from settings_library.redis import RedisDatabase
@@ -71,7 +71,7 @@ async def _sleep_task(sleep_interval: float, on_sleep_events: mock.Mock) -> None
 
     await _assert_on_sleep_done(sleep_events, stop_after=stop_after)
 
-    await cancel_wait_task(task, max_delay=5)
+    await cancel_and_wait(task, max_delay=5)
 
     events_timestamps: tuple[float, ...] = tuple(
         x.args[0].timestamp() for x in sleep_events.call_args_list