Monitor Celery tasks cancellation

Author: giancarloromeo

services/storage/src/simcore_service_storage/modules/celery/_task.py

@@ -7,21 +7,25 @@
 from typing import Any, Concatenate, Final, ParamSpec, TypeVar, overload
 
 from celery import Celery  # type: ignore[import-untyped]
-from celery.contrib.abortable import AbortableTask  # type: ignore[import-untyped]
+from celery.contrib.abortable import (  # type: ignore[import-untyped]
+    AbortableAsyncResult,
+    AbortableTask,
+)
 from pydantic import NonNegativeInt
+from servicelib.async_utils import cancel_wait_task
 
 from . import get_event_loop
 from .errors import encore_celery_transferrable_error
-from .models import TaskId
+from .models import TaskID, TaskId
 from .utils import get_fastapi_app
 
 _logger = logging.getLogger(__name__)
 
 _DEFAULT_TASK_TIMEOUT: Final[timedelta | None] = None
 _DEFAULT_MAX_RETRIES: Final[NonNegativeInt] = 3
 _DEFAULT_WAIT_BEFORE_RETRY: Final[timedelta] = timedelta(seconds=5)
-_DEFAULT_DONT_AUTORETRY_FOR: Final[tuple[type[Exception], ...]] = tuple()
-
+_DEFAULT_DONT_AUTORETRY_FOR: Final[tuple[type[Exception], ...]] = ()
+_DEFAULT_ABORT_TASK_TIMEOUT: Final[timedelta] = timedelta(seconds=0.5)
 
 T = TypeVar("T")
 P = ParamSpec("P")
@@ -43,8 +47,25 @@ def wrapper(task: AbortableTask, *args: P.args, **kwargs: P.kwargs) -> R:
             _logger.debug("task id: %s", task.request.id)
             # NOTE: task.request is a thread local object, so we need to pass the id explicitly
             assert task.request.id is not None  # nosec
+
+            async def run_task(task_id: TaskID) -> R:
+                task_coro = asyncio.create_task(coro(task, task_id, *args, **kwargs))
+
+                can_continue = True
+                while can_continue:
+                    if AbortableAsyncResult(task_id).is_aborted():
+                        _logger.warning("Task %s was aborted by user.", task_id)
+                        await cancel_wait_task(task_coro, max_delay=5)  # to constant
+                        raise asyncio.CancelledError
+                    if task_coro.done():
+                        break
+
+                    await asyncio.sleep(_DEFAULT_ABORT_TASK_TIMEOUT.total_seconds())
+
+                return task_coro.result()
+
             return asyncio.run_coroutine_threadsafe(
-                coro(task, task.request.id, *args, **kwargs),
+                run_task(task.request.id),
                 get_event_loop(fastapi_app),
             ).result()
 

services/storage/tests/unit/test_async_jobs.py

@@ -240,7 +240,6 @@ async def test_async_jobs_workflow(
 @pytest.mark.parametrize(
     "exposed_rpc_start",
     [
-        rpc_sync_job.__name__,
         rpc_async_job.__name__,
     ],
 )
@@ -259,7 +258,7 @@ async def test_async_jobs_cancel(
         user_id=user_id,
         product_name=product_name,
         action=Action.SLEEP,
-        payload=10,
+        payload=60 * 10,  # test hangs if not cancelled properly
     )
 
     await async_jobs.cancel(
@@ -283,6 +282,30 @@ async def test_async_jobs_cancel(
             job_id_data=job_id_data,
         )
 
+    async_job_get, job_id_data = await _start_task_via_rpc(
+        storage_rabbitmq_rpc_client,
+        rpc_task_name=exposed_rpc_start,
+        user_id=user_id,
+        product_name=product_name,
+        action=Action.ECHO,
+        payload="bla",
+    )
+
+    await _wait_for_job(
+        storage_rabbitmq_rpc_client,
+        async_job_get=async_job_get,
+        job_id_data=job_id_data,
+        stop_after=timedelta(seconds=15),
+    )
+
+    async_job_result = await async_jobs.result(
+        storage_rabbitmq_rpc_client,
+        rpc_namespace=STORAGE_RPC_NAMESPACE,
+        job_id=async_job_get.job_id,
+        job_id_data=job_id_data,
+    )
+    assert async_job_result.result == "bla"
+
 
 @pytest.mark.parametrize(
     "exposed_rpc_start",