fix cancelling code

Author: sanderegg

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

@@ -2,11 +2,12 @@
 import logging
 from collections import deque
 from collections.abc import Awaitable, Callable
-from contextlib import suppress
 from dataclasses import dataclass
 from functools import wraps
 from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar
 
+from common_library.async_tools import cancel_wait_task
+
 from . import tracing
 from .utils_profiling_middleware import dont_profile, is_profiling, profile_context
 
@@ -54,9 +55,7 @@ async def _safe_cancel(context: Context) -> None:
     try:
         await context.in_queue.put(None)
         if context.task is not None:
-            context.task.cancel()
-            with suppress(asyncio.CancelledError):
-                await context.task
+            await cancel_wait_task(context.task, max_delay=None)
     except RuntimeError as e:
         if "Event loop is closed" in f"{e}":
             _logger.warning("event loop is closed and could not cancel %s", context)

packages/service-library/src/servicelib/deferred_tasks/_worker_tracker.py

@@ -57,6 +57,11 @@ async def handle_run(
             result_to_return = TaskResultSuccess(value=task_result)
         except asyncio.CancelledError:
             result_to_return = TaskResultCancelledError()
+            current_task = asyncio.current_task()
+            assert current_task is not None  # nosec
+            if current_task.cancelling() > 0:
+                # owner function is being cancelled -> propagate cancellation
+                raise
         except Exception as e:  # pylint:disable=broad-exception-caught
             result_to_return = TaskResultError(
                 error=_format_exception(e),

packages/service-library/src/servicelib/fastapi/requests_decorators.py

@@ -4,6 +4,7 @@
 from functools import wraps
 from typing import Any, Protocol
 
+from common_library.async_tools import cancel_wait_task
 from fastapi import Request, status
 from fastapi.exceptions import HTTPException
 
@@ -13,8 +14,7 @@
 class _HandlerWithRequestArg(Protocol):
     __name__: str
 
-    async def __call__(self, request: Request, *args: Any, **kwargs: Any) -> Any:
-        ...
+    async def __call__(self, request: Request, *args: Any, **kwargs: Any) -> Any: ...
 
 
 def _validate_signature(handler: _HandlerWithRequestArg):
@@ -75,13 +75,8 @@ async def wrapper(request: Request, *args, **kwargs):
 
         # One has completed, cancel the other
         for t in pending:
-            t.cancel()
-
             try:
-                await asyncio.wait_for(t, timeout=3)
-
-            except asyncio.CancelledError:
-                pass
+                await cancel_wait_task(t, max_delay=3)
             except Exception:  # pylint: disable=broad-except
                 if t is handler_task:
                     raise

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

@@ -335,6 +335,11 @@ async def _tasks_monitor(self) -> None:
                     result_field = ResultField(
                         str_error=dumps(TaskCancelledError(task_id=task_id))
                     )
+                    current_task = asyncio.current_task()
+                    assert current_task is not None  # nosec
+                    if current_task.cancelling() > 0:
+                        # owner function is being cancelled -> propagate cancellation
+                        raise
                 except Exception as e:  # pylint:disable=broad-except
                     allowed_errors = TaskRegistry.get_allowed_errors(
                         task_data.registered_task_name

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

@@ -1,4 +1,4 @@
-""" General utils
+"""General utils
 
 IMPORTANT: lowest level module
    I order to avoid cyclic dependences, please
@@ -245,7 +245,7 @@ async def limited_as_completed(
                         future.set_name(f"{tasks_group_prefix}-{future.get_name()}")
                     pending_futures.add(future)
 
-                except (StopIteration, StopAsyncIteration):  # noqa: PERF203
+                except (StopIteration, StopAsyncIteration):
                     completed_all_awaitables = True
             if not pending_futures:
                 return
@@ -294,8 +294,7 @@ async def limited_gather(
     log: logging.Logger = _DEFAULT_LOGGER,
     limit: int = _DEFAULT_LIMITED_CONCURRENCY,
     tasks_group_prefix: str | None = None,
-) -> list[T]:
-    ...
+) -> list[T]: ...
 
 
 @overload
@@ -305,8 +304,7 @@ async def limited_gather(
     log: logging.Logger = _DEFAULT_LOGGER,
     limit: int = _DEFAULT_LIMITED_CONCURRENCY,
     tasks_group_prefix: str | None = None,
-) -> list[T | BaseException]:
-    ...
+) -> list[T | BaseException]: ...
 
 
 async def limited_gather(

services/catalog/src/simcore_service_catalog/core/background_tasks.py

@@ -203,7 +203,7 @@ async def _sync_services_task(app: FastAPI) -> None:
 
             await asyncio.sleep(app.state.settings.CATALOG_BACKGROUND_TASK_REST_TIME)
 
-        except asyncio.CancelledError:  # noqa: PERF203
+        except asyncio.CancelledError:
             # task is stopped
             _logger.info("registry syncing task cancelled")
             raise

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

@@ -1,6 +1,5 @@
 # pylint: disable=relative-beyond-top-level
 
-import asyncio
 import logging
 from copy import deepcopy
 from math import floor
@@ -138,8 +137,6 @@ async def observing_single_service(
     try:
         await _apply_observation_cycle(scheduler, scheduler_data)
         logger.debug("completed observation cycle of %s", f"{service_name=}")
-    except asyncio.CancelledError:  # pylint: disable=try-except-raise
-        raise  # pragma: no cover
     except Exception as exc:  # pylint: disable=broad-except
         service_name = scheduler_data.service_name
 

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

@@ -14,7 +14,6 @@
 """
 
 import asyncio
-import contextlib
 import functools
 import logging
 import time
@@ -131,9 +130,7 @@ async def shutdown(self) -> None:
         if self._trigger_observation_queue_task is not None:
             await self._trigger_observation_queue.put(None)
 
-            self._trigger_observation_queue_task.cancel()
-            with contextlib.suppress(asyncio.CancelledError):
-                await self._trigger_observation_queue_task
+            await cancel_wait_task(self._trigger_observation_queue_task, max_delay=None)
             self._trigger_observation_queue_task = None
             self._trigger_observation_queue = Queue()