From ee16811a2cd8bdae8796036d455073c9853a5bee Mon Sep 17 00:00:00 2001 From: Pedro Crespo-Valero <32402063+pcrespov@users.noreply.github.com> Date: Mon, 23 Jun 2025 17:50:36 +0200 Subject: [PATCH 01/13] =?UTF-8?q?=E2=9C=A8=20[Async]=20Add=20cancel=5Fand?= =?UTF-8?q?=5Fwait=20function=20to=20manage=20task=20cancellation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/common_library/async_tools.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/packages/common-library/src/common_library/async_tools.py b/packages/common-library/src/common_library/async_tools.py index 205de066851..262079dff18 100644 --- a/packages/common-library/src/common_library/async_tools.py +++ b/packages/common-library/src/common_library/async_tools.py @@ -62,3 +62,22 @@ async def maybe_await( return await obj assert not isawaitable(obj) # nosec return obj + + +async def cancel_and_wait(task: asyncio.Task) -> None: + """Cancels the given task and waits for it to finish. + + Accounts for the case where the parent function is being cancelled + and the task is cancelled as a result. In that case, it suppresses the + `asyncio.CancelledError` if the task was cancelled, but propagates it + if the task was not cancelled (i.e., it was still running when the parent + function was cancelled). + """ + task.cancel() + try: + await asyncio.shield(task) + except asyncio.CancelledError: + if not task.cancelled(): + # parent function is being cancelled -> propagate cancellation + raise + # else: task was cancelled, suppress From c1320da609c237c64c66d9c485c1cec303710ab6 Mon Sep 17 00:00:00 2001 From: Pedro Crespo-Valero <32402063+pcrespov@users.noreply.github.com> Date: Mon, 23 Jun 2025 17:53:22 +0200 Subject: [PATCH 02/13] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20Refactor:=20Replace?= =?UTF-8?q?=20manual=20task=20cancellation=20with=20cancel=5Fand=5Fwait=20?= =?UTF-8?q?in=20garbage=20collector=20tasks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../garbage_collector/_tasks_api_keys.py | 7 ++----- .../garbage_collector/_tasks_core.py | 17 +++++++---------- .../garbage_collector/_tasks_trash.py | 9 +++------ .../garbage_collector/_tasks_users.py | 7 ++----- .../payments/_tasks.py | 7 ++----- 5 files changed, 16 insertions(+), 31 deletions(-) diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py index b992d25b387..6f4d690eca3 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py @@ -8,6 +8,7 @@ from collections.abc import AsyncIterator, Callable from aiohttp import web +from common_library.async_utils import cancel_and_wait from tenacity import retry from tenacity.before_sleep import before_sleep_log from tenacity.wait import wait_exponential @@ -67,10 +68,6 @@ async def _cleanup_ctx_fun( yield # tear-down - task.cancel() - try: - await task - except asyncio.CancelledError: - assert task.cancelled() # nosec + await cancel_and_wait(task) return _cleanup_ctx_fun diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_core.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_core.py index dfb7237d97f..44ea83ffd3c 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_core.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_core.py @@ -1,4 +1,4 @@ -""" Setup and running of periodic background task +"""Setup and running of periodic background task Specifics of the gc implementation should go into garbage_collector_core.py @@ -9,6 +9,7 @@ from collections.abc import AsyncGenerator from aiohttp import web +from common_library.async_tools import cancel_and_wait from servicelib.logging_utils import log_context from ._core import collect_garbage @@ -49,18 +50,14 @@ async def run_background_task(app: web.Application) -> AsyncGenerator: # TEAR-DOWN ----- # controlled cancelation of the gc task - try: - _logger.info("Stopping garbage collector...") + _logger.info("Stopping garbage collector...") - ack = gc_bg_task.cancel() - assert ack # nosec + ack = gc_bg_task.cancel() + assert ack # nosec - app[_GC_TASK_CONFIG]["force_stop"] = True + app[_GC_TASK_CONFIG]["force_stop"] = True - await gc_bg_task - - except asyncio.CancelledError: - assert gc_bg_task.cancelled() # nosec + await cancel_and_wait(gc_bg_task) async def _collect_garbage_periodically(app: web.Application): diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py index 46df72c0a70..8355c2666b6 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py @@ -1,5 +1,5 @@ """ - Scheduled tasks addressing users +Scheduled tasks addressing users """ @@ -8,6 +8,7 @@ from collections.abc import AsyncIterator, Callable from aiohttp import web +from common_library.async_utils import cancel_and_wait from servicelib.logging_utils import log_context from tenacity import retry from tenacity.before_sleep import before_sleep_log @@ -55,10 +56,6 @@ async def _cleanup_ctx_fun( yield # tear-down - task.cancel() - try: - await task - except asyncio.CancelledError: - assert task.cancelled() # nosec + await cancel_and_wait(task) return _cleanup_ctx_fun diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py index 3b834c71ab7..dc64a1b3060 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py @@ -8,6 +8,7 @@ from collections.abc import AsyncIterator, Callable from aiohttp import web +from common_library.async_utils import cancel_and_wait from models_library.users import UserID from servicelib.logging_utils import get_log_record_extra, log_context from tenacity import retry @@ -107,10 +108,6 @@ async def _cleanup_ctx_fun( yield # tear-down - task.cancel() - try: - await task - except asyncio.CancelledError: - assert task.cancelled() # nosec + await cancel_and_wait(task) return _cleanup_ctx_fun diff --git a/services/web/server/src/simcore_service_webserver/payments/_tasks.py b/services/web/server/src/simcore_service_webserver/payments/_tasks.py index b87465f5f3e..067a20531ca 100644 --- a/services/web/server/src/simcore_service_webserver/payments/_tasks.py +++ b/services/web/server/src/simcore_service_webserver/payments/_tasks.py @@ -5,6 +5,7 @@ from typing import Any from aiohttp import web +from common_library.async_utils import cancel_and_wait from models_library.api_schemas_webserver.wallets import PaymentID, PaymentMethodID from pydantic import HttpUrl, TypeAdapter from servicelib.aiohttp.typing_extension import CleanupContextFunc @@ -143,10 +144,6 @@ async def _cleanup_ctx_fun( yield # tear-down - task.cancel() - try: - await task - except asyncio.CancelledError: - assert task.cancelled() # nosec + await cancel_and_wait(task) return _cleanup_ctx_fun From 2732448f4ce76402e7ffccb06dc770b29c7152fb Mon Sep 17 00:00:00 2001 From: Pedro Crespo-Valero <32402063+pcrespov@users.noreply.github.com> Date: Mon, 23 Jun 2025 17:53:31 +0200 Subject: [PATCH 03/13] adds draft tests --- .../common-library/tests/test_async_tools.py | 56 ++++++++++++++++++- 1 file changed, 55 insertions(+), 1 deletion(-) diff --git a/packages/common-library/tests/test_async_tools.py b/packages/common-library/tests/test_async_tools.py index 850945d39b2..f16d98d3080 100644 --- a/packages/common-library/tests/test_async_tools.py +++ b/packages/common-library/tests/test_async_tools.py @@ -3,7 +3,7 @@ from typing import Any import pytest -from common_library.async_tools import make_async, maybe_await +from common_library.async_tools import cancel_and_wait, make_async, maybe_await @make_async() @@ -93,3 +93,57 @@ def fetchone(self) -> Any: # pylint: disable=no-self-use sync_result = await maybe_await(SyncResultProxy().fetchone()) assert sync_result == {"id": 2, "name": "test2"} + + +async def test_cancel_and_wait(): + state = {"started": False, "cancelled": False, "cleaned_up": False} + + async def coro(): + try: + state["started"] = True + await asyncio.sleep(5) + except asyncio.CancelledError: + state["cancelled"] = True + raise + finally: + state["cleaned_up"] = True + + task = asyncio.create_task(coro()) + await asyncio.sleep(0.1) # Let coro start + + await cancel_and_wait(task) + + assert task.done() + assert task.cancelled() + assert state["started"] + assert state["cancelled"] + assert state["cleaned_up"] + + +async def test_cancel_and_wait_propagates_external_cancel(): + """ + This test ensures that if the caller of cancel_and_wait is cancelled, + the CancelledError is not swallowed. + """ + + async def inner_coro(): + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + await asyncio.sleep(0.1) # simulate cleanup + raise + + task = asyncio.create_task(inner_coro()) + + async def outer_coro(): + await cancel_and_wait(task) + + # Cancel the wrapper after a short delay + outer_task = asyncio.create_task(outer_coro()) + await asyncio.sleep(0.1) + outer_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await outer_task + + assert task.cancelled() From a8b339e59651dfe2c8ecf70bd9498e295c800fe4 Mon Sep 17 00:00:00 2001 From: Pedro Crespo-Valero <32402063+pcrespov@users.noreply.github.com> Date: Mon, 23 Jun 2025 18:17:47 +0200 Subject: [PATCH 04/13] fixes tests --- .../src/common_library/async_tools.py | 19 +++++++----- .../common-library/tests/test_async_tools.py | 30 ++++++++++++++----- 2 files changed, 35 insertions(+), 14 deletions(-) diff --git a/packages/common-library/src/common_library/async_tools.py b/packages/common-library/src/common_library/async_tools.py index 262079dff18..439c1f7f455 100644 --- a/packages/common-library/src/common_library/async_tools.py +++ b/packages/common-library/src/common_library/async_tools.py @@ -1,10 +1,13 @@ import asyncio import functools +import logging from collections.abc import Awaitable, Callable from concurrent.futures import Executor from inspect import isawaitable from typing import ParamSpec, TypeVar, overload +_logger = logging.getLogger(__name__) + R = TypeVar("R") P = ParamSpec("P") @@ -67,17 +70,19 @@ async def maybe_await( async def cancel_and_wait(task: asyncio.Task) -> None: """Cancels the given task and waits for it to finish. - Accounts for the case where the parent function is being cancelled - and the task is cancelled as a result. In that case, it suppresses the - `asyncio.CancelledError` if the task was cancelled, but propagates it - if the task was not cancelled (i.e., it was still running when the parent - function was cancelled). + Accounts for the case where the tasks's owner function is being cancelled """ task.cancel() try: + # NOTE shield ensures that cancellation of the caller function won’t stop you + # from observing the cancellation/finalization of task. await asyncio.shield(task) except asyncio.CancelledError: if not task.cancelled(): - # parent function is being cancelled -> propagate cancellation + # task owner function is being cancelled -> propagate cancellation raise - # else: task was cancelled, suppress + # else: task cancellation is complete, we can safely ignore it + _logger.debug( + "Task %s cancellation is complete", + task.get_name(), + ) diff --git a/packages/common-library/tests/test_async_tools.py b/packages/common-library/tests/test_async_tools.py index f16d98d3080..99d41a9a924 100644 --- a/packages/common-library/tests/test_async_tools.py +++ b/packages/common-library/tests/test_async_tools.py @@ -13,7 +13,8 @@ def sync_function(x: int, y: int) -> int: @make_async() def sync_function_with_exception() -> None: - raise ValueError("This is an error!") + msg = "This is an error!" + raise ValueError(msg) @pytest.mark.asyncio @@ -126,17 +127,23 @@ async def test_cancel_and_wait_propagates_external_cancel(): the CancelledError is not swallowed. """ - async def inner_coro(): + async def coro(): try: - await asyncio.sleep(10) + await asyncio.sleep(4) except asyncio.CancelledError: - await asyncio.sleep(0.1) # simulate cleanup + await asyncio.sleep(1) # simulate cleanup raise - task = asyncio.create_task(inner_coro()) + inner_task = asyncio.create_task(coro()) async def outer_coro(): - await cancel_and_wait(task) + try: + await cancel_and_wait(inner_task) + except asyncio.CancelledError: + assert ( + not inner_task.cancelled() + ), "Internal Task should not be cancelled yet (shielded)" + raise # Cancel the wrapper after a short delay outer_task = asyncio.create_task(outer_coro()) @@ -146,4 +153,13 @@ async def outer_coro(): with pytest.raises(asyncio.CancelledError): await outer_task - assert task.cancelled() + # Ensure the task was cancelled + assert inner_task.cancelled() is False, "Task should not be cancelled initially" + + done_event = asyncio.Event() + + def on_done(_): + done_event.set() + + inner_task.add_done_callback(on_done) + await done_event.wait() From 33679b824ef738a2f75633cb8c58b6b506e01388 Mon Sep 17 00:00:00 2001 From: Pedro Crespo-Valero <32402063+pcrespov@users.noreply.github.com> Date: Mon, 23 Jun 2025 18:21:21 +0200 Subject: [PATCH 05/13] wrong imports --- .../garbage_collector/_tasks_api_keys.py | 2 +- .../simcore_service_webserver/garbage_collector/_tasks_trash.py | 2 +- .../simcore_service_webserver/garbage_collector/_tasks_users.py | 2 +- .../web/server/src/simcore_service_webserver/payments/_tasks.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py index 6f4d690eca3..a3d7d88b42c 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py @@ -8,7 +8,7 @@ from collections.abc import AsyncIterator, Callable from aiohttp import web -from common_library.async_utils import cancel_and_wait +from common_library.async_tools import cancel_and_wait from tenacity import retry from tenacity.before_sleep import before_sleep_log from tenacity.wait import wait_exponential diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py index 8355c2666b6..5eb47d55e1f 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py @@ -8,7 +8,7 @@ from collections.abc import AsyncIterator, Callable from aiohttp import web -from common_library.async_utils import cancel_and_wait +from common_library.async_tools import cancel_and_wait from servicelib.logging_utils import log_context from tenacity import retry from tenacity.before_sleep import before_sleep_log diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py index dc64a1b3060..0876d28ad23 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py @@ -8,7 +8,7 @@ from collections.abc import AsyncIterator, Callable from aiohttp import web -from common_library.async_utils import cancel_and_wait +from common_library.async_tools import cancel_and_wait from models_library.users import UserID from servicelib.logging_utils import get_log_record_extra, log_context from tenacity import retry diff --git a/services/web/server/src/simcore_service_webserver/payments/_tasks.py b/services/web/server/src/simcore_service_webserver/payments/_tasks.py index 067a20531ca..d385eb47314 100644 --- a/services/web/server/src/simcore_service_webserver/payments/_tasks.py +++ b/services/web/server/src/simcore_service_webserver/payments/_tasks.py @@ -5,7 +5,7 @@ from typing import Any from aiohttp import web -from common_library.async_utils import cancel_and_wait +from common_library.async_tools import cancel_and_wait from models_library.api_schemas_webserver.wallets import PaymentID, PaymentMethodID from pydantic import HttpUrl, TypeAdapter from servicelib.aiohttp.typing_extension import CleanupContextFunc From 6af9144073557b48232666d107d1a0cfa459ebe5 Mon Sep 17 00:00:00 2001 From: Pedro Crespo-Valero <32402063+pcrespov@users.noreply.github.com> Date: Mon, 23 Jun 2025 18:24:07 +0200 Subject: [PATCH 06/13] doc --- packages/common-library/src/common_library/async_tools.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/packages/common-library/src/common_library/async_tools.py b/packages/common-library/src/common_library/async_tools.py index 439c1f7f455..523a128e61c 100644 --- a/packages/common-library/src/common_library/async_tools.py +++ b/packages/common-library/src/common_library/async_tools.py @@ -68,19 +68,22 @@ async def maybe_await( async def cancel_and_wait(task: asyncio.Task) -> None: - """Cancels the given task and waits for it to finish. + """Cancels the given task and waits for it to complete. - Accounts for the case where the tasks's owner function is being cancelled + 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 # from observing the cancellation/finalization of task. await asyncio.shield(task) + except asyncio.CancelledError: if not task.cancelled(): # task owner function is being cancelled -> propagate cancellation raise + # else: task cancellation is complete, we can safely ignore it _logger.debug( "Task %s cancellation is complete", From 1674bed13535df95948dbc8b09fe5de4dbe5ca48 Mon Sep 17 00:00:00 2001 From: Pedro Crespo-Valero <32402063+pcrespov@users.noreply.github.com> Date: Mon, 23 Jun 2025 18:28:11 +0200 Subject: [PATCH 07/13] doc --- packages/common-library/tests/test_async_tools.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/packages/common-library/tests/test_async_tools.py b/packages/common-library/tests/test_async_tools.py index 99d41a9a924..826d0d259d0 100644 --- a/packages/common-library/tests/test_async_tools.py +++ b/packages/common-library/tests/test_async_tools.py @@ -1,4 +1,5 @@ import asyncio +import time from concurrent.futures import ThreadPoolExecutor from typing import Any @@ -98,11 +99,12 @@ def fetchone(self) -> Any: # pylint: disable=no-self-use async def test_cancel_and_wait(): state = {"started": False, "cancelled": False, "cleaned_up": False} + SLEEP_TIME = 5 # seconds async def coro(): try: state["started"] = True - await asyncio.sleep(5) + await asyncio.sleep(SLEEP_TIME) except asyncio.CancelledError: state["cancelled"] = True raise @@ -112,8 +114,11 @@ async def coro(): task = asyncio.create_task(coro()) await asyncio.sleep(0.1) # Let coro start + start = time.time() await cancel_and_wait(task) + elapsed = time.time() - start + assert elapsed < SLEEP_TIME, "Task should be cancelled quickly" assert task.done() assert task.cancelled() assert state["started"] @@ -142,7 +147,7 @@ async def outer_coro(): except asyncio.CancelledError: assert ( not inner_task.cancelled() - ), "Internal Task should not be cancelled yet (shielded)" + ), "Internal Task DOES NOT RAISE CancelledError" raise # Cancel the wrapper after a short delay From 2371c6432a89ff630d996428d61dfe0865cb0e7d Mon Sep 17 00:00:00 2001 From: Pedro Crespo-Valero <32402063+pcrespov@users.noreply.github.com> Date: Mon, 23 Jun 2025 18:33:34 +0200 Subject: [PATCH 08/13] @sanderegg review: using it in cancel_wait_task --- packages/service-library/src/servicelib/async_utils.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/packages/service-library/src/servicelib/async_utils.py b/packages/service-library/src/servicelib/async_utils.py index c6466df0a70..1c07a60e7a0 100644 --- a/packages/service-library/src/servicelib/async_utils.py +++ b/packages/service-library/src/servicelib/async_utils.py @@ -1,5 +1,4 @@ import asyncio -import contextlib import datetime import logging from collections import deque @@ -9,6 +8,8 @@ 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 @@ -244,8 +245,5 @@ async def cancel_wait_task( up the cancellation. If None it waits forever. :raises TimeoutError: raised if cannot cancel the task. """ - - task.cancel() async with asyncio.timeout(max_delay): - with contextlib.suppress(asyncio.CancelledError): - await task + await cancel_and_wait(task) From c826c8f9f49b47b5ac59351f57700d17b1b886a9 Mon Sep 17 00:00:00 2001 From: Pedro Crespo-Valero <32402063+pcrespov@users.noreply.github.com> Date: Tue, 24 Jun 2025 13:29:05 +0200 Subject: [PATCH 09/13] @sanderegg review: replace cancel_wait_task by cancel_and_wait --- .../celery-library/src/celery_library/task.py | 4 +-- .../src/common_library/async_tools.py | 33 +++++++++++++++---- .../common-library/tests/test_async_tools.py | 25 ++++++++++++++ .../src/servicelib/async_utils.py | 18 ---------- .../src/servicelib/background_task.py | 4 +-- .../src/servicelib/long_running_tasks/task.py | 4 +-- .../src/servicelib/redis/_client.py | 4 +-- ...test_rabbitmq_rpc_interfaces_async_jobs.py | 4 +-- .../tests/redis/test_project_lock.py | 4 +-- .../tests/test_background_task.py | 4 +-- .../tests/test_background_task_utils.py | 4 +-- .../services/volumes_manager.py | 6 ++-- .../core/_prometheus_instrumentation.py | 4 +-- .../core/health_checker.py | 6 ++-- .../modules/auto_scaling_task.py | 4 +-- .../modules/buffer_machines_pool_task.py | 4 +-- .../modules/clusters_management_task.py | 4 +-- .../rabbitmq_worker_plugin.py | 4 +-- .../modules/comp_scheduler/_manager.py | 4 +-- .../scheduler/_core/_scheduler.py | 6 ++-- .../utils/base_distributed_identifier.py | 4 +-- .../registry_proxy.py | 4 +-- .../services/status_monitor/_monitor.py | 10 +++--- .../modules/outputs/_manager.py | 4 +-- .../modules/prometheus_metrics.py | 4 +-- .../modules/resource_tracking/_core.py | 4 +-- .../modules/system_monitor/_disk_usage.py | 4 +-- .../services/background_tasks_setup.py | 4 +-- .../services/fire_and_forget_setup.py | 4 +-- .../clients/postgres/_liveness.py | 4 +-- ...und_task_periodic_heartbeat_check_setup.py | 4 +-- .../services/fire_and_forget_setup.py | 4 +-- .../simcore_service_storage/dsm_cleaner.py | 4 +-- .../simcore_service_storage/utils/s3_utils.py | 4 +-- .../licenses/_itis_vip_syncer_service.py | 4 +-- 35 files changed, 122 insertions(+), 94 deletions(-) diff --git a/packages/celery-library/src/celery_library/task.py b/packages/celery-library/src/celery_library/task.py index f14771cf207..4eb37be5aa7 100644 --- a/packages/celery-library/src/celery_library/task.py +++ b/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(), ) diff --git a/packages/common-library/src/common_library/async_tools.py b/packages/common-library/src/common_library/async_tools.py index 523a128e61c..e9ffb1b914d 100644 --- a/packages/common-library/src/common_library/async_tools.py +++ b/packages/common-library/src/common_library/async_tools.py @@ -67,17 +67,38 @@ 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. - Accounts for the case where the tasks's owner function is being cancelled. - In that case, it propagates the cancellation exception upstream. + Arguments: + task -- task to be canceled + + + Keyword Arguments: + max_delay -- duration (in seconds) to wait before giving + up the cancellation. This timeout should be an upper bound to the + time needed for the task to cleanup after being canceled and + avoids that the cancellation takes forever. If None the timeout is not + set. (default: {None}) + + Raises: + TimeoutError: raised if cannot cancel the task. + CancelledError: raised ONLY if owner is being cancelled. """ + task.cancel() + assert task.cancelling() # nosec + assert not task.cancelled() # nosec + try: - # 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( + # NOTE shield ensures that cancellation of the caller function won't stop you + # from observing the cancellation/finalization of task. + asyncio.wait_for(task, timeout=max_delay) + ) except asyncio.CancelledError: if not task.cancelled(): diff --git a/packages/common-library/tests/test_async_tools.py b/packages/common-library/tests/test_async_tools.py index 826d0d259d0..2639ffd4ee3 100644 --- a/packages/common-library/tests/test_async_tools.py +++ b/packages/common-library/tests/test_async_tools.py @@ -168,3 +168,28 @@ def on_done(_): inner_task.add_done_callback(on_done) await done_event.wait() + + +async def test_cancel_and_wait_timeout_on_slow_cleanup(): + """Test that cancel_and_wait raises TimeoutError when cleanup takes longer than max_delay""" + + CLEANUP_TIME = 2 # seconds + + async def slow_cleanup_coro(): + try: + await asyncio.sleep(10) # Long running task + except asyncio.CancelledError: + # Simulate slow cleanup that exceeds max_delay! + await asyncio.sleep(CLEANUP_TIME) + raise + + task = asyncio.create_task(slow_cleanup_coro()) + await asyncio.sleep(0.1) # Let the task start + + # Cancel with a max_delay shorter than cleanup time + with pytest.raises(TimeoutError): + await cancel_and_wait( + task, max_delay=CLEANUP_TIME / 10 + ) # 0.2 seconds < 2 seconds cleanup + + assert task.cancelled() diff --git a/packages/service-library/src/servicelib/async_utils.py b/packages/service-library/src/servicelib/async_utils.py index 1c07a60e7a0..b4aa978ec3a 100644 --- a/packages/service-library/src/servicelib/async_utils.py +++ b/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) diff --git a/packages/service-library/src/servicelib/background_task.py b/packages/service-library/src/servicelib/background_task.py index 508f34b99ee..94f62b50169 100644 --- a/packages/service-library/src/servicelib/background_task.py +++ b/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)) diff --git a/packages/service-library/src/servicelib/long_running_tasks/task.py b/packages/service-library/src/servicelib/long_running_tasks/task.py index c89e94ce476..71978bb18f1 100644 --- a/packages/service-library/src/servicelib/long_running_tasks/task.py +++ b/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 ) diff --git a/packages/service-library/src/servicelib/redis/_client.py b/packages/service-library/src/servicelib/redis/_client.py index c2a08154110..e59f1d21252 100644 --- a/packages/service-library/src/servicelib/redis/_client.py +++ b/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 ) diff --git a/packages/service-library/tests/rabbitmq/test_rabbitmq_rpc_interfaces_async_jobs.py b/packages/service-library/tests/rabbitmq/test_rabbitmq_rpc_interfaces_async_jobs.py index 72ecc9a8aa6..68d55f239ec 100644 --- a/packages/service-library/tests/rabbitmq/test_rabbitmq_rpc_interfaces_async_jobs.py +++ b/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"]) diff --git a/packages/service-library/tests/redis/test_project_lock.py b/packages/service-library/tests/redis/test_project_lock.py index aa9d7fd1c74..04bfa99f176 100644 --- a/packages/service-library/tests/redis/test_project_lock.py +++ b/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) diff --git a/packages/service-library/tests/test_background_task.py b/packages/service-library/tests/test_background_task.py index 8c508bf8979..958f15b7f9c 100644 --- a/packages/service-library/tests/test_background_task.py +++ b/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) ) diff --git a/packages/service-library/tests/test_background_task_utils.py b/packages/service-library/tests/test_background_task_utils.py index 9a03a6c3541..c703592146b 100644 --- a/packages/service-library/tests/test_background_task_utils.py +++ b/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 diff --git a/services/agent/src/simcore_service_agent/services/volumes_manager.py b/services/agent/src/simcore_service_agent/services/volumes_manager.py index 860ab86d0e2..bba3b8178b8 100644 --- a/services/agent/src/simcore_service_agent/services/volumes_manager.py +++ b/services/agent/src/simcore_service_agent/services/volumes_manager.py @@ -6,10 +6,10 @@ import arrow from aiodocker.docker import Docker +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI from models_library.projects_nodes_io import NodeID from pydantic import NonNegativeFloat -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.fastapi.app_state import SingletonInAppStateMixin from servicelib.logging_utils import log_context @@ -61,10 +61,10 @@ async def shutdown(self) -> None: await self.docker.close() if self._task_bookkeeping: - await cancel_wait_task(self._task_bookkeeping) + await cancel_and_wait(self._task_bookkeeping) if self._task_periodic_volume_cleanup: - await cancel_wait_task(self._task_periodic_volume_cleanup) + await cancel_and_wait(self._task_periodic_volume_cleanup) async def _bookkeeping_task(self) -> None: with log_context(_logger, logging.DEBUG, "volume bookkeeping"): diff --git a/services/api-server/src/simcore_service_api_server/core/_prometheus_instrumentation.py b/services/api-server/src/simcore_service_api_server/core/_prometheus_instrumentation.py index f19bac34a76..995900a3a77 100644 --- a/services/api-server/src/simcore_service_api_server/core/_prometheus_instrumentation.py +++ b/services/api-server/src/simcore_service_api_server/core/_prometheus_instrumentation.py @@ -4,10 +4,10 @@ from datetime import timedelta from typing import Final, cast +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI from prometheus_client import CollectorRegistry, Gauge from pydantic import PositiveInt -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.fastapi.monitoring import ( setup_prometheus_instrumentation as setup_rest_instrumentation, @@ -85,7 +85,7 @@ async def on_startup() -> None: async def on_shutdown() -> None: assert app.state.instrumentation_task # nosec with log_catch(_logger, reraise=False): - await cancel_wait_task(app.state.instrumentation_task) + await cancel_and_wait(app.state.instrumentation_task) app.add_event_handler("startup", on_startup) app.add_event_handler("shutdown", on_shutdown) diff --git a/services/api-server/src/simcore_service_api_server/core/health_checker.py b/services/api-server/src/simcore_service_api_server/core/health_checker.py index b5a5180b12b..3d99c38b562 100644 --- a/services/api-server/src/simcore_service_api_server/core/health_checker.py +++ b/services/api-server/src/simcore_service_api_server/core/health_checker.py @@ -5,11 +5,11 @@ from typing import Annotated, Final, cast from uuid import uuid4 +from common_library.async_tools import cancel_and_wait from fastapi import Depends, FastAPI from models_library.rabbitmq_messages import LoggerRabbitMessage from models_library.users import UserID from pydantic import NonNegativeInt, PositiveFloat, PositiveInt -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.fastapi.dependencies import get_app from servicelib.logging_utils import log_catch @@ -63,7 +63,7 @@ async def setup(self, health_check_task_period_seconds: PositiveFloat): async def teardown(self): if self._background_task: with log_catch(_logger, reraise=False): - await cancel_wait_task( + await cancel_and_wait( self._background_task, max_delay=self._timeout_seconds ) await self._log_distributor.deregister(job_id=self._dummy_job_id) @@ -95,7 +95,7 @@ async def _background_task_method(self): self._dummy_queue.get(), timeout=self._timeout_seconds ) self._health_check_failure_count = 0 - except asyncio.TimeoutError: + except TimeoutError: self._increment_health_check_failure_count() diff --git a/services/autoscaling/src/simcore_service_autoscaling/modules/auto_scaling_task.py b/services/autoscaling/src/simcore_service_autoscaling/modules/auto_scaling_task.py index 5ebc6a190f8..756c716a2c2 100644 --- a/services/autoscaling/src/simcore_service_autoscaling/modules/auto_scaling_task.py +++ b/services/autoscaling/src/simcore_service_autoscaling/modules/auto_scaling_task.py @@ -2,8 +2,8 @@ from collections.abc import Awaitable, Callable from typing import Final +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.redis import exclusive @@ -44,7 +44,7 @@ async def _startup() -> None: def on_app_shutdown(app: FastAPI) -> Callable[[], Awaitable[None]]: async def _stop() -> None: - await cancel_wait_task(app.state.autoscaler_task) + await cancel_and_wait(app.state.autoscaler_task) return _stop diff --git a/services/autoscaling/src/simcore_service_autoscaling/modules/buffer_machines_pool_task.py b/services/autoscaling/src/simcore_service_autoscaling/modules/buffer_machines_pool_task.py index 2985e2ffcc4..42dbd99c1f0 100644 --- a/services/autoscaling/src/simcore_service_autoscaling/modules/buffer_machines_pool_task.py +++ b/services/autoscaling/src/simcore_service_autoscaling/modules/buffer_machines_pool_task.py @@ -2,8 +2,8 @@ from collections.abc import Awaitable, Callable from typing import Final +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.redis import exclusive @@ -44,7 +44,7 @@ async def _startup() -> None: def on_app_shutdown(app: FastAPI) -> Callable[[], Awaitable[None]]: async def _stop() -> None: if hasattr(app.state, "buffers_pool_task"): - await cancel_wait_task(app.state.buffers_pool_task) + await cancel_and_wait(app.state.buffers_pool_task) return _stop diff --git a/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters_management_task.py b/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters_management_task.py index c540d7b160f..5c9affe3e60 100644 --- a/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters_management_task.py +++ b/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters_management_task.py @@ -2,8 +2,8 @@ import logging from collections.abc import Awaitable, Callable +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.redis import exclusive @@ -37,7 +37,7 @@ async def _startup() -> None: def on_app_shutdown(app: FastAPI) -> Callable[[], Awaitable[None]]: async def _stop() -> None: - await cancel_wait_task(app.state.clusters_cleaning_task, max_delay=5) + await cancel_and_wait(app.state.clusters_cleaning_task, max_delay=5) return _stop diff --git a/services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq_worker_plugin.py b/services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq_worker_plugin.py index ba4936284d7..90458fbc977 100644 --- a/services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq_worker_plugin.py +++ b/services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq_worker_plugin.py @@ -6,7 +6,7 @@ from typing import Final import distributed -from servicelib.async_utils import cancel_wait_task +from common_library.async_tools import cancel_and_wait from servicelib.logging_utils import log_catch, log_context from servicelib.rabbitmq import RabbitMQClient, wait_till_rabbitmq_responsive from servicelib.rabbitmq._models import RabbitMessage @@ -104,7 +104,7 @@ async def _() -> None: # Cancel the message processor task if self._message_processor: with log_catch(_logger, reraise=False): - await cancel_wait_task(self._message_processor, max_delay=5) + await cancel_and_wait(self._message_processor, max_delay=5) self._message_processor = None # close client diff --git a/services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_manager.py b/services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_manager.py index 09478d4d02e..0287e92f402 100644 --- a/services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_manager.py +++ b/services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_manager.py @@ -2,10 +2,10 @@ from typing import Final import networkx as nx +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI from models_library.projects import ProjectID from models_library.users import UserID -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.exception_utils import suppress_exceptions from servicelib.logging_utils import log_context @@ -202,4 +202,4 @@ async def setup_manager(app: FastAPI) -> None: async def shutdown_manager(app: FastAPI) -> None: - await cancel_wait_task(app.state.scheduler_manager) + await cancel_and_wait(app.state.scheduler_manager) diff --git a/services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_scheduler.py b/services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_scheduler.py index ea36680240d..41467eebe19 100644 --- a/services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_scheduler.py +++ b/services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_scheduler.py @@ -23,6 +23,7 @@ from typing import Final import arrow +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI from models_library.api_schemas_directorv2.dynamic_services import ( DynamicServiceCreate, @@ -41,7 +42,6 @@ from models_library.users import UserID from models_library.wallets import WalletID from pydantic import NonNegativeFloat -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.long_running_tasks.models import ProgressCallback, TaskProgress from servicelib.redis import RedisClientsManager, exclusive @@ -125,7 +125,7 @@ async def shutdown(self) -> None: self._to_observe = {} if self._scheduler_task is not None: - await cancel_wait_task(self._scheduler_task, max_delay=5) + await cancel_and_wait(self._scheduler_task, max_delay=5) self._scheduler_task = None if self._trigger_observation_queue_task is not None: @@ -364,7 +364,7 @@ async def mark_service_for_removal( self._service_observation_task[service_name] ) if isinstance(service_task, asyncio.Task): - await cancel_wait_task(service_task, max_delay=10) + await cancel_and_wait(service_task, max_delay=10) if skip_observation_recreation: return diff --git a/services/director-v2/src/simcore_service_director_v2/utils/base_distributed_identifier.py b/services/director-v2/src/simcore_service_director_v2/utils/base_distributed_identifier.py index ea685777a0d..7a94c4e71fa 100644 --- a/services/director-v2/src/simcore_service_director_v2/utils/base_distributed_identifier.py +++ b/services/director-v2/src/simcore_service_director_v2/utils/base_distributed_identifier.py @@ -4,8 +4,8 @@ from datetime import timedelta from typing import Final, Generic, TypeVar +from common_library.async_tools import cancel_and_wait from pydantic import NonNegativeInt -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.logging_utils import log_catch, log_context from servicelib.redis import RedisClientSDK @@ -76,7 +76,7 @@ async def setup(self) -> None: async def shutdown(self) -> None: if self._cleanup_task: - await cancel_wait_task(self._cleanup_task, max_delay=5) + await cancel_and_wait(self._cleanup_task, max_delay=5) @classmethod def class_path(cls) -> str: diff --git a/services/director/src/simcore_service_director/registry_proxy.py b/services/director/src/simcore_service_director/registry_proxy.py index 56b5d812f8c..964560e701d 100644 --- a/services/director/src/simcore_service_director/registry_proxy.py +++ b/services/director/src/simcore_service_director/registry_proxy.py @@ -8,9 +8,9 @@ import httpx from aiocache import Cache, SimpleMemoryCache # type: ignore[import-untyped] +from common_library.async_tools import cancel_and_wait from common_library.json_serialization import json_loads from fastapi import FastAPI, status -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.fastapi.client_session import get_client_session from servicelib.logging_utils import log_catch, log_context @@ -255,7 +255,7 @@ async def on_startup() -> None: async def on_shutdown() -> None: if app.state.auto_cache_task: - await cancel_wait_task(app.state.auto_cache_task) + await cancel_and_wait(app.state.auto_cache_task) app.add_event_handler("startup", on_startup) app.add_event_handler("shutdown", on_shutdown) diff --git a/services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/status_monitor/_monitor.py b/services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/status_monitor/_monitor.py index 750b0dbdc63..6107474eb33 100644 --- a/services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/status_monitor/_monitor.py +++ b/services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/status_monitor/_monitor.py @@ -5,10 +5,10 @@ from typing import Final import arrow +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI from models_library.projects_nodes_io import NodeID from pydantic import NonNegativeFloat, NonNegativeInt -from servicelib.async_utils import cancel_wait_task from servicelib.background_task_utils import exclusive_periodic from servicelib.utils import limited_gather from settings_library.redis import RedisDatabase @@ -79,9 +79,9 @@ async def _worker_check_services_require_status_update(self) -> None: # NOTE: this worker runs on only once across all instances of the scheduler - models: dict[ - NodeID, TrackedServiceModel - ] = await service_tracker.get_all_tracked_services(self.app) + models: dict[NodeID, TrackedServiceModel] = ( + await service_tracker.get_all_tracked_services(self.app) + ) to_remove: list[NodeID] = [] to_start: list[NodeID] = [] @@ -149,4 +149,4 @@ async def _periodic_check_services_require_status_update() -> None: async def shutdown(self) -> None: if getattr(self.app.state, "status_monitor_background_task", None): - await cancel_wait_task(self.app.state.status_monitor_background_task) + await cancel_and_wait(self.app.state.status_monitor_background_task) diff --git a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/outputs/_manager.py b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/outputs/_manager.py index f29f26358e2..4459150c6c0 100644 --- a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/outputs/_manager.py +++ b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/outputs/_manager.py @@ -6,12 +6,12 @@ from datetime import timedelta from functools import partial +from common_library.async_tools import cancel_and_wait from common_library.errors_classes import OsparcErrorMixin from fastapi import FastAPI from models_library.rabbitmq_messages import ProgressType from pydantic import PositiveFloat from servicelib import progress_bar -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.logging_utils import log_catch, log_context from simcore_sdk.node_ports_common.file_io_utils import LogRedirectCB @@ -204,7 +204,7 @@ async def shutdown(self) -> None: with log_context(_logger, logging.INFO, f"{OutputsManager.__name__} shutdown"): await self._uploading_task_cancel() if self._task_scheduler_worker is not None: - await cancel_wait_task( + await cancel_and_wait( self._task_scheduler_worker, max_delay=self.task_monitor_interval_s ) diff --git a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/prometheus_metrics.py b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/prometheus_metrics.py index eb7ad93ed9e..348efa5f278 100644 --- a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/prometheus_metrics.py +++ b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/prometheus_metrics.py @@ -6,10 +6,10 @@ from typing import Final import arrow +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI, status from models_library.callbacks_mapping import CallbacksMapping, UserServiceCommand from pydantic import BaseModel, NonNegativeFloat, NonNegativeInt -from servicelib.async_utils import cancel_wait_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_wait_task( + await cancel_and_wait( self._metrics_recovery_task, max_delay=_TASK_CANCELLATION_TIMEOUT_S ) diff --git a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/resource_tracking/_core.py b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/resource_tracking/_core.py index eecbfd2089e..466f9cdf0ba 100644 --- a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/resource_tracking/_core.py +++ b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/resource_tracking/_core.py @@ -2,6 +2,7 @@ import logging from typing import Final +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI from models_library.generated_models.docker_rest_api import ContainerState from models_library.rabbitmq_messages import ( @@ -14,7 +15,6 @@ from models_library.services import ServiceType from models_library.services_creation import CreateServiceMetricsAdditionalParams from pydantic import NonNegativeFloat -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.logging_utils import log_context @@ -62,7 +62,7 @@ async def _start_heart_beat_task(app: FastAPI) -> None: async def stop_heart_beat_task(app: FastAPI) -> None: resource_tracking: ResourceTrackingState = app.state.resource_tracking if resource_tracking.heart_beat_task: - await cancel_wait_task( + await cancel_and_wait( resource_tracking.heart_beat_task, max_delay=_STOP_WORKER_TIMEOUT_S ) diff --git a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/system_monitor/_disk_usage.py b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/system_monitor/_disk_usage.py index d2148842ef5..df8f3990d15 100644 --- a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/system_monitor/_disk_usage.py +++ b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/system_monitor/_disk_usage.py @@ -7,6 +7,7 @@ from typing import Final import psutil +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI from models_library.api_schemas_dynamic_sidecar.telemetry import ( DiskUsage, @@ -14,7 +15,6 @@ ) from models_library.projects_nodes_io import NodeID from models_library.users import UserID -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.logging_utils import log_context from servicelib.utils import logged_gather @@ -185,7 +185,7 @@ async def setup(self) -> None: async def shutdown(self) -> None: if self._monitor_task: - await cancel_wait_task(self._monitor_task) + await cancel_and_wait(self._monitor_task) def set_disk_usage_for_path(self, overwrite_usage: dict[str, DiskUsage]) -> None: """ diff --git a/services/efs-guardian/src/simcore_service_efs_guardian/services/background_tasks_setup.py b/services/efs-guardian/src/simcore_service_efs_guardian/services/background_tasks_setup.py index e1480f84b20..f6602899340 100644 --- a/services/efs-guardian/src/simcore_service_efs_guardian/services/background_tasks_setup.py +++ b/services/efs-guardian/src/simcore_service_efs_guardian/services/background_tasks_setup.py @@ -3,8 +3,8 @@ from collections.abc import Awaitable, Callable from datetime import timedelta +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI -from servicelib.async_utils import cancel_wait_task from servicelib.background_task_utils import exclusive_periodic from servicelib.logging_utils import log_catch, log_context @@ -50,7 +50,7 @@ async def _stop() -> None: ): assert _app # nosec if _app.state.efs_guardian_removal_policy_background_task: - await cancel_wait_task( + await cancel_and_wait( _app.state.efs_guardian_removal_policy_background_task ) diff --git a/services/efs-guardian/src/simcore_service_efs_guardian/services/fire_and_forget_setup.py b/services/efs-guardian/src/simcore_service_efs_guardian/services/fire_and_forget_setup.py index a38411f56a1..94fb8985cd3 100644 --- a/services/efs-guardian/src/simcore_service_efs_guardian/services/fire_and_forget_setup.py +++ b/services/efs-guardian/src/simcore_service_efs_guardian/services/fire_and_forget_setup.py @@ -1,8 +1,8 @@ import logging from collections.abc import Awaitable, Callable +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI -from servicelib.async_utils import cancel_wait_task from servicelib.logging_utils import log_catch, log_context _logger = logging.getLogger(__name__) @@ -28,7 +28,7 @@ async def _stop() -> None: assert _app # nosec if _app.state.efs_guardian_fire_and_forget_tasks: for task in _app.state.efs_guardian_fire_and_forget_tasks: - await cancel_wait_task(task) + await cancel_and_wait(task) return _stop diff --git a/services/notifications/src/simcore_service_notifications/clients/postgres/_liveness.py b/services/notifications/src/simcore_service_notifications/clients/postgres/_liveness.py index 57bc7a40076..51372a1ee48 100644 --- a/services/notifications/src/simcore_service_notifications/clients/postgres/_liveness.py +++ b/services/notifications/src/simcore_service_notifications/clients/postgres/_liveness.py @@ -3,9 +3,9 @@ from datetime import timedelta from typing import Final +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI from models_library.healthchecks import IsResponsive, LivenessResult -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.db_asyncpg_utils import check_postgres_liveness from servicelib.fastapi.db_asyncpg_engine import get_engine @@ -40,4 +40,4 @@ async def setup(self) -> None: async def teardown(self) -> None: if self._task is not None: with log_catch(_logger, reraise=False): - await cancel_wait_task(self._task, max_delay=5) + await cancel_and_wait(self._task, max_delay=5) diff --git a/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/background_task_periodic_heartbeat_check_setup.py b/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/background_task_periodic_heartbeat_check_setup.py index abaefe1e9b7..8287d83a57b 100644 --- a/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/background_task_periodic_heartbeat_check_setup.py +++ b/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/background_task_periodic_heartbeat_check_setup.py @@ -3,8 +3,8 @@ from collections.abc import Awaitable, Callable from typing import TypedDict +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI -from servicelib.async_utils import cancel_wait_task from servicelib.background_task_utils import exclusive_periodic from servicelib.logging_utils import log_catch, log_context @@ -69,7 +69,7 @@ async def _stop() -> None: ): assert _app # nosec if _app.state.rut_background_task__periodic_check_of_running_services: - await cancel_wait_task( + await cancel_and_wait( _app.state.rut_background_task__periodic_check_of_running_services ) diff --git a/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/fire_and_forget_setup.py b/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/fire_and_forget_setup.py index 2523a069974..0c8c8091c02 100644 --- a/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/fire_and_forget_setup.py +++ b/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/fire_and_forget_setup.py @@ -1,8 +1,8 @@ import logging from collections.abc import Awaitable, Callable +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI -from servicelib.async_utils import cancel_wait_task from servicelib.logging_utils import log_catch, log_context _logger = logging.getLogger(__name__) @@ -32,7 +32,7 @@ async def _stop() -> None: assert _app # nosec if _app.state.rut_fire_and_forget_tasks: for task in _app.state.rut_fire_and_forget_tasks: - await cancel_wait_task(task) + await cancel_and_wait(task) return _stop diff --git a/services/storage/src/simcore_service_storage/dsm_cleaner.py b/services/storage/src/simcore_service_storage/dsm_cleaner.py index d09c83e4f5d..45458fef336 100644 --- a/services/storage/src/simcore_service_storage/dsm_cleaner.py +++ b/services/storage/src/simcore_service_storage/dsm_cleaner.py @@ -23,8 +23,8 @@ from datetime import timedelta from typing import cast +from common_library.async_tools import cancel_and_wait from fastapi import FastAPI -from servicelib.async_utils import cancel_wait_task from servicelib.background_task_utils import exclusive_periodic from servicelib.logging_utils import log_context @@ -66,7 +66,7 @@ async def _periodic_dsm_clean() -> None: async def _on_shutdown() -> None: assert isinstance(app.state.dsm_cleaner_task, asyncio.Task) # nosec - await cancel_wait_task(app.state.dsm_cleaner_task) + await cancel_and_wait(app.state.dsm_cleaner_task) app.add_event_handler("startup", _on_startup) app.add_event_handler("shutdown", _on_shutdown) diff --git a/services/storage/src/simcore_service_storage/utils/s3_utils.py b/services/storage/src/simcore_service_storage/utils/s3_utils.py index 3fcb17d0c45..25e0f3020c3 100644 --- a/services/storage/src/simcore_service_storage/utils/s3_utils.py +++ b/services/storage/src/simcore_service_storage/utils/s3_utils.py @@ -4,8 +4,8 @@ from collections import defaultdict from dataclasses import dataclass, field +from common_library.async_tools import cancel_and_wait from pydantic import ByteSize, TypeAdapter -from servicelib.async_utils import cancel_wait_task from servicelib.background_task import create_periodic_task from servicelib.progress_bar import ProgressBarData @@ -39,7 +39,7 @@ async def __aexit__(self, exc_type, exc_value, traceback) -> None: self.finalize_transfer() await asyncio.sleep(0) assert self._async_update_periodic_task # nosec - await cancel_wait_task(self._async_update_periodic_task) + await cancel_and_wait(self._async_update_periodic_task) async def _async_update(self) -> None: await self._update_task_event.wait() diff --git a/services/web/server/src/simcore_service_webserver/licenses/_itis_vip_syncer_service.py b/services/web/server/src/simcore_service_webserver/licenses/_itis_vip_syncer_service.py index 5f04cea6a3c..1b84807a679 100644 --- a/services/web/server/src/simcore_service_webserver/licenses/_itis_vip_syncer_service.py +++ b/services/web/server/src/simcore_service_webserver/licenses/_itis_vip_syncer_service.py @@ -4,9 +4,9 @@ from datetime import timedelta from aiohttp import web +from common_library.async_tools import cancel_and_wait from httpx import AsyncClient from models_library.licenses import LicensedResourceType -from servicelib.async_utils import cancel_wait_task from servicelib.background_task_utils import exclusive_periodic from servicelib.logging_utils import log_catch, log_context @@ -115,6 +115,6 @@ async def _periodic_sync() -> None: yield - await cancel_wait_task(background_task) + await cancel_and_wait(background_task) app.cleanup_ctx.append(_lifespan) From e6654b6a32a9872c7465a73118c17a2c07c3fd20 Mon Sep 17 00:00:00 2001 From: Pedro Crespo-Valero <32402063+pcrespov@users.noreply.github.com> Date: Tue, 24 Jun 2025 13:54:28 +0200 Subject: [PATCH 10/13] @sanderegg review: moving to async_tools --- .../src/common_library/async_tools.py | 24 ++++++++++++++++-- .../common-library/tests/test_async_tools.py | 25 ++++++++++++++++++- .../src/servicelib/async_utils.py | 21 +--------------- .../src/servicelib/background_task.py | 2 +- .../service-library/tests/test_async_utils.py | 19 -------------- 5 files changed, 48 insertions(+), 43 deletions(-) diff --git a/packages/common-library/src/common_library/async_tools.py b/packages/common-library/src/common_library/async_tools.py index e9ffb1b914d..f343ffc4fd4 100644 --- a/packages/common-library/src/common_library/async_tools.py +++ b/packages/common-library/src/common_library/async_tools.py @@ -1,10 +1,12 @@ import asyncio +import datetime import functools import logging -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Coroutine from concurrent.futures import Executor +from functools import wraps from inspect import isawaitable -from typing import ParamSpec, TypeVar, overload +from typing import Any, ParamSpec, TypeVar, overload _logger = logging.getLogger(__name__) @@ -110,3 +112,21 @@ async def cancel_and_wait( "Task %s cancellation is complete", task.get_name(), ) + + +def delayed_start( + delay: datetime.timedelta, +) -> Callable[ + [Callable[P, Coroutine[Any, Any, R]]], Callable[P, Coroutine[Any, Any, R]] +]: + def _decorator( + func: Callable[P, Coroutine[Any, Any, R]], + ) -> Callable[P, Coroutine[Any, Any, R]]: + @wraps(func) + async def _wrapper(*args: P.args, **kwargs: P.kwargs) -> R: + await asyncio.sleep(delay.total_seconds()) + return await func(*args, **kwargs) + + return _wrapper + + return _decorator diff --git a/packages/common-library/tests/test_async_tools.py b/packages/common-library/tests/test_async_tools.py index 2639ffd4ee3..bcebc3438a0 100644 --- a/packages/common-library/tests/test_async_tools.py +++ b/packages/common-library/tests/test_async_tools.py @@ -1,10 +1,16 @@ import asyncio import time from concurrent.futures import ThreadPoolExecutor +from datetime import timedelta from typing import Any import pytest -from common_library.async_tools import cancel_and_wait, make_async, maybe_await +from common_library.async_tools import ( + cancel_and_wait, + delayed_start, + make_async, + maybe_await, +) @make_async() @@ -193,3 +199,20 @@ async def slow_cleanup_coro(): ) # 0.2 seconds < 2 seconds cleanup assert task.cancelled() + + +async def test_with_delay(): + @delayed_start(timedelta(seconds=0.2)) + async def decorated_awaitable() -> int: + return 42 + + assert await decorated_awaitable() == 42 + + async def another_awaitable() -> int: + return 42 + + decorated_another_awaitable = delayed_start(timedelta(seconds=0.2))( + another_awaitable + ) + + assert await decorated_another_awaitable() == 42 diff --git a/packages/service-library/src/servicelib/async_utils.py b/packages/service-library/src/servicelib/async_utils.py index b4aa978ec3a..d2c62ba55ff 100644 --- a/packages/service-library/src/servicelib/async_utils.py +++ b/packages/service-library/src/servicelib/async_utils.py @@ -1,8 +1,7 @@ import asyncio -import datetime import logging from collections import deque -from collections.abc import Awaitable, Callable, Coroutine +from collections.abc import Awaitable, Callable from contextlib import suppress from dataclasses import dataclass from functools import wraps @@ -211,21 +210,3 @@ async def worker(in_q: Queue[QueueElement], out_q: Queue) -> None: return wrapper return decorator - - -def delayed_start( - delay: datetime.timedelta, -) -> Callable[ - [Callable[P, Coroutine[Any, Any, R]]], Callable[P, Coroutine[Any, Any, R]] -]: - def _decorator( - func: Callable[P, Coroutine[Any, Any, R]], - ) -> Callable[P, Coroutine[Any, Any, R]]: - @wraps(func) - async def _wrapper(*args: P.args, **kwargs: P.kwargs) -> R: - await asyncio.sleep(delay.total_seconds()) - return await func(*args, **kwargs) - - return _wrapper - - return _decorator diff --git a/packages/service-library/src/servicelib/background_task.py b/packages/service-library/src/servicelib/background_task.py index 94f62b50169..b37ff0e0c2b 100644 --- a/packages/service-library/src/servicelib/background_task.py +++ b/packages/service-library/src/servicelib/background_task.py @@ -6,10 +6,10 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine from typing import Any, Final, ParamSpec, TypeVar +from common_library.async_tools import cancel_and_wait, delayed_start from tenacity import TryAgain, before_sleep_log, retry, retry_if_exception_type from tenacity.wait import wait_fixed -from .async_utils import cancel_and_wait, delayed_start from .logging_utils import log_catch, log_context _logger = logging.getLogger(__name__) diff --git a/packages/service-library/tests/test_async_utils.py b/packages/service-library/tests/test_async_utils.py index 9bb1b4fff45..e7164417fc6 100644 --- a/packages/service-library/tests/test_async_utils.py +++ b/packages/service-library/tests/test_async_utils.py @@ -7,7 +7,6 @@ import random from collections import deque from dataclasses import dataclass -from datetime import timedelta from time import time from typing import Any @@ -15,7 +14,6 @@ from faker import Faker from servicelib.async_utils import ( _sequential_jobs_contexts, - delayed_start, run_sequentially_in_context, ) @@ -225,20 +223,3 @@ 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_with_delay(): - @delayed_start(timedelta(seconds=0.2)) - async def decorated_awaitable() -> int: - return 42 - - assert await decorated_awaitable() == 42 - - async def another_awaitable() -> int: - return 42 - - decorated_another_awaitable = delayed_start(timedelta(seconds=0.2))( - another_awaitable - ) - - assert await decorated_another_awaitable() == 42 From bc456aa353b53717a236f7880704d1b25169af93 Mon Sep 17 00:00:00 2001 From: Pedro Crespo-Valero <32402063+pcrespov@users.noreply.github.com> Date: Tue, 24 Jun 2025 14:04:34 +0200 Subject: [PATCH 11/13] rename --- packages/celery-library/src/celery_library/task.py | 4 ++-- packages/common-library/src/common_library/async_tools.py | 2 +- packages/common-library/tests/test_async_tools.py | 8 ++++---- .../service-library/src/servicelib/background_task.py | 4 ++-- .../src/servicelib/long_running_tasks/task.py | 4 ++-- packages/service-library/src/servicelib/redis/_client.py | 4 ++-- .../rabbitmq/test_rabbitmq_rpc_interfaces_async_jobs.py | 4 ++-- packages/service-library/tests/redis/test_project_lock.py | 4 ++-- packages/service-library/tests/test_background_task.py | 7 +++++-- .../service-library/tests/test_background_task_utils.py | 4 ++-- .../src/simcore_service_agent/services/volumes_manager.py | 6 +++--- .../core/_prometheus_instrumentation.py | 4 ++-- .../src/simcore_service_api_server/core/health_checker.py | 4 ++-- .../modules/auto_scaling_task.py | 4 ++-- .../modules/buffer_machines_pool_task.py | 4 ++-- .../modules/clusters_management_task.py | 4 ++-- .../rabbitmq_worker_plugin.py | 6 ++++-- .../modules/comp_scheduler/_manager.py | 4 ++-- .../modules/dynamic_sidecar/scheduler/_core/_scheduler.py | 6 +++--- .../utils/base_distributed_identifier.py | 4 ++-- .../src/simcore_service_director/registry_proxy.py | 4 ++-- .../services/status_monitor/_monitor.py | 6 ++++-- .../modules/outputs/_manager.py | 4 ++-- .../modules/prometheus_metrics.py | 4 ++-- .../modules/resource_tracking/_core.py | 4 ++-- .../modules/system_monitor/_disk_usage.py | 4 ++-- .../services/background_tasks_setup.py | 4 ++-- .../services/fire_and_forget_setup.py | 4 ++-- .../clients/postgres/_liveness.py | 4 ++-- .../background_task_periodic_heartbeat_check_setup.py | 4 ++-- .../services/fire_and_forget_setup.py | 4 ++-- .../storage/src/simcore_service_storage/dsm_cleaner.py | 4 ++-- .../storage/src/simcore_service_storage/utils/s3_utils.py | 4 ++-- .../garbage_collector/_tasks_api_keys.py | 4 ++-- .../garbage_collector/_tasks_core.py | 4 ++-- .../garbage_collector/_tasks_trash.py | 4 ++-- .../garbage_collector/_tasks_users.py | 4 ++-- .../licenses/_itis_vip_syncer_service.py | 4 ++-- .../src/simcore_service_webserver/payments/_tasks.py | 4 ++-- 39 files changed, 88 insertions(+), 81 deletions(-) diff --git a/packages/celery-library/src/celery_library/task.py b/packages/celery-library/src/celery_library/task.py index 4eb37be5aa7..dafc74a4f4c 100644 --- a/packages/celery-library/src/celery_library/task.py +++ b/packages/celery-library/src/celery_library/task.py @@ -12,7 +12,7 @@ AbortableTask, ) from celery.exceptions import Ignore # type: ignore[import-untyped] -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from pydantic import NonNegativeInt from servicelib.celery.models import TaskID @@ -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_and_wait( + await cancel_and_shielded_wait( main_task, max_delay=_DEFAULT_CANCEL_TASK_TIMEOUT.total_seconds(), ) diff --git a/packages/common-library/src/common_library/async_tools.py b/packages/common-library/src/common_library/async_tools.py index f343ffc4fd4..3cd4cf5acff 100644 --- a/packages/common-library/src/common_library/async_tools.py +++ b/packages/common-library/src/common_library/async_tools.py @@ -69,7 +69,7 @@ async def maybe_await( return obj -async def cancel_and_wait( +async def cancel_and_shielded_wait( task: asyncio.Task, *, max_delay: float | None = None ) -> None: """Cancels the given task and waits for it to complete. diff --git a/packages/common-library/tests/test_async_tools.py b/packages/common-library/tests/test_async_tools.py index bcebc3438a0..774f529ae5a 100644 --- a/packages/common-library/tests/test_async_tools.py +++ b/packages/common-library/tests/test_async_tools.py @@ -6,7 +6,7 @@ import pytest from common_library.async_tools import ( - cancel_and_wait, + cancel_and_shielded_wait, delayed_start, make_async, maybe_await, @@ -121,7 +121,7 @@ async def coro(): await asyncio.sleep(0.1) # Let coro start start = time.time() - await cancel_and_wait(task) + await cancel_and_shielded_wait(task) elapsed = time.time() - start assert elapsed < SLEEP_TIME, "Task should be cancelled quickly" @@ -149,7 +149,7 @@ async def coro(): async def outer_coro(): try: - await cancel_and_wait(inner_task) + await cancel_and_shielded_wait(inner_task) except asyncio.CancelledError: assert ( not inner_task.cancelled() @@ -194,7 +194,7 @@ async def slow_cleanup_coro(): # Cancel with a max_delay shorter than cleanup time with pytest.raises(TimeoutError): - await cancel_and_wait( + await cancel_and_shielded_wait( task, max_delay=CLEANUP_TIME / 10 ) # 0.2 seconds < 2 seconds cleanup diff --git a/packages/service-library/src/servicelib/background_task.py b/packages/service-library/src/servicelib/background_task.py index b37ff0e0c2b..d2c897fefbf 100644 --- a/packages/service-library/src/servicelib/background_task.py +++ b/packages/service-library/src/servicelib/background_task.py @@ -6,7 +6,7 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine from typing import Any, Final, ParamSpec, TypeVar -from common_library.async_tools import cancel_and_wait, delayed_start +from common_library.async_tools import cancel_and_shielded_wait, delayed_start from tenacity import TryAgain, before_sleep_log, retry, retry_if_exception_type from tenacity.wait import wait_fixed @@ -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_and_wait(asyncio_task, max_delay=stop_timeout)) + await cancel_and_shielded_wait(asyncio_task, max_delay=stop_timeout) diff --git a/packages/service-library/src/servicelib/long_running_tasks/task.py b/packages/service-library/src/servicelib/long_running_tasks/task.py index 71978bb18f1..c4549868086 100644 --- a/packages/service-library/src/servicelib/long_running_tasks/task.py +++ b/packages/service-library/src/servicelib/long_running_tasks/task.py @@ -9,7 +9,7 @@ from typing import Any, Final, Protocol, TypeAlias from uuid import uuid4 -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from models_library.api_schemas_long_running_tasks.base import TaskProgress from pydantic import PositiveFloat from servicelib.background_task import create_periodic_task @@ -105,7 +105,7 @@ async def teardown(self) -> None: if self._stale_tasks_monitor_task: with log_catch(_logger, reraise=False): - await cancel_and_wait( + await cancel_and_shielded_wait( self._stale_tasks_monitor_task, max_delay=_CANCEL_TASK_TIMEOUT ) diff --git a/packages/service-library/src/servicelib/redis/_client.py b/packages/service-library/src/servicelib/redis/_client.py index e59f1d21252..54d7c3afce1 100644 --- a/packages/service-library/src/servicelib/redis/_client.py +++ b/packages/service-library/src/servicelib/redis/_client.py @@ -8,11 +8,11 @@ import redis.asyncio as aioredis import redis.exceptions +from common_library.async_tools import cancel_and_shielded_wait from redis.asyncio.lock import Lock from redis.asyncio.retry import Retry from redis.backoff import ExponentialBackoff -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_and_wait( + await cancel_and_shielded_wait( self._health_check_task, max_delay=_HEALTHCHECK_TASK_TIMEOUT_S ) diff --git a/packages/service-library/tests/rabbitmq/test_rabbitmq_rpc_interfaces_async_jobs.py b/packages/service-library/tests/rabbitmq/test_rabbitmq_rpc_interfaces_async_jobs.py index 68d55f239ec..168e91fa37f 100644 --- a/packages/service-library/tests/rabbitmq/test_rabbitmq_rpc_interfaces_async_jobs.py +++ b/packages/service-library/tests/rabbitmq/test_rabbitmq_rpc_interfaces_async_jobs.py @@ -4,7 +4,7 @@ from dataclasses import dataclass, field import pytest -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from faker import Faker from models_library.api_schemas_rpc_async_jobs.async_jobs import ( AsyncJobGet, @@ -137,7 +137,7 @@ async def setup(self) -> None: yield for task in fake_server.tasks: - await cancel_and_wait(task) + await cancel_and_shielded_wait(task) @pytest.mark.parametrize("method", ["result", "status", "cancel"]) diff --git a/packages/service-library/tests/redis/test_project_lock.py b/packages/service-library/tests/redis/test_project_lock.py index 04bfa99f176..297104d624b 100644 --- a/packages/service-library/tests/redis/test_project_lock.py +++ b/packages/service-library/tests/redis/test_project_lock.py @@ -10,7 +10,7 @@ from uuid import UUID import pytest -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from faker import Faker from models_library.projects import ProjectID from models_library.projects_access import Owner @@ -141,4 +141,4 @@ async def _locked_fct() -> None: with pytest.raises(ProjectLockError): await _locked_fct() - await cancel_and_wait(task1) + await cancel_and_shielded_wait(task1) diff --git a/packages/service-library/tests/test_background_task.py b/packages/service-library/tests/test_background_task.py index 958f15b7f9c..34ab5277fdf 100644 --- a/packages/service-library/tests/test_background_task.py +++ b/packages/service-library/tests/test_background_task.py @@ -13,7 +13,7 @@ from unittest.mock import AsyncMock import pytest -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from faker import Faker from pytest_mock.plugin import MockerFixture from servicelib.background_task import create_periodic_task, periodic, periodic_task @@ -78,7 +78,10 @@ async def _creator( yield _creator # cleanup await asyncio.gather( - *(cancel_and_wait(t, max_delay=stop_task_timeout) for t in created_tasks) + *( + cancel_and_shielded_wait(t, max_delay=stop_task_timeout) + for t in created_tasks + ) ) diff --git a/packages/service-library/tests/test_background_task_utils.py b/packages/service-library/tests/test_background_task_utils.py index c703592146b..6d64e381ac8 100644 --- a/packages/service-library/tests/test_background_task_utils.py +++ b/packages/service-library/tests/test_background_task_utils.py @@ -13,7 +13,7 @@ import arrow import pytest -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_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_and_wait(task, max_delay=5) + await cancel_and_shielded_wait(task, max_delay=5) events_timestamps: tuple[float, ...] = tuple( x.args[0].timestamp() for x in sleep_events.call_args_list diff --git a/services/agent/src/simcore_service_agent/services/volumes_manager.py b/services/agent/src/simcore_service_agent/services/volumes_manager.py index bba3b8178b8..8ee39c04d7c 100644 --- a/services/agent/src/simcore_service_agent/services/volumes_manager.py +++ b/services/agent/src/simcore_service_agent/services/volumes_manager.py @@ -6,7 +6,7 @@ import arrow from aiodocker.docker import Docker -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from models_library.projects_nodes_io import NodeID from pydantic import NonNegativeFloat @@ -61,10 +61,10 @@ async def shutdown(self) -> None: await self.docker.close() if self._task_bookkeeping: - await cancel_and_wait(self._task_bookkeeping) + await cancel_and_shielded_wait(self._task_bookkeeping) if self._task_periodic_volume_cleanup: - await cancel_and_wait(self._task_periodic_volume_cleanup) + await cancel_and_shielded_wait(self._task_periodic_volume_cleanup) async def _bookkeeping_task(self) -> None: with log_context(_logger, logging.DEBUG, "volume bookkeeping"): diff --git a/services/api-server/src/simcore_service_api_server/core/_prometheus_instrumentation.py b/services/api-server/src/simcore_service_api_server/core/_prometheus_instrumentation.py index 995900a3a77..b4c3248e7f2 100644 --- a/services/api-server/src/simcore_service_api_server/core/_prometheus_instrumentation.py +++ b/services/api-server/src/simcore_service_api_server/core/_prometheus_instrumentation.py @@ -4,7 +4,7 @@ from datetime import timedelta from typing import Final, cast -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from prometheus_client import CollectorRegistry, Gauge from pydantic import PositiveInt @@ -85,7 +85,7 @@ async def on_startup() -> None: async def on_shutdown() -> None: assert app.state.instrumentation_task # nosec with log_catch(_logger, reraise=False): - await cancel_and_wait(app.state.instrumentation_task) + await cancel_and_shielded_wait(app.state.instrumentation_task) app.add_event_handler("startup", on_startup) app.add_event_handler("shutdown", on_shutdown) diff --git a/services/api-server/src/simcore_service_api_server/core/health_checker.py b/services/api-server/src/simcore_service_api_server/core/health_checker.py index 3d99c38b562..11ce7ea0e9c 100644 --- a/services/api-server/src/simcore_service_api_server/core/health_checker.py +++ b/services/api-server/src/simcore_service_api_server/core/health_checker.py @@ -5,7 +5,7 @@ from typing import Annotated, Final, cast from uuid import uuid4 -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import Depends, FastAPI from models_library.rabbitmq_messages import LoggerRabbitMessage from models_library.users import UserID @@ -63,7 +63,7 @@ async def setup(self, health_check_task_period_seconds: PositiveFloat): async def teardown(self): if self._background_task: with log_catch(_logger, reraise=False): - await cancel_and_wait( + await cancel_and_shielded_wait( self._background_task, max_delay=self._timeout_seconds ) await self._log_distributor.deregister(job_id=self._dummy_job_id) diff --git a/services/autoscaling/src/simcore_service_autoscaling/modules/auto_scaling_task.py b/services/autoscaling/src/simcore_service_autoscaling/modules/auto_scaling_task.py index 756c716a2c2..51bc5f7b2f0 100644 --- a/services/autoscaling/src/simcore_service_autoscaling/modules/auto_scaling_task.py +++ b/services/autoscaling/src/simcore_service_autoscaling/modules/auto_scaling_task.py @@ -2,7 +2,7 @@ from collections.abc import Awaitable, Callable from typing import Final -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from servicelib.background_task import create_periodic_task from servicelib.redis import exclusive @@ -44,7 +44,7 @@ async def _startup() -> None: def on_app_shutdown(app: FastAPI) -> Callable[[], Awaitable[None]]: async def _stop() -> None: - await cancel_and_wait(app.state.autoscaler_task) + await cancel_and_shielded_wait(app.state.autoscaler_task) return _stop diff --git a/services/autoscaling/src/simcore_service_autoscaling/modules/buffer_machines_pool_task.py b/services/autoscaling/src/simcore_service_autoscaling/modules/buffer_machines_pool_task.py index 42dbd99c1f0..531e5286b76 100644 --- a/services/autoscaling/src/simcore_service_autoscaling/modules/buffer_machines_pool_task.py +++ b/services/autoscaling/src/simcore_service_autoscaling/modules/buffer_machines_pool_task.py @@ -2,7 +2,7 @@ from collections.abc import Awaitable, Callable from typing import Final -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from servicelib.background_task import create_periodic_task from servicelib.redis import exclusive @@ -44,7 +44,7 @@ async def _startup() -> None: def on_app_shutdown(app: FastAPI) -> Callable[[], Awaitable[None]]: async def _stop() -> None: if hasattr(app.state, "buffers_pool_task"): - await cancel_and_wait(app.state.buffers_pool_task) + await cancel_and_shielded_wait(app.state.buffers_pool_task) return _stop diff --git a/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters_management_task.py b/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters_management_task.py index 5c9affe3e60..96ca89e6b5c 100644 --- a/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters_management_task.py +++ b/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters_management_task.py @@ -2,7 +2,7 @@ import logging from collections.abc import Awaitable, Callable -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from servicelib.background_task import create_periodic_task from servicelib.redis import exclusive @@ -37,7 +37,7 @@ async def _startup() -> None: def on_app_shutdown(app: FastAPI) -> Callable[[], Awaitable[None]]: async def _stop() -> None: - await cancel_and_wait(app.state.clusters_cleaning_task, max_delay=5) + await cancel_and_shielded_wait(app.state.clusters_cleaning_task, max_delay=5) return _stop diff --git a/services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq_worker_plugin.py b/services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq_worker_plugin.py index 90458fbc977..1cf3aa59ab0 100644 --- a/services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq_worker_plugin.py +++ b/services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq_worker_plugin.py @@ -6,7 +6,7 @@ from typing import Final import distributed -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from servicelib.logging_utils import log_catch, log_context from servicelib.rabbitmq import RabbitMQClient, wait_till_rabbitmq_responsive from servicelib.rabbitmq._models import RabbitMessage @@ -104,7 +104,9 @@ async def _() -> None: # Cancel the message processor task if self._message_processor: with log_catch(_logger, reraise=False): - await cancel_and_wait(self._message_processor, max_delay=5) + await cancel_and_shielded_wait( + self._message_processor, max_delay=5 + ) self._message_processor = None # close client diff --git a/services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_manager.py b/services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_manager.py index 0287e92f402..556f000d741 100644 --- a/services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_manager.py +++ b/services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_manager.py @@ -2,7 +2,7 @@ from typing import Final import networkx as nx -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from models_library.projects import ProjectID from models_library.users import UserID @@ -202,4 +202,4 @@ async def setup_manager(app: FastAPI) -> None: async def shutdown_manager(app: FastAPI) -> None: - await cancel_and_wait(app.state.scheduler_manager) + await cancel_and_shielded_wait(app.state.scheduler_manager) diff --git a/services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_scheduler.py b/services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_scheduler.py index 41467eebe19..9bd71b1f466 100644 --- a/services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_scheduler.py +++ b/services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_scheduler.py @@ -23,7 +23,7 @@ from typing import Final import arrow -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from models_library.api_schemas_directorv2.dynamic_services import ( DynamicServiceCreate, @@ -125,7 +125,7 @@ async def shutdown(self) -> None: self._to_observe = {} if self._scheduler_task is not None: - await cancel_and_wait(self._scheduler_task, max_delay=5) + await cancel_and_shielded_wait(self._scheduler_task, max_delay=5) self._scheduler_task = None if self._trigger_observation_queue_task is not None: @@ -364,7 +364,7 @@ async def mark_service_for_removal( self._service_observation_task[service_name] ) if isinstance(service_task, asyncio.Task): - await cancel_and_wait(service_task, max_delay=10) + await cancel_and_shielded_wait(service_task, max_delay=10) if skip_observation_recreation: return diff --git a/services/director-v2/src/simcore_service_director_v2/utils/base_distributed_identifier.py b/services/director-v2/src/simcore_service_director_v2/utils/base_distributed_identifier.py index 7a94c4e71fa..72052da58d9 100644 --- a/services/director-v2/src/simcore_service_director_v2/utils/base_distributed_identifier.py +++ b/services/director-v2/src/simcore_service_director_v2/utils/base_distributed_identifier.py @@ -4,7 +4,7 @@ from datetime import timedelta from typing import Final, Generic, TypeVar -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from pydantic import NonNegativeInt from servicelib.background_task import create_periodic_task from servicelib.logging_utils import log_catch, log_context @@ -76,7 +76,7 @@ async def setup(self) -> None: async def shutdown(self) -> None: if self._cleanup_task: - await cancel_and_wait(self._cleanup_task, max_delay=5) + await cancel_and_shielded_wait(self._cleanup_task, max_delay=5) @classmethod def class_path(cls) -> str: diff --git a/services/director/src/simcore_service_director/registry_proxy.py b/services/director/src/simcore_service_director/registry_proxy.py index 964560e701d..4639f50fa1e 100644 --- a/services/director/src/simcore_service_director/registry_proxy.py +++ b/services/director/src/simcore_service_director/registry_proxy.py @@ -8,7 +8,7 @@ import httpx from aiocache import Cache, SimpleMemoryCache # type: ignore[import-untyped] -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from common_library.json_serialization import json_loads from fastapi import FastAPI, status from servicelib.background_task import create_periodic_task @@ -255,7 +255,7 @@ async def on_startup() -> None: async def on_shutdown() -> None: if app.state.auto_cache_task: - await cancel_and_wait(app.state.auto_cache_task) + await cancel_and_shielded_wait(app.state.auto_cache_task) app.add_event_handler("startup", on_startup) app.add_event_handler("shutdown", on_shutdown) diff --git a/services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/status_monitor/_monitor.py b/services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/status_monitor/_monitor.py index 6107474eb33..effdc857391 100644 --- a/services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/status_monitor/_monitor.py +++ b/services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/status_monitor/_monitor.py @@ -5,7 +5,7 @@ from typing import Final import arrow -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from models_library.projects_nodes_io import NodeID from pydantic import NonNegativeFloat, NonNegativeInt @@ -149,4 +149,6 @@ async def _periodic_check_services_require_status_update() -> None: async def shutdown(self) -> None: if getattr(self.app.state, "status_monitor_background_task", None): - await cancel_and_wait(self.app.state.status_monitor_background_task) + await cancel_and_shielded_wait( + self.app.state.status_monitor_background_task + ) diff --git a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/outputs/_manager.py b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/outputs/_manager.py index 4459150c6c0..596ff80383a 100644 --- a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/outputs/_manager.py +++ b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/outputs/_manager.py @@ -6,7 +6,7 @@ from datetime import timedelta from functools import partial -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from common_library.errors_classes import OsparcErrorMixin from fastapi import FastAPI from models_library.rabbitmq_messages import ProgressType @@ -204,7 +204,7 @@ async def shutdown(self) -> None: with log_context(_logger, logging.INFO, f"{OutputsManager.__name__} shutdown"): await self._uploading_task_cancel() if self._task_scheduler_worker is not None: - await cancel_and_wait( + await cancel_and_shielded_wait( self._task_scheduler_worker, max_delay=self.task_monitor_interval_s ) diff --git a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/prometheus_metrics.py b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/prometheus_metrics.py index 348efa5f278..3a5a0703641 100644 --- a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/prometheus_metrics.py +++ b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/prometheus_metrics.py @@ -6,7 +6,7 @@ from typing import Final import arrow -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI, status from models_library.callbacks_mapping import CallbacksMapping, UserServiceCommand from pydantic import BaseModel, NonNegativeFloat, NonNegativeInt @@ -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_and_wait( + await cancel_and_shielded_wait( self._metrics_recovery_task, max_delay=_TASK_CANCELLATION_TIMEOUT_S ) diff --git a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/resource_tracking/_core.py b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/resource_tracking/_core.py index 466f9cdf0ba..96b224f1dab 100644 --- a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/resource_tracking/_core.py +++ b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/resource_tracking/_core.py @@ -2,7 +2,7 @@ import logging from typing import Final -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from models_library.generated_models.docker_rest_api import ContainerState from models_library.rabbitmq_messages import ( @@ -62,7 +62,7 @@ async def _start_heart_beat_task(app: FastAPI) -> None: async def stop_heart_beat_task(app: FastAPI) -> None: resource_tracking: ResourceTrackingState = app.state.resource_tracking if resource_tracking.heart_beat_task: - await cancel_and_wait( + await cancel_and_shielded_wait( resource_tracking.heart_beat_task, max_delay=_STOP_WORKER_TIMEOUT_S ) diff --git a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/system_monitor/_disk_usage.py b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/system_monitor/_disk_usage.py index df8f3990d15..bcf0436466d 100644 --- a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/system_monitor/_disk_usage.py +++ b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/system_monitor/_disk_usage.py @@ -7,7 +7,7 @@ from typing import Final import psutil -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from models_library.api_schemas_dynamic_sidecar.telemetry import ( DiskUsage, @@ -185,7 +185,7 @@ async def setup(self) -> None: async def shutdown(self) -> None: if self._monitor_task: - await cancel_and_wait(self._monitor_task) + await cancel_and_shielded_wait(self._monitor_task) def set_disk_usage_for_path(self, overwrite_usage: dict[str, DiskUsage]) -> None: """ diff --git a/services/efs-guardian/src/simcore_service_efs_guardian/services/background_tasks_setup.py b/services/efs-guardian/src/simcore_service_efs_guardian/services/background_tasks_setup.py index f6602899340..a916cb57aad 100644 --- a/services/efs-guardian/src/simcore_service_efs_guardian/services/background_tasks_setup.py +++ b/services/efs-guardian/src/simcore_service_efs_guardian/services/background_tasks_setup.py @@ -3,7 +3,7 @@ from collections.abc import Awaitable, Callable from datetime import timedelta -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from servicelib.background_task_utils import exclusive_periodic from servicelib.logging_utils import log_catch, log_context @@ -50,7 +50,7 @@ async def _stop() -> None: ): assert _app # nosec if _app.state.efs_guardian_removal_policy_background_task: - await cancel_and_wait( + await cancel_and_shielded_wait( _app.state.efs_guardian_removal_policy_background_task ) diff --git a/services/efs-guardian/src/simcore_service_efs_guardian/services/fire_and_forget_setup.py b/services/efs-guardian/src/simcore_service_efs_guardian/services/fire_and_forget_setup.py index 94fb8985cd3..fb941974f67 100644 --- a/services/efs-guardian/src/simcore_service_efs_guardian/services/fire_and_forget_setup.py +++ b/services/efs-guardian/src/simcore_service_efs_guardian/services/fire_and_forget_setup.py @@ -1,7 +1,7 @@ import logging from collections.abc import Awaitable, Callable -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from servicelib.logging_utils import log_catch, log_context @@ -28,7 +28,7 @@ async def _stop() -> None: assert _app # nosec if _app.state.efs_guardian_fire_and_forget_tasks: for task in _app.state.efs_guardian_fire_and_forget_tasks: - await cancel_and_wait(task) + await cancel_and_shielded_wait(task) return _stop diff --git a/services/notifications/src/simcore_service_notifications/clients/postgres/_liveness.py b/services/notifications/src/simcore_service_notifications/clients/postgres/_liveness.py index 51372a1ee48..48ec1cbe2de 100644 --- a/services/notifications/src/simcore_service_notifications/clients/postgres/_liveness.py +++ b/services/notifications/src/simcore_service_notifications/clients/postgres/_liveness.py @@ -3,7 +3,7 @@ from datetime import timedelta from typing import Final -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from models_library.healthchecks import IsResponsive, LivenessResult from servicelib.background_task import create_periodic_task @@ -40,4 +40,4 @@ async def setup(self) -> None: async def teardown(self) -> None: if self._task is not None: with log_catch(_logger, reraise=False): - await cancel_and_wait(self._task, max_delay=5) + await cancel_and_shielded_wait(self._task, max_delay=5) diff --git a/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/background_task_periodic_heartbeat_check_setup.py b/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/background_task_periodic_heartbeat_check_setup.py index 8287d83a57b..96f7357397d 100644 --- a/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/background_task_periodic_heartbeat_check_setup.py +++ b/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/background_task_periodic_heartbeat_check_setup.py @@ -3,7 +3,7 @@ from collections.abc import Awaitable, Callable from typing import TypedDict -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from servicelib.background_task_utils import exclusive_periodic from servicelib.logging_utils import log_catch, log_context @@ -69,7 +69,7 @@ async def _stop() -> None: ): assert _app # nosec if _app.state.rut_background_task__periodic_check_of_running_services: - await cancel_and_wait( + await cancel_and_shielded_wait( _app.state.rut_background_task__periodic_check_of_running_services ) diff --git a/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/fire_and_forget_setup.py b/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/fire_and_forget_setup.py index 0c8c8091c02..4c3f8d14bee 100644 --- a/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/fire_and_forget_setup.py +++ b/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/fire_and_forget_setup.py @@ -1,7 +1,7 @@ import logging from collections.abc import Awaitable, Callable -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from servicelib.logging_utils import log_catch, log_context @@ -32,7 +32,7 @@ async def _stop() -> None: assert _app # nosec if _app.state.rut_fire_and_forget_tasks: for task in _app.state.rut_fire_and_forget_tasks: - await cancel_and_wait(task) + await cancel_and_shielded_wait(task) return _stop diff --git a/services/storage/src/simcore_service_storage/dsm_cleaner.py b/services/storage/src/simcore_service_storage/dsm_cleaner.py index 45458fef336..fc9b95da3ee 100644 --- a/services/storage/src/simcore_service_storage/dsm_cleaner.py +++ b/services/storage/src/simcore_service_storage/dsm_cleaner.py @@ -23,7 +23,7 @@ from datetime import timedelta from typing import cast -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from fastapi import FastAPI from servicelib.background_task_utils import exclusive_periodic from servicelib.logging_utils import log_context @@ -66,7 +66,7 @@ async def _periodic_dsm_clean() -> None: async def _on_shutdown() -> None: assert isinstance(app.state.dsm_cleaner_task, asyncio.Task) # nosec - await cancel_and_wait(app.state.dsm_cleaner_task) + await cancel_and_shielded_wait(app.state.dsm_cleaner_task) app.add_event_handler("startup", _on_startup) app.add_event_handler("shutdown", _on_shutdown) diff --git a/services/storage/src/simcore_service_storage/utils/s3_utils.py b/services/storage/src/simcore_service_storage/utils/s3_utils.py index 25e0f3020c3..0698a1c8296 100644 --- a/services/storage/src/simcore_service_storage/utils/s3_utils.py +++ b/services/storage/src/simcore_service_storage/utils/s3_utils.py @@ -4,7 +4,7 @@ from collections import defaultdict from dataclasses import dataclass, field -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from pydantic import ByteSize, TypeAdapter from servicelib.background_task import create_periodic_task from servicelib.progress_bar import ProgressBarData @@ -39,7 +39,7 @@ async def __aexit__(self, exc_type, exc_value, traceback) -> None: self.finalize_transfer() await asyncio.sleep(0) assert self._async_update_periodic_task # nosec - await cancel_and_wait(self._async_update_periodic_task) + await cancel_and_shielded_wait(self._async_update_periodic_task) async def _async_update(self) -> None: await self._update_task_event.wait() diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py index a3d7d88b42c..4da67aa66c6 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py @@ -8,7 +8,7 @@ from collections.abc import AsyncIterator, Callable from aiohttp import web -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from tenacity import retry from tenacity.before_sleep import before_sleep_log from tenacity.wait import wait_exponential @@ -68,6 +68,6 @@ async def _cleanup_ctx_fun( yield # tear-down - await cancel_and_wait(task) + await cancel_and_shielded_wait(task) return _cleanup_ctx_fun diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_core.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_core.py index 44ea83ffd3c..9b79dfe6ad9 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_core.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_core.py @@ -9,7 +9,7 @@ from collections.abc import AsyncGenerator from aiohttp import web -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from servicelib.logging_utils import log_context from ._core import collect_garbage @@ -57,7 +57,7 @@ async def run_background_task(app: web.Application) -> AsyncGenerator: app[_GC_TASK_CONFIG]["force_stop"] = True - await cancel_and_wait(gc_bg_task) + await cancel_and_shielded_wait(gc_bg_task) async def _collect_garbage_periodically(app: web.Application): diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py index 5eb47d55e1f..ca2e811a4ed 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py @@ -8,7 +8,7 @@ from collections.abc import AsyncIterator, Callable from aiohttp import web -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from servicelib.logging_utils import log_context from tenacity import retry from tenacity.before_sleep import before_sleep_log @@ -56,6 +56,6 @@ async def _cleanup_ctx_fun( yield # tear-down - await cancel_and_wait(task) + await cancel_and_shielded_wait(task) return _cleanup_ctx_fun diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py index 0876d28ad23..bbaf0a359c0 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py @@ -8,7 +8,7 @@ from collections.abc import AsyncIterator, Callable from aiohttp import web -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from models_library.users import UserID from servicelib.logging_utils import get_log_record_extra, log_context from tenacity import retry @@ -108,6 +108,6 @@ async def _cleanup_ctx_fun( yield # tear-down - await cancel_and_wait(task) + await cancel_and_shielded_wait(task) return _cleanup_ctx_fun diff --git a/services/web/server/src/simcore_service_webserver/licenses/_itis_vip_syncer_service.py b/services/web/server/src/simcore_service_webserver/licenses/_itis_vip_syncer_service.py index 1b84807a679..3347b5e870f 100644 --- a/services/web/server/src/simcore_service_webserver/licenses/_itis_vip_syncer_service.py +++ b/services/web/server/src/simcore_service_webserver/licenses/_itis_vip_syncer_service.py @@ -4,7 +4,7 @@ from datetime import timedelta from aiohttp import web -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from httpx import AsyncClient from models_library.licenses import LicensedResourceType from servicelib.background_task_utils import exclusive_periodic @@ -115,6 +115,6 @@ async def _periodic_sync() -> None: yield - await cancel_and_wait(background_task) + await cancel_and_shielded_wait(background_task) app.cleanup_ctx.append(_lifespan) diff --git a/services/web/server/src/simcore_service_webserver/payments/_tasks.py b/services/web/server/src/simcore_service_webserver/payments/_tasks.py index d385eb47314..a90291a3d9b 100644 --- a/services/web/server/src/simcore_service_webserver/payments/_tasks.py +++ b/services/web/server/src/simcore_service_webserver/payments/_tasks.py @@ -5,7 +5,7 @@ from typing import Any from aiohttp import web -from common_library.async_tools import cancel_and_wait +from common_library.async_tools import cancel_and_shielded_wait from models_library.api_schemas_webserver.wallets import PaymentID, PaymentMethodID from pydantic import HttpUrl, TypeAdapter from servicelib.aiohttp.typing_extension import CleanupContextFunc @@ -144,6 +144,6 @@ async def _cleanup_ctx_fun( yield # tear-down - await cancel_and_wait(task) + await cancel_and_shielded_wait(task) return _cleanup_ctx_fun From 039469ce3f4e9b2b7ca2bc1624f7ccfa629edd8c Mon Sep 17 00:00:00 2001 From: Pedro Crespo-Valero <32402063+pcrespov@users.noreply.github.com> Date: Tue, 24 Jun 2025 16:52:15 +0200 Subject: [PATCH 12/13] rename --- packages/celery-library/src/celery_library/task.py | 4 ++-- packages/common-library/src/common_library/async_tools.py | 4 ++-- packages/common-library/tests/test_async_tools.py | 8 ++++---- .../service-library/src/servicelib/background_task.py | 4 ++-- .../src/servicelib/long_running_tasks/task.py | 4 ++-- packages/service-library/src/servicelib/redis/_client.py | 4 ++-- .../rabbitmq/test_rabbitmq_rpc_interfaces_async_jobs.py | 4 ++-- packages/service-library/tests/redis/test_project_lock.py | 4 ++-- packages/service-library/tests/test_background_task.py | 7 ++----- .../service-library/tests/test_background_task_utils.py | 4 ++-- .../src/simcore_service_agent/services/volumes_manager.py | 6 +++--- .../core/_prometheus_instrumentation.py | 4 ++-- .../src/simcore_service_api_server/core/health_checker.py | 4 ++-- .../modules/auto_scaling_task.py | 4 ++-- .../modules/buffer_machines_pool_task.py | 4 ++-- .../modules/clusters_management_task.py | 4 ++-- .../rabbitmq_worker_plugin.py | 6 ++---- .../modules/comp_scheduler/_manager.py | 4 ++-- .../modules/dynamic_sidecar/scheduler/_core/_scheduler.py | 6 +++--- .../utils/base_distributed_identifier.py | 4 ++-- .../src/simcore_service_director/registry_proxy.py | 4 ++-- .../services/status_monitor/_monitor.py | 6 ++---- .../modules/outputs/_manager.py | 4 ++-- .../modules/prometheus_metrics.py | 4 ++-- .../modules/resource_tracking/_core.py | 4 ++-- .../modules/system_monitor/_disk_usage.py | 4 ++-- .../services/background_tasks_setup.py | 4 ++-- .../services/fire_and_forget_setup.py | 4 ++-- .../clients/postgres/_liveness.py | 4 ++-- .../background_task_periodic_heartbeat_check_setup.py | 4 ++-- .../services/fire_and_forget_setup.py | 4 ++-- .../storage/src/simcore_service_storage/dsm_cleaner.py | 4 ++-- .../storage/src/simcore_service_storage/utils/s3_utils.py | 4 ++-- .../garbage_collector/_tasks_api_keys.py | 4 ++-- .../garbage_collector/_tasks_core.py | 4 ++-- .../garbage_collector/_tasks_trash.py | 4 ++-- .../garbage_collector/_tasks_users.py | 4 ++-- .../licenses/_itis_vip_syncer_service.py | 4 ++-- .../src/simcore_service_webserver/payments/_tasks.py | 4 ++-- 39 files changed, 82 insertions(+), 89 deletions(-) diff --git a/packages/celery-library/src/celery_library/task.py b/packages/celery-library/src/celery_library/task.py index dafc74a4f4c..075e10036bc 100644 --- a/packages/celery-library/src/celery_library/task.py +++ b/packages/celery-library/src/celery_library/task.py @@ -12,7 +12,7 @@ AbortableTask, ) from celery.exceptions import Ignore # type: ignore[import-untyped] -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from pydantic import NonNegativeInt from servicelib.celery.models import TaskID @@ -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_and_shielded_wait( + await cancel_wait_task( main_task, max_delay=_DEFAULT_CANCEL_TASK_TIMEOUT.total_seconds(), ) diff --git a/packages/common-library/src/common_library/async_tools.py b/packages/common-library/src/common_library/async_tools.py index 3cd4cf5acff..ff0b61fa1d2 100644 --- a/packages/common-library/src/common_library/async_tools.py +++ b/packages/common-library/src/common_library/async_tools.py @@ -69,10 +69,10 @@ async def maybe_await( return obj -async def cancel_and_shielded_wait( +async def cancel_wait_task( task: asyncio.Task, *, max_delay: float | None = None ) -> None: - """Cancels the given task and waits for it to complete. + """Cancels the given task and waits for it to complete Arguments: task -- task to be canceled diff --git a/packages/common-library/tests/test_async_tools.py b/packages/common-library/tests/test_async_tools.py index 774f529ae5a..9e9e081056f 100644 --- a/packages/common-library/tests/test_async_tools.py +++ b/packages/common-library/tests/test_async_tools.py @@ -6,7 +6,7 @@ import pytest from common_library.async_tools import ( - cancel_and_shielded_wait, + cancel_wait_task, delayed_start, make_async, maybe_await, @@ -121,7 +121,7 @@ async def coro(): await asyncio.sleep(0.1) # Let coro start start = time.time() - await cancel_and_shielded_wait(task) + await cancel_wait_task(task) elapsed = time.time() - start assert elapsed < SLEEP_TIME, "Task should be cancelled quickly" @@ -149,7 +149,7 @@ async def coro(): async def outer_coro(): try: - await cancel_and_shielded_wait(inner_task) + await cancel_wait_task(inner_task) except asyncio.CancelledError: assert ( not inner_task.cancelled() @@ -194,7 +194,7 @@ async def slow_cleanup_coro(): # Cancel with a max_delay shorter than cleanup time with pytest.raises(TimeoutError): - await cancel_and_shielded_wait( + await cancel_wait_task( task, max_delay=CLEANUP_TIME / 10 ) # 0.2 seconds < 2 seconds cleanup diff --git a/packages/service-library/src/servicelib/background_task.py b/packages/service-library/src/servicelib/background_task.py index d2c897fefbf..506dd314ea1 100644 --- a/packages/service-library/src/servicelib/background_task.py +++ b/packages/service-library/src/servicelib/background_task.py @@ -6,7 +6,7 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine from typing import Any, Final, ParamSpec, TypeVar -from common_library.async_tools import cancel_and_shielded_wait, delayed_start +from common_library.async_tools import cancel_wait_task, delayed_start from tenacity import TryAgain, before_sleep_log, retry, retry_if_exception_type from tenacity.wait import wait_fixed @@ -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 cancel_and_shielded_wait(asyncio_task, max_delay=stop_timeout) + await cancel_wait_task(asyncio_task, max_delay=stop_timeout) diff --git a/packages/service-library/src/servicelib/long_running_tasks/task.py b/packages/service-library/src/servicelib/long_running_tasks/task.py index c4549868086..a6007e2059a 100644 --- a/packages/service-library/src/servicelib/long_running_tasks/task.py +++ b/packages/service-library/src/servicelib/long_running_tasks/task.py @@ -9,7 +9,7 @@ from typing import Any, Final, Protocol, TypeAlias from uuid import uuid4 -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from models_library.api_schemas_long_running_tasks.base import TaskProgress from pydantic import PositiveFloat from servicelib.background_task import create_periodic_task @@ -105,7 +105,7 @@ async def teardown(self) -> None: if self._stale_tasks_monitor_task: with log_catch(_logger, reraise=False): - await cancel_and_shielded_wait( + await cancel_wait_task( self._stale_tasks_monitor_task, max_delay=_CANCEL_TASK_TIMEOUT ) diff --git a/packages/service-library/src/servicelib/redis/_client.py b/packages/service-library/src/servicelib/redis/_client.py index 54d7c3afce1..e961a6a73e9 100644 --- a/packages/service-library/src/servicelib/redis/_client.py +++ b/packages/service-library/src/servicelib/redis/_client.py @@ -8,7 +8,7 @@ import redis.asyncio as aioredis import redis.exceptions -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from redis.asyncio.lock import Lock from redis.asyncio.retry import Retry from redis.backoff import ExponentialBackoff @@ -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_and_shielded_wait( + await cancel_wait_task( self._health_check_task, max_delay=_HEALTHCHECK_TASK_TIMEOUT_S ) diff --git a/packages/service-library/tests/rabbitmq/test_rabbitmq_rpc_interfaces_async_jobs.py b/packages/service-library/tests/rabbitmq/test_rabbitmq_rpc_interfaces_async_jobs.py index 168e91fa37f..40455ee6d7f 100644 --- a/packages/service-library/tests/rabbitmq/test_rabbitmq_rpc_interfaces_async_jobs.py +++ b/packages/service-library/tests/rabbitmq/test_rabbitmq_rpc_interfaces_async_jobs.py @@ -4,7 +4,7 @@ from dataclasses import dataclass, field import pytest -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from faker import Faker from models_library.api_schemas_rpc_async_jobs.async_jobs import ( AsyncJobGet, @@ -137,7 +137,7 @@ async def setup(self) -> None: yield for task in fake_server.tasks: - await cancel_and_shielded_wait(task) + await cancel_wait_task(task) @pytest.mark.parametrize("method", ["result", "status", "cancel"]) diff --git a/packages/service-library/tests/redis/test_project_lock.py b/packages/service-library/tests/redis/test_project_lock.py index 297104d624b..03fe4f0e462 100644 --- a/packages/service-library/tests/redis/test_project_lock.py +++ b/packages/service-library/tests/redis/test_project_lock.py @@ -10,7 +10,7 @@ from uuid import UUID import pytest -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from faker import Faker from models_library.projects import ProjectID from models_library.projects_access import Owner @@ -141,4 +141,4 @@ async def _locked_fct() -> None: with pytest.raises(ProjectLockError): await _locked_fct() - await cancel_and_shielded_wait(task1) + await cancel_wait_task(task1) diff --git a/packages/service-library/tests/test_background_task.py b/packages/service-library/tests/test_background_task.py index 34ab5277fdf..715aaec92b6 100644 --- a/packages/service-library/tests/test_background_task.py +++ b/packages/service-library/tests/test_background_task.py @@ -13,7 +13,7 @@ from unittest.mock import AsyncMock import pytest -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from faker import Faker from pytest_mock.plugin import MockerFixture from servicelib.background_task import create_periodic_task, periodic, periodic_task @@ -78,10 +78,7 @@ async def _creator( yield _creator # cleanup await asyncio.gather( - *( - cancel_and_shielded_wait(t, max_delay=stop_task_timeout) - for t in created_tasks - ) + *(cancel_wait_task(t, max_delay=stop_task_timeout) for t in created_tasks) ) diff --git a/packages/service-library/tests/test_background_task_utils.py b/packages/service-library/tests/test_background_task_utils.py index 6d64e381ac8..92d0337eca7 100644 --- a/packages/service-library/tests/test_background_task_utils.py +++ b/packages/service-library/tests/test_background_task_utils.py @@ -13,7 +13,7 @@ import arrow import pytest -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task 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_and_shielded_wait(task, max_delay=5) + await cancel_wait_task(task, max_delay=5) events_timestamps: tuple[float, ...] = tuple( x.args[0].timestamp() for x in sleep_events.call_args_list diff --git a/services/agent/src/simcore_service_agent/services/volumes_manager.py b/services/agent/src/simcore_service_agent/services/volumes_manager.py index 8ee39c04d7c..1ef6ef1d0cb 100644 --- a/services/agent/src/simcore_service_agent/services/volumes_manager.py +++ b/services/agent/src/simcore_service_agent/services/volumes_manager.py @@ -6,7 +6,7 @@ import arrow from aiodocker.docker import Docker -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from models_library.projects_nodes_io import NodeID from pydantic import NonNegativeFloat @@ -61,10 +61,10 @@ async def shutdown(self) -> None: await self.docker.close() if self._task_bookkeeping: - await cancel_and_shielded_wait(self._task_bookkeeping) + await cancel_wait_task(self._task_bookkeeping) if self._task_periodic_volume_cleanup: - await cancel_and_shielded_wait(self._task_periodic_volume_cleanup) + await cancel_wait_task(self._task_periodic_volume_cleanup) async def _bookkeeping_task(self) -> None: with log_context(_logger, logging.DEBUG, "volume bookkeeping"): diff --git a/services/api-server/src/simcore_service_api_server/core/_prometheus_instrumentation.py b/services/api-server/src/simcore_service_api_server/core/_prometheus_instrumentation.py index b4c3248e7f2..68bef9b369b 100644 --- a/services/api-server/src/simcore_service_api_server/core/_prometheus_instrumentation.py +++ b/services/api-server/src/simcore_service_api_server/core/_prometheus_instrumentation.py @@ -4,7 +4,7 @@ from datetime import timedelta from typing import Final, cast -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from prometheus_client import CollectorRegistry, Gauge from pydantic import PositiveInt @@ -85,7 +85,7 @@ async def on_startup() -> None: async def on_shutdown() -> None: assert app.state.instrumentation_task # nosec with log_catch(_logger, reraise=False): - await cancel_and_shielded_wait(app.state.instrumentation_task) + await cancel_wait_task(app.state.instrumentation_task) app.add_event_handler("startup", on_startup) app.add_event_handler("shutdown", on_shutdown) diff --git a/services/api-server/src/simcore_service_api_server/core/health_checker.py b/services/api-server/src/simcore_service_api_server/core/health_checker.py index 11ce7ea0e9c..1e882569f39 100644 --- a/services/api-server/src/simcore_service_api_server/core/health_checker.py +++ b/services/api-server/src/simcore_service_api_server/core/health_checker.py @@ -5,7 +5,7 @@ from typing import Annotated, Final, cast from uuid import uuid4 -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import Depends, FastAPI from models_library.rabbitmq_messages import LoggerRabbitMessage from models_library.users import UserID @@ -63,7 +63,7 @@ async def setup(self, health_check_task_period_seconds: PositiveFloat): async def teardown(self): if self._background_task: with log_catch(_logger, reraise=False): - await cancel_and_shielded_wait( + await cancel_wait_task( self._background_task, max_delay=self._timeout_seconds ) await self._log_distributor.deregister(job_id=self._dummy_job_id) diff --git a/services/autoscaling/src/simcore_service_autoscaling/modules/auto_scaling_task.py b/services/autoscaling/src/simcore_service_autoscaling/modules/auto_scaling_task.py index 51bc5f7b2f0..b3d5ab8e193 100644 --- a/services/autoscaling/src/simcore_service_autoscaling/modules/auto_scaling_task.py +++ b/services/autoscaling/src/simcore_service_autoscaling/modules/auto_scaling_task.py @@ -2,7 +2,7 @@ from collections.abc import Awaitable, Callable from typing import Final -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from servicelib.background_task import create_periodic_task from servicelib.redis import exclusive @@ -44,7 +44,7 @@ async def _startup() -> None: def on_app_shutdown(app: FastAPI) -> Callable[[], Awaitable[None]]: async def _stop() -> None: - await cancel_and_shielded_wait(app.state.autoscaler_task) + await cancel_wait_task(app.state.autoscaler_task) return _stop diff --git a/services/autoscaling/src/simcore_service_autoscaling/modules/buffer_machines_pool_task.py b/services/autoscaling/src/simcore_service_autoscaling/modules/buffer_machines_pool_task.py index 531e5286b76..625c0170476 100644 --- a/services/autoscaling/src/simcore_service_autoscaling/modules/buffer_machines_pool_task.py +++ b/services/autoscaling/src/simcore_service_autoscaling/modules/buffer_machines_pool_task.py @@ -2,7 +2,7 @@ from collections.abc import Awaitable, Callable from typing import Final -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from servicelib.background_task import create_periodic_task from servicelib.redis import exclusive @@ -44,7 +44,7 @@ async def _startup() -> None: def on_app_shutdown(app: FastAPI) -> Callable[[], Awaitable[None]]: async def _stop() -> None: if hasattr(app.state, "buffers_pool_task"): - await cancel_and_shielded_wait(app.state.buffers_pool_task) + await cancel_wait_task(app.state.buffers_pool_task) return _stop diff --git a/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters_management_task.py b/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters_management_task.py index 96ca89e6b5c..70e80c550f8 100644 --- a/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters_management_task.py +++ b/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters_management_task.py @@ -2,7 +2,7 @@ import logging from collections.abc import Awaitable, Callable -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from servicelib.background_task import create_periodic_task from servicelib.redis import exclusive @@ -37,7 +37,7 @@ async def _startup() -> None: def on_app_shutdown(app: FastAPI) -> Callable[[], Awaitable[None]]: async def _stop() -> None: - await cancel_and_shielded_wait(app.state.clusters_cleaning_task, max_delay=5) + await cancel_wait_task(app.state.clusters_cleaning_task, max_delay=5) return _stop diff --git a/services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq_worker_plugin.py b/services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq_worker_plugin.py index 1cf3aa59ab0..ef288fea483 100644 --- a/services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq_worker_plugin.py +++ b/services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq_worker_plugin.py @@ -6,7 +6,7 @@ from typing import Final import distributed -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from servicelib.logging_utils import log_catch, log_context from servicelib.rabbitmq import RabbitMQClient, wait_till_rabbitmq_responsive from servicelib.rabbitmq._models import RabbitMessage @@ -104,9 +104,7 @@ async def _() -> None: # Cancel the message processor task if self._message_processor: with log_catch(_logger, reraise=False): - await cancel_and_shielded_wait( - self._message_processor, max_delay=5 - ) + await cancel_wait_task(self._message_processor, max_delay=5) self._message_processor = None # close client diff --git a/services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_manager.py b/services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_manager.py index 556f000d741..430bb4a871e 100644 --- a/services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_manager.py +++ b/services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_manager.py @@ -2,7 +2,7 @@ from typing import Final import networkx as nx -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from models_library.projects import ProjectID from models_library.users import UserID @@ -202,4 +202,4 @@ async def setup_manager(app: FastAPI) -> None: async def shutdown_manager(app: FastAPI) -> None: - await cancel_and_shielded_wait(app.state.scheduler_manager) + await cancel_wait_task(app.state.scheduler_manager) diff --git a/services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_scheduler.py b/services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_scheduler.py index 9bd71b1f466..d6618d6cbfd 100644 --- a/services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_scheduler.py +++ b/services/director-v2/src/simcore_service_director_v2/modules/dynamic_sidecar/scheduler/_core/_scheduler.py @@ -23,7 +23,7 @@ from typing import Final import arrow -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from models_library.api_schemas_directorv2.dynamic_services import ( DynamicServiceCreate, @@ -125,7 +125,7 @@ async def shutdown(self) -> None: self._to_observe = {} if self._scheduler_task is not None: - await cancel_and_shielded_wait(self._scheduler_task, max_delay=5) + await cancel_wait_task(self._scheduler_task, max_delay=5) self._scheduler_task = None if self._trigger_observation_queue_task is not None: @@ -364,7 +364,7 @@ async def mark_service_for_removal( self._service_observation_task[service_name] ) if isinstance(service_task, asyncio.Task): - await cancel_and_shielded_wait(service_task, max_delay=10) + await cancel_wait_task(service_task, max_delay=10) if skip_observation_recreation: return diff --git a/services/director-v2/src/simcore_service_director_v2/utils/base_distributed_identifier.py b/services/director-v2/src/simcore_service_director_v2/utils/base_distributed_identifier.py index 72052da58d9..25d5dca72f3 100644 --- a/services/director-v2/src/simcore_service_director_v2/utils/base_distributed_identifier.py +++ b/services/director-v2/src/simcore_service_director_v2/utils/base_distributed_identifier.py @@ -4,7 +4,7 @@ from datetime import timedelta from typing import Final, Generic, TypeVar -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from pydantic import NonNegativeInt from servicelib.background_task import create_periodic_task from servicelib.logging_utils import log_catch, log_context @@ -76,7 +76,7 @@ async def setup(self) -> None: async def shutdown(self) -> None: if self._cleanup_task: - await cancel_and_shielded_wait(self._cleanup_task, max_delay=5) + await cancel_wait_task(self._cleanup_task, max_delay=5) @classmethod def class_path(cls) -> str: diff --git a/services/director/src/simcore_service_director/registry_proxy.py b/services/director/src/simcore_service_director/registry_proxy.py index 4639f50fa1e..177460f915c 100644 --- a/services/director/src/simcore_service_director/registry_proxy.py +++ b/services/director/src/simcore_service_director/registry_proxy.py @@ -8,7 +8,7 @@ import httpx from aiocache import Cache, SimpleMemoryCache # type: ignore[import-untyped] -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from common_library.json_serialization import json_loads from fastapi import FastAPI, status from servicelib.background_task import create_periodic_task @@ -255,7 +255,7 @@ async def on_startup() -> None: async def on_shutdown() -> None: if app.state.auto_cache_task: - await cancel_and_shielded_wait(app.state.auto_cache_task) + await cancel_wait_task(app.state.auto_cache_task) app.add_event_handler("startup", on_startup) app.add_event_handler("shutdown", on_shutdown) diff --git a/services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/status_monitor/_monitor.py b/services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/status_monitor/_monitor.py index effdc857391..5e05384c990 100644 --- a/services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/status_monitor/_monitor.py +++ b/services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/status_monitor/_monitor.py @@ -5,7 +5,7 @@ from typing import Final import arrow -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from models_library.projects_nodes_io import NodeID from pydantic import NonNegativeFloat, NonNegativeInt @@ -149,6 +149,4 @@ async def _periodic_check_services_require_status_update() -> None: async def shutdown(self) -> None: if getattr(self.app.state, "status_monitor_background_task", None): - await cancel_and_shielded_wait( - self.app.state.status_monitor_background_task - ) + await cancel_wait_task(self.app.state.status_monitor_background_task) diff --git a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/outputs/_manager.py b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/outputs/_manager.py index 596ff80383a..367d5dc224a 100644 --- a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/outputs/_manager.py +++ b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/outputs/_manager.py @@ -6,7 +6,7 @@ from datetime import timedelta from functools import partial -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from common_library.errors_classes import OsparcErrorMixin from fastapi import FastAPI from models_library.rabbitmq_messages import ProgressType @@ -204,7 +204,7 @@ async def shutdown(self) -> None: with log_context(_logger, logging.INFO, f"{OutputsManager.__name__} shutdown"): await self._uploading_task_cancel() if self._task_scheduler_worker is not None: - await cancel_and_shielded_wait( + await cancel_wait_task( self._task_scheduler_worker, max_delay=self.task_monitor_interval_s ) diff --git a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/prometheus_metrics.py b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/prometheus_metrics.py index 3a5a0703641..7a6c47f6649 100644 --- a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/prometheus_metrics.py +++ b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/prometheus_metrics.py @@ -6,7 +6,7 @@ from typing import Final import arrow -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI, status from models_library.callbacks_mapping import CallbacksMapping, UserServiceCommand from pydantic import BaseModel, NonNegativeFloat, NonNegativeInt @@ -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_and_shielded_wait( + await cancel_wait_task( self._metrics_recovery_task, max_delay=_TASK_CANCELLATION_TIMEOUT_S ) diff --git a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/resource_tracking/_core.py b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/resource_tracking/_core.py index 96b224f1dab..d61402e6d28 100644 --- a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/resource_tracking/_core.py +++ b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/resource_tracking/_core.py @@ -2,7 +2,7 @@ import logging from typing import Final -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from models_library.generated_models.docker_rest_api import ContainerState from models_library.rabbitmq_messages import ( @@ -62,7 +62,7 @@ async def _start_heart_beat_task(app: FastAPI) -> None: async def stop_heart_beat_task(app: FastAPI) -> None: resource_tracking: ResourceTrackingState = app.state.resource_tracking if resource_tracking.heart_beat_task: - await cancel_and_shielded_wait( + await cancel_wait_task( resource_tracking.heart_beat_task, max_delay=_STOP_WORKER_TIMEOUT_S ) diff --git a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/system_monitor/_disk_usage.py b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/system_monitor/_disk_usage.py index bcf0436466d..fc6941d1a54 100644 --- a/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/system_monitor/_disk_usage.py +++ b/services/dynamic-sidecar/src/simcore_service_dynamic_sidecar/modules/system_monitor/_disk_usage.py @@ -7,7 +7,7 @@ from typing import Final import psutil -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from models_library.api_schemas_dynamic_sidecar.telemetry import ( DiskUsage, @@ -185,7 +185,7 @@ async def setup(self) -> None: async def shutdown(self) -> None: if self._monitor_task: - await cancel_and_shielded_wait(self._monitor_task) + await cancel_wait_task(self._monitor_task) def set_disk_usage_for_path(self, overwrite_usage: dict[str, DiskUsage]) -> None: """ diff --git a/services/efs-guardian/src/simcore_service_efs_guardian/services/background_tasks_setup.py b/services/efs-guardian/src/simcore_service_efs_guardian/services/background_tasks_setup.py index a916cb57aad..ae88f0fdb84 100644 --- a/services/efs-guardian/src/simcore_service_efs_guardian/services/background_tasks_setup.py +++ b/services/efs-guardian/src/simcore_service_efs_guardian/services/background_tasks_setup.py @@ -3,7 +3,7 @@ from collections.abc import Awaitable, Callable from datetime import timedelta -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from servicelib.background_task_utils import exclusive_periodic from servicelib.logging_utils import log_catch, log_context @@ -50,7 +50,7 @@ async def _stop() -> None: ): assert _app # nosec if _app.state.efs_guardian_removal_policy_background_task: - await cancel_and_shielded_wait( + await cancel_wait_task( _app.state.efs_guardian_removal_policy_background_task ) diff --git a/services/efs-guardian/src/simcore_service_efs_guardian/services/fire_and_forget_setup.py b/services/efs-guardian/src/simcore_service_efs_guardian/services/fire_and_forget_setup.py index fb941974f67..5ca03f7bd9e 100644 --- a/services/efs-guardian/src/simcore_service_efs_guardian/services/fire_and_forget_setup.py +++ b/services/efs-guardian/src/simcore_service_efs_guardian/services/fire_and_forget_setup.py @@ -1,7 +1,7 @@ import logging from collections.abc import Awaitable, Callable -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from servicelib.logging_utils import log_catch, log_context @@ -28,7 +28,7 @@ async def _stop() -> None: assert _app # nosec if _app.state.efs_guardian_fire_and_forget_tasks: for task in _app.state.efs_guardian_fire_and_forget_tasks: - await cancel_and_shielded_wait(task) + await cancel_wait_task(task) return _stop diff --git a/services/notifications/src/simcore_service_notifications/clients/postgres/_liveness.py b/services/notifications/src/simcore_service_notifications/clients/postgres/_liveness.py index 48ec1cbe2de..6d26fd83e93 100644 --- a/services/notifications/src/simcore_service_notifications/clients/postgres/_liveness.py +++ b/services/notifications/src/simcore_service_notifications/clients/postgres/_liveness.py @@ -3,7 +3,7 @@ from datetime import timedelta from typing import Final -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from models_library.healthchecks import IsResponsive, LivenessResult from servicelib.background_task import create_periodic_task @@ -40,4 +40,4 @@ async def setup(self) -> None: async def teardown(self) -> None: if self._task is not None: with log_catch(_logger, reraise=False): - await cancel_and_shielded_wait(self._task, max_delay=5) + await cancel_wait_task(self._task, max_delay=5) diff --git a/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/background_task_periodic_heartbeat_check_setup.py b/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/background_task_periodic_heartbeat_check_setup.py index 96f7357397d..a747cd5d476 100644 --- a/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/background_task_periodic_heartbeat_check_setup.py +++ b/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/background_task_periodic_heartbeat_check_setup.py @@ -3,7 +3,7 @@ from collections.abc import Awaitable, Callable from typing import TypedDict -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from servicelib.background_task_utils import exclusive_periodic from servicelib.logging_utils import log_catch, log_context @@ -69,7 +69,7 @@ async def _stop() -> None: ): assert _app # nosec if _app.state.rut_background_task__periodic_check_of_running_services: - await cancel_and_shielded_wait( + await cancel_wait_task( _app.state.rut_background_task__periodic_check_of_running_services ) diff --git a/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/fire_and_forget_setup.py b/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/fire_and_forget_setup.py index 4c3f8d14bee..a1e7db5ac30 100644 --- a/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/fire_and_forget_setup.py +++ b/services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/services/fire_and_forget_setup.py @@ -1,7 +1,7 @@ import logging from collections.abc import Awaitable, Callable -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from servicelib.logging_utils import log_catch, log_context @@ -32,7 +32,7 @@ async def _stop() -> None: assert _app # nosec if _app.state.rut_fire_and_forget_tasks: for task in _app.state.rut_fire_and_forget_tasks: - await cancel_and_shielded_wait(task) + await cancel_wait_task(task) return _stop diff --git a/services/storage/src/simcore_service_storage/dsm_cleaner.py b/services/storage/src/simcore_service_storage/dsm_cleaner.py index fc9b95da3ee..6194d61a835 100644 --- a/services/storage/src/simcore_service_storage/dsm_cleaner.py +++ b/services/storage/src/simcore_service_storage/dsm_cleaner.py @@ -23,7 +23,7 @@ from datetime import timedelta from typing import cast -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from fastapi import FastAPI from servicelib.background_task_utils import exclusive_periodic from servicelib.logging_utils import log_context @@ -66,7 +66,7 @@ async def _periodic_dsm_clean() -> None: async def _on_shutdown() -> None: assert isinstance(app.state.dsm_cleaner_task, asyncio.Task) # nosec - await cancel_and_shielded_wait(app.state.dsm_cleaner_task) + await cancel_wait_task(app.state.dsm_cleaner_task) app.add_event_handler("startup", _on_startup) app.add_event_handler("shutdown", _on_shutdown) diff --git a/services/storage/src/simcore_service_storage/utils/s3_utils.py b/services/storage/src/simcore_service_storage/utils/s3_utils.py index 0698a1c8296..0ebec26a6bb 100644 --- a/services/storage/src/simcore_service_storage/utils/s3_utils.py +++ b/services/storage/src/simcore_service_storage/utils/s3_utils.py @@ -4,7 +4,7 @@ from collections import defaultdict from dataclasses import dataclass, field -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from pydantic import ByteSize, TypeAdapter from servicelib.background_task import create_periodic_task from servicelib.progress_bar import ProgressBarData @@ -39,7 +39,7 @@ async def __aexit__(self, exc_type, exc_value, traceback) -> None: self.finalize_transfer() await asyncio.sleep(0) assert self._async_update_periodic_task # nosec - await cancel_and_shielded_wait(self._async_update_periodic_task) + await cancel_wait_task(self._async_update_periodic_task) async def _async_update(self) -> None: await self._update_task_event.wait() diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py index 4da67aa66c6..73c2d570ebf 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py @@ -8,7 +8,7 @@ from collections.abc import AsyncIterator, Callable from aiohttp import web -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from tenacity import retry from tenacity.before_sleep import before_sleep_log from tenacity.wait import wait_exponential @@ -68,6 +68,6 @@ async def _cleanup_ctx_fun( yield # tear-down - await cancel_and_shielded_wait(task) + await cancel_wait_task(task) return _cleanup_ctx_fun diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_core.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_core.py index 9b79dfe6ad9..255c9d006dc 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_core.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_core.py @@ -9,7 +9,7 @@ from collections.abc import AsyncGenerator from aiohttp import web -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from servicelib.logging_utils import log_context from ._core import collect_garbage @@ -57,7 +57,7 @@ async def run_background_task(app: web.Application) -> AsyncGenerator: app[_GC_TASK_CONFIG]["force_stop"] = True - await cancel_and_shielded_wait(gc_bg_task) + await cancel_wait_task(gc_bg_task) async def _collect_garbage_periodically(app: web.Application): diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py index ca2e811a4ed..121b8b79ee0 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py @@ -8,7 +8,7 @@ from collections.abc import AsyncIterator, Callable from aiohttp import web -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from servicelib.logging_utils import log_context from tenacity import retry from tenacity.before_sleep import before_sleep_log @@ -56,6 +56,6 @@ async def _cleanup_ctx_fun( yield # tear-down - await cancel_and_shielded_wait(task) + await cancel_wait_task(task) return _cleanup_ctx_fun diff --git a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py index bbaf0a359c0..26cdb9e053c 100644 --- a/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py +++ b/services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py @@ -8,7 +8,7 @@ from collections.abc import AsyncIterator, Callable from aiohttp import web -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from models_library.users import UserID from servicelib.logging_utils import get_log_record_extra, log_context from tenacity import retry @@ -108,6 +108,6 @@ async def _cleanup_ctx_fun( yield # tear-down - await cancel_and_shielded_wait(task) + await cancel_wait_task(task) return _cleanup_ctx_fun diff --git a/services/web/server/src/simcore_service_webserver/licenses/_itis_vip_syncer_service.py b/services/web/server/src/simcore_service_webserver/licenses/_itis_vip_syncer_service.py index 3347b5e870f..7cca49731e0 100644 --- a/services/web/server/src/simcore_service_webserver/licenses/_itis_vip_syncer_service.py +++ b/services/web/server/src/simcore_service_webserver/licenses/_itis_vip_syncer_service.py @@ -4,7 +4,7 @@ from datetime import timedelta from aiohttp import web -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from httpx import AsyncClient from models_library.licenses import LicensedResourceType from servicelib.background_task_utils import exclusive_periodic @@ -115,6 +115,6 @@ async def _periodic_sync() -> None: yield - await cancel_and_shielded_wait(background_task) + await cancel_wait_task(background_task) app.cleanup_ctx.append(_lifespan) diff --git a/services/web/server/src/simcore_service_webserver/payments/_tasks.py b/services/web/server/src/simcore_service_webserver/payments/_tasks.py index a90291a3d9b..5e6cd74c05b 100644 --- a/services/web/server/src/simcore_service_webserver/payments/_tasks.py +++ b/services/web/server/src/simcore_service_webserver/payments/_tasks.py @@ -5,7 +5,7 @@ from typing import Any from aiohttp import web -from common_library.async_tools import cancel_and_shielded_wait +from common_library.async_tools import cancel_wait_task from models_library.api_schemas_webserver.wallets import PaymentID, PaymentMethodID from pydantic import HttpUrl, TypeAdapter from servicelib.aiohttp.typing_extension import CleanupContextFunc @@ -144,6 +144,6 @@ async def _cleanup_ctx_fun( yield # tear-down - await cancel_and_shielded_wait(task) + await cancel_wait_task(task) return _cleanup_ctx_fun From b715ab07e8bb848961f6f269ae0b9b7a983f2d9a Mon Sep 17 00:00:00 2001 From: Pedro Crespo-Valero <32402063+pcrespov@users.noreply.github.com> Date: Tue, 24 Jun 2025 17:00:23 +0200 Subject: [PATCH 13/13] fixes on canceling --- packages/common-library/src/common_library/async_tools.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/common-library/src/common_library/async_tools.py b/packages/common-library/src/common_library/async_tools.py index ff0b61fa1d2..8d4a276deab 100644 --- a/packages/common-library/src/common_library/async_tools.py +++ b/packages/common-library/src/common_library/async_tools.py @@ -90,7 +90,10 @@ async def cancel_wait_task( CancelledError: raised ONLY if owner is being cancelled. """ - task.cancel() + cancelling = task.cancel() + if not cancelling: + return # task was alredy cancelled + assert task.cancelling() # nosec assert not task.cancelled() # nosec