fix: use lifespan

Author: giancarloromeo

packages/celery-library/src/celery_library/signals.py

@@ -26,8 +26,6 @@ def _init(startup_complete_event: threading.Event) -> None:
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
 
-        shutdown_event = asyncio.Event()
-
         async def _setup_task_manager():
             assert sender.app  # nosec
             assert isinstance(sender.app, Celery)  # nosec
@@ -42,9 +40,7 @@ async def _setup_task_manager():
         app_server.event_loop = loop
 
         loop.run_until_complete(_setup_task_manager())
-        loop.run_until_complete(
-            app_server.startup(startup_complete_event, shutdown_event)
-        )
+        loop.run_until_complete(app_server.lifespan(startup_complete_event))
 
     thread = threading.Thread(
         group=None,
@@ -63,4 +59,4 @@ def on_worker_shutdown(sender, **_kwargs) -> None:
         assert isinstance(sender.app, Celery)
         app_server = get_app_server(sender.app)
 
-        app_server.event_loop.run_until_complete(app_server.shutdown())
+        app_server.shutdown_event.set()

packages/service-library/src/servicelib/celery/app_server.py

@@ -12,7 +12,7 @@
 class BaseAppServer(ABC, Generic[T]):
     def __init__(self, app: T) -> None:
         self._app: T = app
-        self._shutdown_event: asyncio.Event | None = None
+        self._shutdown_event: asyncio.Event = asyncio.Event()
 
     @property
     def app(self) -> T:
@@ -26,6 +26,10 @@ def event_loop(self) -> AbstractEventLoop:
     def event_loop(self, loop: AbstractEventLoop) -> None:
         self._event_loop = loop
 
+    @property
+    def shutdown_event(self) -> asyncio.Event:
+        return self._shutdown_event
+
     @property
     def task_manager(self) -> TaskManager:
         return self._task_manager
@@ -35,23 +39,8 @@ def task_manager(self, manager: TaskManager) -> None:
         self._task_manager = manager
 
     @abstractmethod
-    async def on_startup(self) -> None:
-        raise NotImplementedError
-
-    async def startup(
-        self, startup_completed_event: threading.Event, shutdown_event: asyncio.Event
+    async def lifespan(
+        self,
+        startup_completed_event: threading.Event,
     ) -> None:
-        self._shutdown_event = shutdown_event
-        await self.on_startup()
-        startup_completed_event.set()
-        await self._shutdown_event.wait()
-
-    @abstractmethod
-    async def on_shutdown(self) -> None:
         raise NotImplementedError
-
-    async def shutdown(self) -> None:
-        if self._shutdown_event is not None:
-            self._shutdown_event.set()
-
-        await self.on_shutdown()

packages/service-library/src/servicelib/fastapi/celery/app_server.py

@@ -1,4 +1,7 @@
+import asyncio
 import datetime
+import logging
+import threading
 from typing import Final
 
 from asgi_lifespan import LifespanManager
@@ -8,21 +11,23 @@
 
 _SHUTDOWN_TIMEOUT: Final[float] = datetime.timedelta(seconds=10).total_seconds()
 
+_logger = logging.getLogger(__name__)
+
 
 class FastAPIAppServer(BaseAppServer[FastAPI]):
     def __init__(self, app: FastAPI):
         super().__init__(app)
         self._lifespan_manager: LifespanManager | None = None
 
-    async def on_startup(self) -> None:
-        self._lifespan_manager = LifespanManager(
+    async def lifespan(self, startup_completed_event: threading.Event) -> None:
+        async with LifespanManager(
             self.app,
             startup_timeout=None,  # waits for full app initialization (DB migrations, etc.)
             shutdown_timeout=_SHUTDOWN_TIMEOUT,
-        )
-        await self._lifespan_manager.__aenter__()
-
-    async def on_shutdown(self) -> None:
-        if self._lifespan_manager is None:
-            return
-        await self._lifespan_manager.__aexit__(None, None, None)
+        ):
+            try:
+                _logger.info("fastapi app initialized")
+                startup_completed_event.set()
+                await self.shutdown_event.wait()
+            except asyncio.CancelledError:
+                _logger.warning("lifespan task cancelled")