🎨 Refactor background task management: Introduce setup_periodic_task utility for cleaner task setup and teardown

Author: pcrespov

services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_api_keys.py

@@ -3,26 +3,22 @@
 
 """
 
-import asyncio
 import logging
-from collections.abc import AsyncIterator, Callable
+from collections.abc import AsyncIterator
 from datetime import timedelta
 
 from aiohttp import web
-from servicelib.async_utils import cancel_wait_task
 from servicelib.background_task_utils import exclusive_periodic
 from servicelib.logging_utils import log_context
 from simcore_service_webserver.redis import get_redis_lock_manager_client_sdk
 
 from ..api_keys import api_keys_service
+from ._tasks_utils import CleanupContextFunc, setup_periodic_task
 
 _logger = logging.getLogger(__name__)
 
-CleanupContextFunc = Callable[[web.Application], AsyncIterator[None]]
 
-
-async def _run_task(app: web.Application):
-    """Checks expiration dates and updates user status"""
+async def _prune_expired_api_keys(app: web.Application):
     if deleted := await api_keys_service.prune_expired_api_keys(app):
         # broadcast force logout of user_id
         for api_key in deleted:
@@ -46,24 +42,9 @@ async def _cleanup_ctx_fun(app: web.Application) -> AsyncIterator[None]:
         )
         async def _prune_expired_api_keys_periodically() -> None:
             with log_context(_logger, logging.INFO, "Pruning expired API keys"):
-                await _run_task(app)
-
-        # setup
-        task_name = _prune_expired_api_keys_periodically.__name__
-
-        task = asyncio.create_task(
-            _prune_expired_api_keys_periodically(),
-            name=task_name,
-        )
-
-        # prevents premature garbage collection of the task
-        app_task_key = f"tasks.{task_name}"
-        app[app_task_key] = task
-
-        yield
+                await _prune_expired_api_keys(app)
 
-        # tear-down
-        await cancel_wait_task(task)
-        app.pop(app_task_key, None)
+        async for _ in setup_periodic_task(app, _prune_expired_api_keys_periodically):
+            yield
 
     return _cleanup_ctx_fun

services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_core.py

@@ -4,24 +4,21 @@
 Specifics of the gc implementation should go into garbage_collector_core.py
 """
 
-import asyncio
 import logging
-from collections.abc import AsyncIterator, Callable
+from collections.abc import AsyncIterator
 from datetime import timedelta
 
 from aiohttp import web
-from servicelib.async_utils import cancel_wait_task
 from servicelib.background_task_utils import exclusive_periodic
 from servicelib.logging_utils import log_context
 from simcore_service_webserver.redis import get_redis_lock_manager_client_sdk
 
 from ._core import collect_garbage
+from ._tasks_utils import CleanupContextFunc, setup_periodic_task
 from .settings import GarbageCollectorSettings, get_plugin_settings
 
 _logger = logging.getLogger(__name__)
 
-CleanupContextFunc = Callable[[web.Application], AsyncIterator[None]]
-
 
 def create_background_task_for_garbage_collection() -> CleanupContextFunc:
 
@@ -38,22 +35,7 @@ async def _collect_garbage_periodically() -> None:
             with log_context(_logger, logging.INFO, "Garbage collect cycle"):
                 await collect_garbage(app)
 
-        # setup
-        task_name = _collect_garbage_periodically.__name__
-
-        task = asyncio.create_task(
-            _collect_garbage_periodically(),
-            name=task_name,
-        )
-
-        # prevents premature garbage collection of the task
-        app_task_key = f"tasks.{task_name}"
-        app[app_task_key] = task
-
-        yield
-
-        # tear-down
-        await cancel_wait_task(task)
-        app.pop(app_task_key, None)
+        async for _ in setup_periodic_task(app, _collect_garbage_periodically):
+            yield
 
     return _cleanup_ctx_fun

services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_trash.py

