dsm

Author: sanderegg

services/storage/src/simcore_service_storage/dsm_cleaner.py

@@ -20,16 +20,16 @@
 
 import asyncio
 import logging
+from collections.abc import Awaitable, Callable
 from datetime import timedelta
 from typing import cast
 
+from fastapi import FastAPI
 from servicelib.async_utils import cancel_wait_task
 from servicelib.background_task_utils import exclusive_periodic
-from servicelib.logging_utils import log_catch, log_context
 
-from .constants import APP_CONFIG_KEY, APP_DSM_KEY
-from .core.settings import ApplicationSettings
-from .dsm_factory import DataManagerProvider
+from .core.settings import get_application_settings
+from .dsm import get_dsm_provider
 from .modules.redis import get_redis_client
 from .simcore_s3_dsm import SimcoreS3DataManager
 
@@ -40,35 +40,36 @@
 
 async def dsm_cleaner_task(app: FastAPI) -> None:
     _logger.info("starting dsm cleaner task...")
-    dsm: DataManagerProvider = app[APP_DSM_KEY]
+    dsm = get_dsm_provider(app)
     simcore_s3_dsm: SimcoreS3DataManager = cast(
         SimcoreS3DataManager, dsm.get(SimcoreS3DataManager.get_location_id())
     )
     await simcore_s3_dsm.clean_expired_uploads()
 
 
-def setup_dsm_cleaner(app: FastAPI):
-    async def _setup(app: FastAPI):
-        with (
-            log_context(_logger, logging.INFO, msg="setup dsm cleaner"),
-            log_catch(_logger, reraise=False),
-        ):
-            cfg: ApplicationSettings = app[APP_CONFIG_KEY]
-            assert cfg.STORAGE_CLEANER_INTERVAL_S  # nosec
+def setup_dsm_cleaner(app: FastAPI) -> None:
+    async def _on_startup(app: FastAPI) -> None:
+        cfg = get_application_settings(app)
+        assert cfg.STORAGE_CLEANER_INTERVAL_S  # nosec
 
-            @exclusive_periodic(
-                get_redis_client(app),
-                task_interval=timedelta(seconds=cfg.STORAGE_CLEANER_INTERVAL_S),
-                retry_after=timedelta(minutes=5),
-            )
-            async def _periodic_dsm_clean() -> None:
-                await dsm_cleaner_task(app)
+        @exclusive_periodic(
+            get_redis_client(app),
+            task_interval=timedelta(seconds=cfg.STORAGE_CLEANER_INTERVAL_S),
+            retry_after=timedelta(minutes=5),
+        )
+        async def _periodic_dsm_clean() -> None:
+            await dsm_cleaner_task(app)
 
-            storage_background_task = asyncio.create_task(
-                _periodic_dsm_clean(), name=_TASK_NAME_PERIODICALY_CLEAN_DSM
-            )
-            yield
+        app.state.dsm_cleaner_task = asyncio.create_task(
+            _periodic_dsm_clean(), name=_TASK_NAME_PERIODICALY_CLEAN_DSM
+        )
 
-            await cancel_wait_task(storage_background_task)
+    def _on_shutdown(app: FastAPI) -> Callable[[], Awaitable[None]]:
+        async def _stop() -> None:
+            assert isinstance(app.state.dsm_cleaner_task, asyncio.Task)  # nosec
+            await cancel_wait_task(app.state.dsm_cleaner_task)
 
-    app.cleanup_ctx.append(_setup)
+        return _stop
+
+    app.add_event_handler("startup", _on_startup)
+    app.add_event_handler("shutdown", _on_shutdown(app))