added in memory store to TasksManager

Author: Andrei Neagu

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

@@ -28,7 +28,7 @@ async def list_tasks(request: web.Request) -> web.Response:
                 result_href=f"{request.app.router['get_task_result'].url_for(task_id=t.task_id)}",
                 abort_href=f"{request.app.router['cancel_and_delete_task'].url_for(task_id=t.task_id)}",
             )
-            for t in lrt_api.list_tasks(
+            for t in await lrt_api.list_tasks(
                 long_running_manager.tasks_manager,
                 long_running_manager.get_task_context(request),
             )
@@ -41,7 +41,7 @@ async def get_task_status(request: web.Request) -> web.Response:
     path_params = parse_request_path_parameters_as(_PathParam, request)
     long_running_manager = get_long_running_manager(request.app)
 
-    task_status: TaskStatus = lrt_api.get_task_status(
+    task_status: TaskStatus = await lrt_api.get_task_status(
         long_running_manager.tasks_manager,
         long_running_manager.get_task_context(request),
         path_params.task_id,

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

@@ -29,7 +29,7 @@ async def list_tasks(
                 request.url_for("cancel_and_delete_task", task_id=t.task_id)
             ),
         )
-        for t in lrt_api.list_tasks(
+        for t in await lrt_api.list_tasks(
             long_running_manager.tasks_manager, task_context=None
         )
     ]
@@ -51,7 +51,7 @@ async def get_task_status(
     ],
 ) -> TaskStatus:
     assert request  # nosec
-    return lrt_api.get_task_status(
+    return await lrt_api.get_task_status(
         long_running_manager.tasks_manager, task_context=None, task_id=task_id
     )
 

packages/service-library/src/servicelib/long_running_tasks/_store/__init__.py

@@ -0,0 +1,22 @@
+from abc import abstractmethod
+
+from ..models import TaskId, TrackedTask
+
+
+class BaseStore:
+
+    @abstractmethod
+    async def get(self, key: TaskId) -> TrackedTask | None:
+        """Retrieve a tracked task by its key."""
+
+    @abstractmethod
+    async def set(self, key: TaskId, value: TrackedTask) -> None:
+        """Set a tracked task with its key."""
+
+    @abstractmethod
+    async def list(self) -> list[TrackedTask]:
+        """List all tracked tasks."""
+
+    @abstractmethod
+    async def delete(self, key: TaskId) -> None:
+        """Delete a tracked task by its key."""

packages/service-library/src/servicelib/long_running_tasks/_store/base.py

@@ -0,0 +1,19 @@
+from ..models import TaskId, TrackedTask
+from .base import BaseStore
+
+
+class InMemoryStore(BaseStore):
+    def __init__(self):
+        self._store: dict[TaskId, TrackedTask] = {}
+
+    async def get(self, key: TaskId) -> TrackedTask | None:
+        return self._store.get(key, None)
+
+    async def set(self, key: TaskId, value: TrackedTask) -> None:
+        self._store[key] = value
+
+    async def list(self) -> list[TrackedTask]:
+        return list(self._store.values())
+
+    async def delete(self, key: TaskId) -> None:
+        self._store.pop(key, None)

packages/service-library/src/servicelib/long_running_tasks/_store/in_memory.py

@@ -46,7 +46,7 @@ async def start_task(
     Returns:
         TaskId: the task unique identifier
     """
-    return tasks_manager.start_task(
+    return await tasks_manager.start_task(
         registered_task_name,
         unique=unique,
         task_context=task_context,
@@ -56,17 +56,17 @@ async def start_task(
     )
 
 
-def list_tasks(
+async def list_tasks(
     tasks_manager: TasksManager, task_context: TaskContext | None
 ) -> list[TaskBase]:
-    return tasks_manager.list_tasks(with_task_context=task_context)
+    return await tasks_manager.list_tasks(with_task_context=task_context)
 
 
-def get_task_status(
+async def get_task_status(
     tasks_manager: TasksManager, task_context: TaskContext | None, task_id: TaskId
 ) -> TaskStatus:
     """returns the status of a task"""
-    return tasks_manager.get_task_status(
+    return await tasks_manager.get_task_status(
         task_id=task_id, with_task_context=task_context
     )
 
@@ -75,7 +75,7 @@ async def get_task_result(
     tasks_manager: TasksManager, task_context: TaskContext | None, task_id: TaskId
 ) -> Any:
     try:
-        task_result = tasks_manager.get_task_result(
+        task_result = await tasks_manager.get_task_result(
             task_id, with_task_context=task_context
         )
         await tasks_manager.remove_task(

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

@@ -15,6 +15,8 @@
 from servicelib.background_task import create_periodic_task
 from servicelib.logging_utils import log_catch
 
+from ._store.base import BaseStore
+from ._store.in_memory import InMemoryStore
 from .errors import (
     TaskAlreadyRunningError,
     TaskCancelledError,
@@ -38,7 +40,6 @@
 
 RegisteredTaskName: TypeAlias = str
 Namespace: TypeAlias = str
-TrackedTaskGroupDict: TypeAlias = dict[TaskId, TrackedTask]
 TaskContext: TypeAlias = dict[str, Any]
 
 
@@ -68,28 +69,28 @@ async def _await_task(task: asyncio.Task) -> None:
     await task
 
 
-def _get_tasks_to_remove(
-    tracked_tasks: TrackedTaskGroupDict,
+async def _get_tasks_to_remove(
+    tracked_tasks: BaseStore,
     stale_task_detect_timeout_s: PositiveFloat,
 ) -> list[TaskId]:
     utc_now = datetime.datetime.now(tz=datetime.UTC)
 
     tasks_to_remove: list[TaskId] = []
 
-    for task_id, tracked_task in tracked_tasks.items():
+    for tracked_task in await tracked_tasks.list():
         if tracked_task.fire_and_forget:
             continue
 
         if tracked_task.last_status_check is None:
             # the task just added or never received a poll request
             elapsed_from_start = (utc_now - tracked_task.started).seconds
             if elapsed_from_start > stale_task_detect_timeout_s:
-                tasks_to_remove.append(task_id)
+                tasks_to_remove.append(tracked_task.task_id)
         else:
             # the task status was already queried by the client
             elapsed_from_last_poll = (utc_now - tracked_task.last_status_check).seconds
             if elapsed_from_last_poll > stale_task_detect_timeout_s:
-                tasks_to_remove.append(task_id)
+                tasks_to_remove.append(tracked_task.task_id)
     return tasks_to_remove
 
 
@@ -103,10 +104,11 @@ def __init__(
         stale_task_check_interval: datetime.timedelta,
         stale_task_detect_timeout: datetime.timedelta,
         namespace: Namespace = _DEFAULT_NAMESPACE,
+        # TODO: inject a Redis connection
     ):
         self.namespace = namespace
         # Task groups: Every taskname maps to multiple asyncio.Task within TrackedTask model
-        self._tracked_tasks: TrackedTaskGroupDict = {}
+        self._tracked_tasks: BaseStore = InMemoryStore()
 
         self.stale_task_check_interval = stale_task_check_interval
         self.stale_task_detect_timeout_s: PositiveFloat = (
@@ -125,7 +127,7 @@ async def setup(self) -> None:
     async def teardown(self) -> None:
         task_ids_to_remove: deque[TaskId] = deque()
 
-        for tracked_task in self._tracked_tasks.values():
+        for tracked_task in await self._tracked_tasks.list():
             task_ids_to_remove.append(tracked_task.task_id)
 
         for task_id in task_ids_to_remove:
@@ -155,7 +157,7 @@ async def _stale_tasks_monitor_worker(self) -> None:
         # Since we own the client, we assume (for now) this
         # will not be the case.
 
-        tasks_to_remove = _get_tasks_to_remove(
+        tasks_to_remove = await _get_tasks_to_remove(
             self._tracked_tasks, self.stale_task_detect_timeout_s
         )
 
@@ -169,25 +171,28 @@ async def _stale_tasks_monitor_worker(self) -> None:
             _logger.warning(
                 "Removing stale task '%s' with status '%s'",
                 task_id,
-                self.get_task_status(task_id, with_task_context=None).model_dump_json(),
+                (
+                    await self.get_task_status(task_id, with_task_context=None)
+                ).model_dump_json(),
             )
             await self.remove_task(
                 task_id, with_task_context=None, reraise_errors=False
             )
 
-    def list_tasks(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 self._tracked_tasks.values()
+                TaskBase(task_id=task.task_id)
+                for task in (await self._tracked_tasks.list())
             ]
 
         return [
             TaskBase(task_id=task.task_id)
-            for task in self._tracked_tasks.values()
+            for task in (await self._tracked_tasks.list())
             if task.task_context == with_task_context
         ]
 
-    def _add_task(
+    async def _add_task(
         self,
         task: asyncio.Task,
         task_progress: TaskProgress,
@@ -204,24 +209,24 @@ def _add_task(
             task_context=task_context,
             fire_and_forget=fire_and_forget,
         )
-        self._tracked_tasks[task_id] = tracked_task
+        await self._tracked_tasks.set(task_id, tracked_task)
 
         return tracked_task
 
-    def _get_tracked_task(
+    async def _get_tracked_task(
         self, task_id: TaskId, with_task_context: TaskContext | None
     ) -> TrackedTask:
-        if task_id not in self._tracked_tasks:
-            raise TaskNotFoundError(task_id=task_id)
+        task = await self._tracked_tasks.get(task_id)
 
-        task = self._tracked_tasks[task_id]
+        if task is None:
+            raise TaskNotFoundError(task_id=task_id)
 
         if with_task_context and task.task_context != with_task_context:
             raise TaskNotFoundError(task_id=task_id)
 
         return task
 
-    def get_task_status(
+    async def get_task_status(
         self, task_id: TaskId, with_task_context: TaskContext | None
     ) -> TaskStatus:
         """
@@ -230,7 +235,9 @@ def get_task_status(
 
         raises TaskNotFoundError if the task cannot be found
         """
-        tracked_task: TrackedTask = self._get_tracked_task(task_id, with_task_context)
+        tracked_task: TrackedTask = await self._get_tracked_task(
+            task_id, with_task_context
+        )
         tracked_task.last_status_check = datetime.datetime.now(tz=datetime.UTC)
 
         task = tracked_task.task
@@ -244,7 +251,7 @@ def get_task_status(
             }
         )
 
-    def get_task_result(
+    async def get_task_result(
         self, task_id: TaskId, with_task_context: TaskContext | None
     ) -> Any:
         """
@@ -254,7 +261,7 @@ def get_task_result(
         raises TaskCancelledError if the task was cancelled
         raises TaskNotCompletedError if the task is not completed
         """
-        tracked_task = self._get_tracked_task(task_id, with_task_context)
+        tracked_task = await self._get_tracked_task(task_id, with_task_context)
 
         try:
             return tracked_task.task.result()
@@ -273,7 +280,7 @@ async def cancel_task(
 
         raises TaskNotFoundError if the task cannot be found
         """
-        tracked_task = self._get_tracked_task(task_id, with_task_context)
+        tracked_task = await self._get_tracked_task(task_id, with_task_context)
         await self._cancel_tracked_task(tracked_task.task, task_id, reraise_errors=True)
 
     @staticmethod
@@ -317,7 +324,7 @@ async def remove_task(
     ) -> None:
         """cancels and removes task"""
         try:
-            tracked_task = self._get_tracked_task(task_id, with_task_context)
+            tracked_task = await self._get_tracked_task(task_id, with_task_context)
         except TaskNotFoundError:
             if reraise_errors:
                 raise
@@ -327,13 +334,13 @@ async def remove_task(
                 tracked_task.task, task_id, reraise_errors=reraise_errors
             )
         finally:
-            del self._tracked_tasks[task_id]
+            await self._tracked_tasks.delete(task_id)
 
     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}"
 
-    def start_task(
+    async def start_task(
         self,
         registered_task_name: RegisteredTaskName,
         *,
@@ -358,9 +365,10 @@ def start_task(
         task_id = self._get_task_id(task_name, is_unique=unique)
 
         # only one unique task can be running
-        if unique and task_id in self._tracked_tasks:
+        queried_task = await self._tracked_tasks.get(task_id)
+        if unique and queried_task is not None:
             raise TaskAlreadyRunningError(
-                task_name=task_name, managed_task=self._tracked_tasks[task_id]
+                task_name=task_name, managed_task=queried_task
             )
 
         task_progress = TaskProgress.create(task_id=task_id)
@@ -377,7 +385,7 @@ async def _progress_task(progress: TaskProgress, handler: TaskProtocol):
             _progress_task(task_progress, task), name=task_name
         )
 
-        tracked_task = self._add_task(
+        tracked_task = await self._add_task(
             task=async_task,
             task_progress=task_progress,
             task_context=task_context or {},