✨ Enhance lifespan error handling and add lifespan call tracking

Author: pcrespov

packages/service-library/src/servicelib/fastapi/lifespan_utils.py

@@ -1,12 +1,44 @@
+from typing import Final
+
 from common_library.errors_classes import OsparcErrorMixin
+from fastapi import FastAPI
+from fastapi_lifespan_manager import State
 
 
 class LifespanError(OsparcErrorMixin, RuntimeError): ...
 
 
 class LifespanOnStartupError(LifespanError):
-    msg_template = "Failed during startup of {module}"
+    msg_template = "Failed during startup of {lifespan_name}"
 
 
 class LifespanOnShutdownError(LifespanError):
-    msg_template = "Failed during shutdown of {module}"
+    msg_template = "Failed during shutdown of {lifespan_name}"
+
+
+class LifespanAlreadyCalledError(LifespanError):
+    msg_template = "The lifespan '{lifespan_name}' has already been called."
+
+
+_CALLED_LIFESPANS_KEY: Final[str] = "_CALLED_LIFESPANS"
+
+
+def is_lifespan_called(state: State, lifespan_name: str) -> bool:
+    called_lifespans = state.get(_CALLED_LIFESPANS_KEY, set())
+    return lifespan_name in called_lifespans
+
+
+def record_lifespan_called_once(state: State, lifespan_name: str) -> State:
+    """Validates if a lifespan has already been called and records it in the state.
+    Raises LifespanAlreadyCalledError if the lifespan has already been called.
+    """
+    assert not isinstance(  # nosec
+        state, FastAPI
+    ), "TIP: lifespan func has (app, state) positional arguments"
+
+    if is_lifespan_called(state, lifespan_name):
+        raise LifespanAlreadyCalledError(lifespan_name=lifespan_name)
+
+    called_lifespans = state.get(_CALLED_LIFESPANS_KEY, set())
+    called_lifespans.add(lifespan_name)
+    return {_CALLED_LIFESPANS_KEY: called_lifespans}

packages/service-library/tests/fastapi/test_lifespan_utils.py

@@ -16,8 +16,10 @@
 from pytest_mock import MockerFixture
 from pytest_simcore.helpers.logging_tools import log_context
 from servicelib.fastapi.lifespan_utils import (
+    LifespanAlreadyCalledError,
     LifespanOnShutdownError,
     LifespanOnStartupError,
+    record_lifespan_called_once,
 )
 
 
@@ -186,7 +188,7 @@ async def lifespan_failing_on_startup(app: FastAPI) -> AsyncIterator[State]:
                 startup_step(_name)
             except RuntimeError as exc:
                 handle_error(_name, exc)
-                raise LifespanOnStartupError(module=_name) from exc
+                raise LifespanOnStartupError(lifespan_name=_name) from exc
             yield {}
             shutdown_step(_name)
 
@@ -201,7 +203,7 @@ async def lifespan_failing_on_shutdown(app: FastAPI) -> AsyncIterator[State]:
                 shutdown_step(_name)
             except RuntimeError as exc:
                 handle_error(_name, exc)
-                raise LifespanOnShutdownError(module=_name) from exc
+                raise LifespanOnShutdownError(lifespan_name=_name) from exc
 
     return {
         "startup_step": startup_step,
@@ -228,7 +230,7 @@ async def test_app_lifespan_with_error_on_startup(
     assert not failing_lifespan_manager["startup_step"].called
     assert not failing_lifespan_manager["shutdown_step"].called
     assert exception.error_context() == {
-        "module": "lifespan_failing_on_startup",
+        "lifespan_name": "lifespan_failing_on_startup",
         "message": "Failed during startup of lifespan_failing_on_startup",
         "code": "RuntimeError.LifespanError.LifespanOnStartupError",
     }
@@ -250,7 +252,35 @@ async def test_app_lifespan_with_error_on_shutdown(
     assert failing_lifespan_manager["startup_step"].called
     assert not failing_lifespan_manager["shutdown_step"].called
     assert exception.error_context() == {
-        "module": "lifespan_failing_on_shutdown",
+        "lifespan_name": "lifespan_failing_on_shutdown",
         "message": "Failed during shutdown of lifespan_failing_on_shutdown",
         "code": "RuntimeError.LifespanError.LifespanOnShutdownError",
     }
+
+
+async def test_lifespan_called_more_than_once(is_pdb_enabled: bool):
+    state = {}
+
+    app_lifespan = LifespanManager()
+
+    @app_lifespan.add
+    async def _one(_, state: State) -> AsyncIterator[State]:
+        called_state = record_lifespan_called_once(state, "test_lifespan_one")
+        yield {"other": 0, **called_state}
+
+    @app_lifespan.add
+    async def _two(_, state: State) -> AsyncIterator[State]:
+        called_state = record_lifespan_called_once(state, "test_lifespan_two")
+        yield {"something": 0, **called_state}
+
+    app_lifespan.add(_one)  # added "by mistake"
+
+    with pytest.raises(LifespanAlreadyCalledError) as err_info:
+        async with ASGILifespanManager(
+            FastAPI(lifespan=app_lifespan),
+            startup_timeout=None if is_pdb_enabled else 10,
+            shutdown_timeout=None if is_pdb_enabled else 10,
+        ):
+            ...
+
+    assert err_info.value.lifespan_name == "test_lifespan_one"