🎨 Fixed long running tasks cleanup 🚨⚠️ (#8253)

Co-authored-by: Andrei Neagu <[email protected]>

Author: GitHK

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

@@ -1,12 +1,12 @@
 from aiohttp import web
 
-from ...long_running_tasks.base_long_running_manager import BaseLongRunningManager
+from ...long_running_tasks.manager import LongRunningManager
 from ...long_running_tasks.models import TaskContext
 from ._constants import APP_LONG_RUNNING_MANAGER_KEY
 from ._request import get_task_context
 
 
-class AiohttpLongRunningManager(BaseLongRunningManager):
+class AiohttpLongRunningManager(LongRunningManager):
 
     @staticmethod
     def get_task_context(request: web.Request) -> TaskContext:

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

@@ -1,10 +1,10 @@
 from fastapi import Request
 
-from ...long_running_tasks.base_long_running_manager import BaseLongRunningManager
+from ...long_running_tasks.manager import LongRunningManager
 from ...long_running_tasks.models import TaskContext
 
 
-class FastAPILongRunningManager(BaseLongRunningManager):
+class FastAPILongRunningManager(LongRunningManager):
     @staticmethod
     def get_task_context(request: Request) -> TaskContext:
         _ = request

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

@@ -15,15 +15,15 @@
 _logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
-    from .base_long_running_manager import BaseLongRunningManager
+    from .manager import LongRunningManager
 
 
 router = RPCRouter()
 
 
 @router.expose(reraise_if_error_type=(BaseLongRunningError,))
 async def start_task(
-    long_running_manager: "BaseLongRunningManager",
+    long_running_manager: "LongRunningManager",
     *,
     registered_task_name: RegisteredTaskName,
     unique: bool = False,
@@ -44,7 +44,7 @@ async def start_task(
 
 @router.expose(reraise_if_error_type=(BaseLongRunningError,))
 async def list_tasks(
-    long_running_manager: "BaseLongRunningManager", *, task_context: TaskContext
+    long_running_manager: "LongRunningManager", *, task_context: TaskContext
 ) -> list[TaskBase]:
     return await long_running_manager.tasks_manager.list_tasks(
         with_task_context=task_context
@@ -53,7 +53,7 @@ async def list_tasks(
 
 @router.expose(reraise_if_error_type=(BaseLongRunningError,))
 async def get_task_status(
-    long_running_manager: "BaseLongRunningManager",
+    long_running_manager: "LongRunningManager",
     *,
     task_context: TaskContext,
     task_id: TaskId,
@@ -65,7 +65,7 @@ async def get_task_status(
 
 @router.expose(reraise_if_error_type=(BaseLongRunningError, RPCTransferrableTaskError))
 async def get_task_result(
-    long_running_manager: "BaseLongRunningManager",
+    long_running_manager: "LongRunningManager",
     *,
     task_context: TaskContext,
     task_id: TaskId,
@@ -94,7 +94,7 @@ async def get_task_result(
 
 @router.expose(reraise_if_error_type=(BaseLongRunningError,))
 async def remove_task(
-    long_running_manager: "BaseLongRunningManager",
+    long_running_manager: "LongRunningManager",
     *,
     task_context: TaskContext,
     task_id: TaskId,

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

@@ -0,0 +1,46 @@
+import logging
+
+import redis.asyncio as aioredis
+from settings_library.redis import RedisDatabase, RedisSettings
+
+from ..logging_utils import log_context
+from ..redis._client import RedisClientSDK
+from .models import LRTNamespace
+
+_logger = logging.getLogger(__name__)
+
+
+class LongRunningClientHelper:
+    def __init__(self, redis_settings: RedisSettings):
+        self.redis_settings = redis_settings
+
+        self._client: RedisClientSDK | None = None
+
+    async def setup(self) -> None:
+        self._client = RedisClientSDK(
+            self.redis_settings.build_redis_dsn(RedisDatabase.LONG_RUNNING_TASKS),
+            client_name="long_running_tasks_cleanup_client",
+        )
+        await self._client.setup()
+
+    async def shutdown(self) -> None:
+        if self._client:
+            await self._client.shutdown()
+
+    @property
+    def _redis(self) -> aioredis.Redis:
+        assert self._client  # nosec
+        return self._client.redis
+
+    async def cleanup(self, lrt_namespace: LRTNamespace) -> None:
+        """removes Redis keys associated to the LRTNamespace if they exist"""
+        keys_to_remove: list[str] = [
+            x async for x in self._redis.scan_iter(f"{lrt_namespace}*")
+        ]
+        with log_context(
+            _logger,
+            logging.DEBUG,
+            msg=f"Removing {keys_to_remove=} from Redis for {lrt_namespace=}",
+        ):
+            if len(keys_to_remove) > 0:
+                await self._redis.delete(*keys_to_remove)

packages/service-library/src/servicelib/long_running_tasks/base_long_running_manager.py renamed to packages/service-library/src/servicelib/long_running_tasks/manager.py

@@ -11,7 +11,7 @@
 from .task import TasksManager
 
 
-class BaseLongRunningManager(ABC):
+class LongRunningManager(ABC):
     """
     Provides a commond inteface for aiohttp and fastapi services
     """

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

@@ -11,6 +11,7 @@
 from common_library.async_tools import cancel_wait_task
 from models_library.api_schemas_long_running_tasks.base import TaskProgress
 from pydantic import NonNegativeFloat, PositiveFloat
+from servicelib.utils import limited_gather
 from settings_library.redis import RedisDatabase, RedisSettings
 from tenacity import (
     AsyncRetrying,
@@ -21,7 +22,7 @@
 
 from ..background_task import create_periodic_task
 from ..logging_errors import create_troubleshootting_log_kwargs
-from ..logging_utils import log_context
+from ..logging_utils import log_catch, log_context
 from ..redis import RedisClientSDK, exclusive
 from ._redis_store import RedisStore
 from ._serialization import dumps
@@ -50,6 +51,7 @@
 _STATUS_UPDATE_CHECK_INTERNAL: Final[datetime.timedelta] = datetime.timedelta(seconds=1)
 _MAX_EXCLUSIVE_TASK_CANCEL_TIMEOUT: Final[NonNegativeFloat] = 5
 _TASK_REMOVAL_MAX_WAIT: Final[NonNegativeFloat] = 60
+_PARALLEL_TASKS_CANCELLATION: Final[int] = 5
 
 AllowedErrrors: TypeAlias = tuple[type[BaseException], ...]
 
@@ -205,34 +207,40 @@ async def setup(self) -> None:
         await self._started_event_task_tasks_monitor.wait()
 
     async def teardown(self) -> None:
-        # ensure all created tasks are cancelled
-        for tracked_task in await self._tasks_data.list_tasks_data():
-            with suppress(TaskNotFoundError):
+        # stop cancelled_tasks_removal
+        if self._task_cancelled_tasks_removal:
+            await cancel_wait_task(self._task_cancelled_tasks_removal)
+
+        # stopping only tasks that are handled by this manager
+        # otherwise it will cancel long running tasks that were running in diffierent processes
+        async def _remove_local_task(task_data: TaskData) -> None:
+            with log_catch(_logger, reraise=False):
                 await self.remove_task(
-                    tracked_task.task_id,
-                    tracked_task.task_context,
-                    wait_for_removal=True,
+                    task_data.task_id,
+                    task_data.task_context,
+                    wait_for_removal=False,
                 )
+                await self._attempt_to_remove_local_task(task_data.task_id)
+
+        await limited_gather(
+            *[
+                _remove_local_task(tracked_task)
+                for task_id in self._created_tasks
+                if (tracked_task := await self._tasks_data.get_task_data(task_id))
+                is not None
+            ],
+            log=_logger,
+            limit=_PARALLEL_TASKS_CANCELLATION,
+        )
 
-        for task in self._created_tasks.values():
-            _logger.warning(
-                "Task %s was not completed before shutdown, cancelling it",
-                task.get_name(),
-            )
-            await cancel_wait_task(task)
-
-        # stale_tasks_monitor
+        # stop stale_tasks_monitor
         if self._task_stale_tasks_monitor:
             await cancel_wait_task(
                 self._task_stale_tasks_monitor,
                 max_delay=_MAX_EXCLUSIVE_TASK_CANCEL_TIMEOUT,
             )
 
-        # cancelled_tasks_removal
-        if self._task_cancelled_tasks_removal:
-            await cancel_wait_task(self._task_cancelled_tasks_removal)
-
-        # tasks_monitor
+        # stop tasks_monitor
         if self._task_tasks_monitor:
             await cancel_wait_task(self._task_tasks_monitor)
 

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

@@ -11,8 +11,8 @@
 from faker import Faker
 from pytest_mock import MockerFixture
 from servicelib.logging_utils import log_catch
-from servicelib.long_running_tasks.base_long_running_manager import (
-    BaseLongRunningManager,
+from servicelib.long_running_tasks.manager import (
+    LongRunningManager,
 )
 from servicelib.long_running_tasks.models import LRTNamespace, TaskContext
 from servicelib.long_running_tasks.task import TasksManager
@@ -24,7 +24,7 @@
 _logger = logging.getLogger(__name__)
 
 
-class _TestingLongRunningManager(BaseLongRunningManager):
+class _TestingLongRunningManager(LongRunningManager):
     @staticmethod
     def get_task_context(request) -> TaskContext:
         _ = request
@@ -37,16 +37,16 @@ async def get_long_running_manager(
 ) -> AsyncIterator[
     Callable[
         [RedisSettings, RabbitSettings, LRTNamespace | None],
-        Awaitable[BaseLongRunningManager],
+        Awaitable[LongRunningManager],
     ]
 ]:
-    managers: list[BaseLongRunningManager] = []
+    managers: list[LongRunningManager] = []
 
     async def _(
         redis_settings: RedisSettings,
         rabbit_settings: RabbitSettings,
         lrt_namespace: LRTNamespace | None,
-    ) -> BaseLongRunningManager:
+    ) -> LongRunningManager:
         manager = _TestingLongRunningManager(
             stale_task_check_interval=timedelta(seconds=TEST_CHECK_STALE_INTERVAL_S),
             stale_task_detect_timeout=timedelta(seconds=TEST_CHECK_STALE_INTERVAL_S),

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

@@ -0,0 +1,83 @@
+# pylint:disable=redefined-outer-name
+
+from collections.abc import AsyncIterable, Callable
+from contextlib import AbstractAsyncContextManager
+
+import pytest
+from pydantic import TypeAdapter
+from servicelib.long_running_tasks._redis_store import RedisStore
+from servicelib.long_running_tasks.long_running_client_helper import (
+    LongRunningClientHelper,
+)
+from servicelib.long_running_tasks.models import LRTNamespace, TaskData
+from servicelib.redis._client import RedisClientSDK
+from settings_library.redis import RedisDatabase, RedisSettings
+
+
+@pytest.fixture
+def task_data() -> TaskData:
+    return TypeAdapter(TaskData).validate_python(
+        TaskData.model_json_schema()["examples"][0]
+    )
+
+
+@pytest.fixture
+def lrt_namespace() -> LRTNamespace:
+    return "TEST-NAMESPACE"
+
+
+@pytest.fixture
+async def store(
+    use_in_memory_redis: RedisSettings,
+    get_redis_client_sdk: Callable[
+        [RedisDatabase], AbstractAsyncContextManager[RedisClientSDK]
+    ],
+    lrt_namespace: LRTNamespace,
+) -> AsyncIterable[RedisStore]:
+    store = RedisStore(redis_settings=use_in_memory_redis, namespace=lrt_namespace)
+
+    await store.setup()
+    yield store
+    await store.shutdown()
+
+    # triggers cleanup of all redis data
+    async with get_redis_client_sdk(RedisDatabase.LONG_RUNNING_TASKS):
+        pass
+
+
+@pytest.fixture
+async def long_running_client_helper(
+    use_in_memory_redis: RedisSettings,
+) -> AsyncIterable[LongRunningClientHelper]:
+    helper = LongRunningClientHelper(redis_settings=use_in_memory_redis)
+
+    await helper.setup()
+    yield helper
+    await helper.shutdown()
+
+
+async def test_cleanup_namespace(
+    store: RedisStore,
+    task_data: TaskData,
+    long_running_client_helper: LongRunningClientHelper,
+    lrt_namespace: LRTNamespace,
+) -> None:
+    # create entries in both sides
+    await store.add_task_data(task_data.task_id, task_data)
+    await store.mark_task_for_removal(task_data.task_id, task_data.task_context)
+
+    # entries exit
+    assert await store.list_tasks_data() == [task_data]
+    assert await store.list_tasks_to_remove() == {
+        task_data.task_id: task_data.task_context
+    }
+
+    # removes
+    await long_running_client_helper.cleanup(lrt_namespace)
+
+    # entris were removed
+    assert await store.list_tasks_data() == []
+    assert await store.list_tasks_to_remove() == {}
+
+    # ensore it does not raise errors if there is nothing to remove
+    await long_running_client_helper.cleanup(lrt_namespace)