removed unsued

Author: Andrei Neagu

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

@@ -8,11 +8,10 @@
 from ..redis._client import RedisClientSDK
 from ..redis._utils import handle_redis_returns_union_types
 from ..utils import limited_gather
-from .models import LRTNamespace, TaskContext, TaskData, TaskId
+from .models import LRTNamespace, TaskData, TaskId
 
 _STORE_TYPE_TASK_DATA: Final[str] = "TD"
-_STORE_TYPE_CANCELLED_TASKS: Final[str] = "CT"
-_LIST_CONCURRENCY: Final[int] = 2
+_LIST_CONCURRENCY: Final[int] = 3
 
 
 def _to_redis_hash_mapping(data: dict[str, Any]) -> dict[str, str]:
@@ -52,11 +51,6 @@ def _get_redis_key_task_data_match(self) -> str:
     def _get_redis_task_data_key(self, task_id: TaskId) -> str:
         return f"{self.namespace}:{_STORE_TYPE_TASK_DATA}:{task_id}"
 
-    def _get_key_to_remove(self) -> str:
-        return f"{self.namespace}:{_STORE_TYPE_CANCELLED_TASKS}"
-
-    # TaskData
-
     async def get_task_data(self, task_id: TaskId) -> TaskData | None:
         result: dict[str, Any] = await handle_redis_returns_union_types(
             self._redis.hgetall(
@@ -117,28 +111,18 @@ async def delete_task_data(self, task_id: TaskId) -> None:
 
     # to cancel
 
-    async def mark_task_for_removal(
-        self, task_id: TaskId, with_task_context: TaskContext
-    ) -> None:
+    async def mark_task_for_removal(self, task_id: TaskId) -> None:
         await handle_redis_returns_union_types(
             self._redis.hset(
-                self._get_key_to_remove(), task_id, json_dumps(with_task_context)
+                self._get_redis_task_data_key(task_id),
+                mapping=_to_redis_hash_mapping({"marked_for_removal": True}),
             )
         )
 
     async def is_marked_for_removal(self, task_id: TaskId) -> bool:
-        result: bool = await handle_redis_returns_union_types(
-            self._redis.hexists(self._get_key_to_remove(), task_id)
-        )
-        return result
-
-    async def completed_task_removal(self, task_id: TaskId) -> None:
-        await handle_redis_returns_union_types(
-            self._redis.hdel(self._get_key_to_remove(), task_id)
-        )
-
-    async def list_tasks_to_remove(self) -> dict[TaskId, TaskContext]:
-        result: dict[str, str | None] = await handle_redis_returns_union_types(
-            self._redis.hgetall(self._get_key_to_remove())
+        result = await handle_redis_returns_union_types(
+            self._redis.hget(
+                self._get_redis_task_data_key(task_id), "marked_for_removal"
+            )
         )
-        return {task_id: json_loads(context) for task_id, context in result.items()}
+        return False if result is None else json_loads(result)

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

@@ -50,7 +50,7 @@ class TaskData(BaseModel):
     task_id: str
     task_progress: TaskProgress
     # NOTE: this context lifetime is with the tracked task (similar to aiohttp storage concept)
-    task_context: dict[str, Any]
+    task_context: TaskContext
     fire_and_forget: Annotated[
         bool,
         Field(
@@ -78,6 +78,10 @@ class TaskData(BaseModel):
     result_field: Annotated[
         ResultField | None, Field(description="the result of the task")
     ] = None
+    marked_for_removal: Annotated[
+        bool,
+        Field(description=("if True, indicates the task is marked for removal")),
+    ] = False
 
     model_config = ConfigDict(
         arbitrary_types_allowed=True,

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

@@ -300,11 +300,17 @@ async def _cancelled_tasks_removal(self) -> None:
         """
         self._started_event_task_cancelled_tasks_removal.set()
 
-        to_remove = await self._tasks_data.list_tasks_to_remove()
-        for task_id in to_remove:
-            await self._attempt_to_remove_local_task(task_id)
+        tasks_data = await self._tasks_data.list_tasks_data()
+        await limited_gather(
+            *(
+                self._attempt_to_remove_local_task(x.task_id)
+                for x in tasks_data
+                if x.marked_for_removal is True
+            ),
+            limit=_PARALLEL_TASKS_CANCELLATION,
+        )
 
-    async def _tasks_monitor(self) -> None:
+    async def _tasks_monitor(self) -> None:  # noqa: C901
         """
         A task which monitors locally running tasks and updates their status
         in the Redis store when they are done.
@@ -385,30 +391,21 @@ async def _tasks_monitor(self) -> None:
                     updates["result_field"] = result_field
                 await self._tasks_data.update_task_data(task_id, updates=updates)
 
-    async def _list_tasks_unfiltered(
-        self, with_task_context: TaskContext | None
-    ) -> list[TaskBase]:
+    async def list_tasks(self, with_task_context: TaskContext | None) -> list[TaskBase]:
         if not with_task_context:
             return [
                 TaskBase(task_id=task.task_id)
                 for task in (await self._tasks_data.list_tasks_data())
+                if task.marked_for_removal is False
             ]
 
         return [
             TaskBase(task_id=task.task_id)
             for task in (await self._tasks_data.list_tasks_data())
             if task.task_context == with_task_context
+            and task.marked_for_removal is False
         ]
 
-    async def list_tasks(self, with_task_context: TaskContext | None) -> list[TaskBase]:
-        tasks_to_filter = await self._list_tasks_unfiltered(with_task_context)
-
-        if len(tasks_to_filter) == 0:
-            return []
-
-        to_exclude = await self._tasks_data.list_tasks_to_remove()
-        return [r for r in tasks_to_filter if r.task_id not in to_exclude]
-
     async def _get_tracked_task(
         self, task_id: TaskId, with_task_context: TaskContext
     ) -> TaskData:
@@ -488,7 +485,6 @@ async def _attempt_to_remove_local_task(self, task_id: TaskId) -> None:
         task_to_cancel = self._created_tasks.pop(task_id, None)
         if task_to_cancel is not None:
             await cancel_wait_task(task_to_cancel)
-            await self._tasks_data.completed_task_removal(task_id)
             await self._tasks_data.delete_task_data(task_id)
 
     async def remove_task(
@@ -507,9 +503,7 @@ async def remove_task(
 
         tracked_task = await self._get_tracked_task(task_id, with_task_context)
 
-        await self._tasks_data.mark_task_for_removal(
-            tracked_task.task_id, tracked_task.task_context
-        )
+        await self._tasks_data.mark_task_for_removal(tracked_task.task_id)
 
         if not wait_for_removal:
             return

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

@@ -2,6 +2,7 @@
 
 from collections.abc import AsyncIterable, Callable
 from contextlib import AbstractAsyncContextManager
+from copy import deepcopy
 
 import pytest
 from pydantic import TypeAdapter
@@ -50,18 +51,14 @@ async def test_workflow(store: RedisStore, task_data: TaskData) -> None:
     assert await store.list_tasks_data() == []
 
     # cancelled tasks
-    assert await store.list_tasks_to_remove() == {}
+    await store.add_task_data(task_data.task_id, task_data)
 
     assert await store.is_marked_for_removal(task_data.task_id) is False
 
-    await store.mark_task_for_removal(task_data.task_id, task_data.task_context)
+    await store.mark_task_for_removal(task_data.task_id)
 
     assert await store.is_marked_for_removal(task_data.task_id) is True
 
-    assert await store.list_tasks_to_remove() == {
-        task_data.task_id: task_data.task_context
-    }
-
 
 @pytest.fixture
 async def redis_stores(
@@ -93,15 +90,15 @@ async def test_workflow_multiple_redis_stores_with_different_namespaces(
 
     for store in redis_stores:
         assert await store.list_tasks_data() == []
-        assert await store.list_tasks_to_remove() == {}
 
     for store in redis_stores:
         await store.add_task_data(task_data.task_id, task_data)
-        await store.mark_task_for_removal(task_data.task_id, {})
+        await store.mark_task_for_removal(task_data.task_id)
 
+    marked_as_removed_task_data = deepcopy(task_data)
+    marked_as_removed_task_data.marked_for_removal = True
     for store in redis_stores:
-        assert await store.list_tasks_data() == [task_data]
-        assert await store.list_tasks_to_remove() == {task_data.task_id: {}}
+        assert await store.list_tasks_data() == [marked_as_removed_task_data]
 
     for store in redis_stores:
         await store.delete_task_data(task_data.task_id)

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

@@ -2,6 +2,7 @@
 
 from collections.abc import AsyncIterable, Callable
 from contextlib import AbstractAsyncContextManager
+from copy import deepcopy
 
 import pytest
 from pydantic import TypeAdapter
@@ -64,20 +65,18 @@ async def test_cleanup_namespace(
 ) -> 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)
+    await store.mark_task_for_removal(task_data.task_id)
 
     # 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
-    }
+    marked_for_removal = deepcopy(task_data)
+    marked_for_removal.marked_for_removal = True
+    assert await store.list_tasks_data() == [marked_for_removal]
 
     # 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)

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

@@ -549,7 +549,7 @@ async def test__cancelled_tasks_worker_equivalent_of_cancellation_from_a_differe
         task_context=empty_context,
     )
     await long_running_manager.tasks_manager._tasks_data.mark_task_for_removal(  # noqa: SLF001
-        task_id, with_task_context=empty_context
+        task_id
     )
 
     async for attempt in AsyncRetrying(**_RETRY_PARAMS):