fixeed removal and task cancellation

Author: Andrei Neagu

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

@@ -95,9 +95,9 @@ async def start_long_running_task(
             dumps=json_dumps,
         )
     except asyncio.CancelledError:
-        # cancel the task, the client has disconnected
+        # remove the task, the client was disconnected
         if task_id:
-            await lrt_api.cancel_task(
+            await lrt_api.remove_task(
                 long_running_manager.tasks_manager, task_context, task_id
             )
         raise

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

@@ -128,12 +128,11 @@ async def _wait_for_task_result() -> Any:
             exception=e,
         ) from e
     except Exception as e:
-        error = TaskExceptionError(task_id=task_id, exception=e, traceback="")
         _logger.warning(
             create_troubleshootting_log_message(
                 user_error_msg=f"{task_id=} raised an exception",
                 error=e,
                 tip=f"Check the logs of the service responding to '{client.base_url}'",
             )
         )
-        raise error from e
+        raise TaskExceptionError(task_id=task_id, exception=e, traceback="") from e

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

@@ -103,12 +103,5 @@ async def get_task_result(
 async def remove_task(
     tasks_manager: TasksManager, task_context: TaskContext, task_id: TaskId
 ) -> None:
-    """removes / cancels a task"""
+    """cancels and removes the task"""
     await tasks_manager.remove_task(task_id, with_task_context=task_context)
-
-
-async def cancel_task(
-    tasks_manager: TasksManager, task_context: TaskContext, task_id: TaskId
-) -> None:
-    """cancels a task"""
-    await tasks_manager.cancel_task(task_id, with_task_context=task_context)

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

@@ -3,7 +3,6 @@
 import functools
 import inspect
 import logging
-import traceback
 import urllib.parse
 from typing import Any, ClassVar, Final, Protocol, TypeAlias
 from uuid import uuid4
@@ -14,15 +13,13 @@
 from settings_library.redis import RedisDatabase, RedisSettings
 
 from ..background_task import create_periodic_task
-from ..logging_utils import log_catch
 from ..redis import RedisClientSDK, exclusive
 from ._redis_serialization import object_to_string, string_to_object
 from ._store.base import BaseStore
 from ._store.redis import RedisStore
 from .errors import (
     TaskAlreadyRunningError,
     TaskCancelledError,
-    TaskExceptionError,
     TaskNotCompletedError,
     TaskNotFoundError,
     TaskNotRegisteredError,
@@ -171,12 +168,10 @@ async def teardown(self) -> None:
             )
 
         if self._stale_tasks_monitor_task:
-            with log_catch(_logger, reraise=False):
-                await cancel_wait_task(self._stale_tasks_monitor_task)
+            await cancel_wait_task(self._stale_tasks_monitor_task)
 
         if self._cancelled_tasks_removal_task:
-            with log_catch(_logger, reraise=False):
-                await cancel_wait_task(self._cancelled_tasks_removal_task)
+            await cancel_wait_task(self._cancelled_tasks_removal_task)
 
         if self.redis_client_sdk is not None:
             await self.redis_client_sdk.shutdown()
@@ -368,38 +363,20 @@ async def get_task_result(
 
         return string_to_object(tracked_task.result_field.result)
 
-    @staticmethod
     async def _cancel_tracked_task(
-        task: asyncio.Task, task_id: TaskId, *, reraise_errors: bool
+        self, task: asyncio.Task, task_id: TaskId, with_task_context: TaskContext
     ) -> None:
         try:
+            await self._tasks_data.set_as_cancelled(task_id, with_task_context)
+            # TODO: remove this cancellation timeout
             await cancel_wait_task(task, max_delay=_CANCEL_TASK_TIMEOUT)
         except Exception as e:  # pylint:disable=broad-except
-            formatted_traceback = "".join(traceback.format_exception(e))
-            if reraise_errors:
-                raise TaskExceptionError(
-                    task_id=task_id, exception=e, traceback=formatted_traceback
-                ) from e
-
-    async def cancel_task(
-        self, task_id: TaskId, with_task_context: TaskContext
-    ) -> None:
-        """
-        Eventually cancels the task.
-
-        # NOTE: aborts the task:
-        # - Immediately, if the task is running on the current worker.
-        # - Asynchronously (after a short delay), if the task is running on another worker.
-
-        raises TaskNotFoundError if the task cannot be found
-        """
-        await self._tasks_data.set_as_cancelled(task_id, with_task_context)
-        tracked_task = await self._get_tracked_task(task_id, with_task_context)
-        await self._cancel_tracked_task(
-            self._created_tasks[tracked_task.task_id], task_id, reraise_errors=False
-        )
-
-        # wait for task to be removed?
+            _logger.info(
+                "Task %s cancellation failed with error: %s",
+                task_id,
+                e,
+                stack_info=True,
+            )
 
     async def remove_task(
         self,
@@ -417,15 +394,12 @@ async def remove_task(
             return
 
         if tracked_task.task_id in self._created_tasks:
-            await self._cancel_tracked_task(
-                self._created_tasks[tracked_task.task_id],
-                task_id,
-                reraise_errors=reraise_errors,
-            )
-            # task might already be removed via cancellation background task
+            # will have affect in the worker which deals with the removal
             await self._tasks_data.delete_task_data(task_id)
             self._created_tasks.pop(tracked_task.task_id, None)
 
+        # TODO: wait for removal to becompleted here
+
     def _get_task_id(self, task_name: str, *, is_unique: bool) -> TaskId:
         unique_part = "unique" if is_unique else f"{uuid4()}"
         return f"{self.namespace}.{task_name}.{unique_part}"

packages/service-library/tests/long_running_tasks/test_long_running_tasks_task.py

@@ -17,7 +17,6 @@
 from servicelib.long_running_tasks import lrt_api
 from servicelib.long_running_tasks.errors import (
     TaskAlreadyRunningError,
-    TaskCancelledError,
     TaskNotCompletedError,
     TaskNotFoundError,
     TaskNotRegisteredError,
@@ -324,34 +323,6 @@ async def test_get_result_finished_with_error(
         await tasks_manager.get_task_result(task_id, with_task_context=empty_context)
 
 
-async def test_get_result_task_was_cancelled_multiple_times(
-    tasks_manager: TasksManager, empty_context: TaskContext
-):
-    task_id = await lrt_api.start_task(
-        tasks_manager,
-        a_background_task.__name__,
-        raise_when_finished=False,
-        total_sleep=10,
-        task_context=empty_context,
-    )
-    for _ in range(5):
-        await tasks_manager.cancel_task(task_id, with_task_context=empty_context)
-
-    with pytest.raises(  # noqa: PT012
-        TaskCancelledError, match=f"Task {task_id} was cancelled before completing"
-    ):
-        async for attempt in AsyncRetrying(
-            **{
-                **_RETRY_PARAMS,
-                "retry": retry_if_exception_type((TaskNotCompletedError,)),
-            }
-        ):
-            with attempt:
-                await tasks_manager.get_task_result(
-                    task_id, with_task_context=empty_context
-                )
-
-
 async def test_cancel_task_from_different_manager(
     use_in_memory_redis: RedisSettings,
     empty_context: TaskContext,
@@ -436,28 +407,6 @@ async def test_remove_unknown_task(
     )
 
 
-async def test_cancel_task_with_task_context(tasks_manager: TasksManager):
-    TASK_CONTEXT = {"some_context": "some_value"}
-    task_id = await lrt_api.start_task(
-        tasks_manager,
-        a_background_task.__name__,
-        raise_when_finished=False,
-        total_sleep=10,
-        task_context=TASK_CONTEXT,
-    )
-    # getting status fails if wrong task context given
-    with pytest.raises(TaskNotFoundError):
-        await tasks_manager.get_task_status(
-            task_id, with_task_context={"wrong_task_context": 12}
-        )
-    # getting status fails if wrong task context given
-    with pytest.raises(TaskNotFoundError):
-        await tasks_manager.cancel_task(
-            task_id, with_task_context={"wrong_task_context": 12}
-        )
-    await tasks_manager.cancel_task(task_id, with_task_context=TASK_CONTEXT)
-
-
 async def test__cancelled_tasks_worker_equivalent_of_cancellation_from_a_different_process(
     tasks_manager: TasksManager, empty_context: TaskContext
 ):