removal does not wait for task to be removed and supports a timeout

Author: Andrei Neagu

packages/service-library/src/servicelib/aiohttp/long_running_tasks/_routes.py

@@ -75,5 +75,6 @@ async def remove_task(request: web.Request) -> web.Response:
         long_running_manager,
         long_running_manager.get_task_context(request),
         path_params.task_id,
+        wait_for_removal=False,  # frontend does not care about waiting for this
     )
     return web.json_response(status=status.HTTP_204_NO_CONTENT)

packages/service-library/src/servicelib/aiohttp/long_running_tasks/_server.py

@@ -103,6 +103,7 @@ async def start_long_running_task(
                 long_running_manager,
                 task_context,
                 task_id,
+                wait_for_removal=True,
             )
         raise
 

packages/service-library/src/servicelib/fastapi/long_running_tasks/_routes.py

@@ -106,4 +106,5 @@ async def remove_task(
         long_running_manager,
         task_context={},
         task_id=task_id,
+        wait_for_removal=True,  # only used by internal services, they will wait as before
     )

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

@@ -22,7 +22,7 @@
 
 _logger = logging.getLogger(__name__)
 
-_RPC_TIMEOUT_TASK_REMOVAL: Final[PositiveInt] = int(
+_RPC_MAX_CANCELLATION_TIMEOUT: Final[PositiveInt] = int(
     timedelta(minutes=60).total_seconds()
 )
 _RPC_TIMEOUT_SHORT_REQUESTS: Final[PositiveInt] = int(
@@ -134,14 +134,28 @@ async def remove_task(
     *,
     task_context: TaskContext,
     task_id: TaskId,
-    reraise_errors: bool = True,
+    wait_for_removal: bool,
+    reraise_errors: bool,
+    cancellation_timeout: timedelta | None = None,
 ) -> None:
+    timeout_s = (
+        _RPC_MAX_CANCELLATION_TIMEOUT
+        if cancellation_timeout is None
+        else int(cancellation_timeout.total_seconds())
+    )
+
+    # NOTE: task always gets cancelled even if not waiting for it
+    # request will return immediatlye, no need to wait so much
+    if not wait_for_removal:
+        timeout_s = _RPC_TIMEOUT_SHORT_REQUESTS
+
     result = await rabbitmq_rpc_client.request(
         get_rabbit_namespace(namespace),
         TypeAdapter(RPCMethodName).validate_python("remove_task"),
         task_context=task_context,
         task_id=task_id,
+        wait_for_removal=wait_for_removal,
         reraise_errors=reraise_errors,
-        timeout_s=_RPC_TIMEOUT_TASK_REMOVAL,
+        timeout_s=timeout_s,
     )
     assert result is None  # nosec

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

@@ -106,8 +106,12 @@ async def remove_task(
     *,
     task_context: TaskContext,
     task_id: TaskId,
+    wait_for_removal: bool,
     reraise_errors: bool,
 ) -> None:
     await long_running_manager.tasks_manager.remove_task(
-        task_id, with_task_context=task_context, reraise_errors=reraise_errors
+        task_id,
+        with_task_context=task_context,
+        wait_for_removal=wait_for_removal,
+        reraise_errors=reraise_errors,
     )

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

@@ -101,11 +101,15 @@ async def remove_task(
     long_running_manager: BaseLongRunningManager,
     task_context: TaskContext,
     task_id: TaskId,
+    *,
+    wait_for_removal: bool,
 ) -> None:
     """cancels and removes the task"""
     await _lrt_client.remove_task(
         rabbitmq_rpc_client,
         long_running_manager.lrt_namespace,
         task_id=task_id,
         task_context=task_context,
+        wait_for_removal=wait_for_removal,
+        reraise_errors=True,
     )

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

@@ -189,6 +189,7 @@ async def teardown(self) -> None:
                 tracked_task.task_id,
                 tracked_task.task_context,
                 # when closing we do not care about pending errors
+                wait_for_removal=True,
                 reraise_errors=False,
             )
 
@@ -260,7 +261,10 @@ async def _stale_tasks_monitor(self) -> None:
                     ).model_dump_json(),
                 )
                 await self.remove_task(
-                    task_id, with_task_context=task_context, reraise_errors=False
+                    task_id,
+                    with_task_context=task_context,
+                    wait_for_removal=True,
+                    reraise_errors=False,
                 )
 
     async def _cancelled_tasks_removal(self) -> None:
@@ -400,6 +404,7 @@ async def remove_task(
         task_id: TaskId,
         with_task_context: TaskContext,
         *,
+        wait_for_removal: bool,
         reraise_errors: bool = True,
     ) -> None:
         """cancels and removes task"""
@@ -414,6 +419,9 @@ async def remove_task(
             tracked_task.task_id, tracked_task.task_context
         )
 
+        if not wait_for_removal:
+            return
+
         # wait for task to be removed since it might not have been running
         # in this process
         async for attempt in AsyncRetrying(

packages/service-library/tests/aiohttp/long_running_tasks/test_long_running_tasks_with_task_context.py

@@ -29,6 +29,10 @@
 from servicelib.long_running_tasks.task import TaskContext
 from settings_library.rabbit import RabbitSettings
 from settings_library.redis import RedisSettings
+from tenacity.asyncio import AsyncRetrying
+from tenacity.retry import retry_if_exception_type
+from tenacity.stop import stop_after_delay
+from tenacity.wait import wait_fixed
 
 pytest_simcore_core_services_selection = [
     "rabbit",
@@ -176,15 +180,31 @@ async def test_cancel_task(
         task_id=task_id
     )
     # calling cancel without task context should find nothing
-    resp = await client_with_task_context.delete(f"{cancel_url}")
-    await assert_status(resp, status.HTTP_404_NOT_FOUND)
+    # no longer waits for removal to end
+    async for attempt in AsyncRetrying(
+        wait=wait_fixed(0.1),
+        stop=stop_after_delay(5),
+        reraise=True,
+        retry=retry_if_exception_type(AssertionError),
+    ):
+        with attempt:
+            resp = await client_with_task_context.delete(f"{cancel_url}")
+            await assert_status(resp, status.HTTP_404_NOT_FOUND)
     # calling with context should find and delete the task
     resp = await client_with_task_context.delete(
         f"{cancel_url.update_query(task_context)}"
     )
     await assert_status(resp, status.HTTP_204_NO_CONTENT)
     # calling with context a second time should find nothing
-    resp = await client_with_task_context.delete(
-        f"{cancel_url.update_query(task_context)}"
-    )
-    await assert_status(resp, status.HTTP_404_NOT_FOUND)
+    # no longer waits for removal to end
+    async for attempt in AsyncRetrying(
+        wait=wait_fixed(0.1),
+        stop=stop_after_delay(5),
+        reraise=True,
+        retry=retry_if_exception_type(AssertionError),
+    ):
+        with attempt:
+            resp = await client_with_task_context.delete(
+                f"{cancel_url.update_query(task_context)}"
+            )
+            await assert_status(resp, status.HTTP_404_NOT_FOUND)