@@ -3,23 +3,20 @@
 
 """
 
-import asyncio
 import logging
-from collections.abc import AsyncIterator, Callable
+from collections.abc import AsyncIterator
 from datetime import timedelta
 
 from aiohttp import web
-from servicelib.async_utils import cancel_wait_task
 from servicelib.background_task_utils import exclusive_periodic
 from servicelib.logging_utils import log_context
 from simcore_service_webserver.redis import get_redis_lock_manager_client_sdk
 
 from ..trash import trash_service
+from ._tasks_utils import CleanupContextFunc, setup_periodic_task
 
 _logger = logging.getLogger(__name__)
 
-CleanupContextFunc = Callable[[web.Application], AsyncIterator[None]]
-
 
 def create_background_task_to_prune_trash(wait_s: float) -> CleanupContextFunc:
 
@@ -35,22 +32,7 @@ async def _prune_trash_periodically() -> None:
             with log_context(_logger, logging.INFO, "Deleting expired trashed items"):
                 await trash_service.safe_delete_expired_trash_as_admin(app)
 
-        # setup
-        task_name = _prune_trash_periodically.__name__
-
-        task = asyncio.create_task(
-            _prune_trash_periodically(),
-            name=task_name,
-        )
-
-        # prevents premature garbage collection of the task
-        app_task_key = f"tasks.{task_name}"
-        app[app_task_key] = task
-
-        yield
-
-        # tear-down
-        await cancel_wait_task(task)
-        app.pop(app_task_key, None)
+        async for _ in setup_periodic_task(app, _prune_trash_periodically):
+            yield
 
     return _cleanup_ctx_fun

services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_users.py

@@ -3,21 +3,20 @@
 
 """
 
-import asyncio
 import logging
 from collections.abc import AsyncIterator, Callable
 from datetime import timedelta
 
 from aiohttp import web
 from models_library.users import UserID
-from servicelib.async_utils import cancel_wait_task
 from servicelib.background_task_utils import exclusive_periodic
 from servicelib.logging_utils import get_log_record_extra, log_context
 from simcore_service_webserver.redis import get_redis_lock_manager_client_sdk
 
 from ..login import login_service
 from ..security import security_service
 from ..users.api import update_expired_users
+from ._tasks_utils import CleanupContextFunc, setup_periodic_task
 
 _logger = logging.getLogger(__name__)
 
@@ -84,22 +83,7 @@ async def _update_expired_users_periodically() -> None:
             with log_context(_logger, logging.INFO, "Updating expired users"):
                 await _update_expired_users(app)
 
-        # setup
-        task_name = _update_expired_users_periodically.__name__
-
-        task = asyncio.create_task(
-            _update_expired_users_periodically(),
-            name=task_name,
-        )
-
-        # prevents premature garbage collection of the task
-        app_task_key = f"tasks.{task_name}"
-        app[app_task_key] = task
-
-        yield
-
-        # tear-down
-        await cancel_wait_task(task)
-        app.pop(app_task_key, None)
+        async for _ in setup_periodic_task(app, _update_expired_users_periodically):
+            yield
 
     return _cleanup_ctx_fun

services/web/server/src/simcore_service_webserver/garbage_collector/_tasks_utils.py

@@ -0,0 +1,41 @@
+"""
+Common utilities for background task management in garbage collector
+"""
+
+import asyncio
+from collections.abc import AsyncIterator, Callable, Coroutine
+
+from aiohttp import web
+from servicelib.async_utils import cancel_wait_task
+
+CleanupContextFunc = Callable[[web.Application], AsyncIterator[None]]
+
+
+async def setup_periodic_task(
+    app: web.Application,
+    periodic_task_coro: Callable[[], Coroutine[None, None, None]],
+) -> AsyncIterator[None]:
+    """
+    Generic setup and teardown for periodic background tasks.
+
+    Args:
+        app: The aiohttp web application
+        periodic_task_coro: The periodic task coroutine function (already decorated with @exclusive_periodic)
+    """
+    # setup
+    task_name = periodic_task_coro.__name__
+
+    task = asyncio.create_task(
+        periodic_task_coro(),
+        name=task_name,
+    )
+
+    # prevents premature garbage collection of the task
+    app_task_key = f"tasks.{task_name}"
+    app[app_task_key] = task
+
+    yield
+
+    # tear-down
+    await cancel_wait_task(task)
+    app.pop(app_task_key, None)