Fix writer startup with starlette lifecycle (#2300)

* Fix writer startup with starlette lifecycle

* zipp update

* Lint

nucliadb/requirements.lock.txt

@@ -111,4 +111,4 @@ urllib3==1.26.19
 uvicorn==0.18.3
 wrapt==1.16.0
 yarl==1.9.4
-zipp==3.17.0
+zipp==3.19.2

nucliadb/src/nucliadb/common/context/fastapi.py

@@ -38,7 +38,7 @@ async def inject_app_context(app: FastAPI):
             route.app.state.context = context
 
     await context.initialize()
-    yield
+    yield context
     await context.finalize()
 
 

nucliadb/src/nucliadb/writer/app.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-import functools
 import importlib.metadata
 
 from fastapi import FastAPI
@@ -28,7 +27,6 @@
 from starlette.requests import ClientDisconnect
 from starlette.responses import HTMLResponse
 
-from nucliadb.common.context.fastapi import get_app_context
 from nucliadb.writer import API_PREFIX
 from nucliadb.writer.api.v1.router import api as api_v1
 from nucliadb.writer.lifecycle import lifespan
@@ -96,17 +94,4 @@ async def homepage(request):
     # Use raw starlette routes to avoid unnecessary overhead
     application.add_route("/", homepage)
 
-    maybe_configure_back_pressure(application)
     return application
-
-
-def maybe_configure_back_pressure(application: FastAPI):
-    from nucliadb.writer.back_pressure import start_materializer, stop_materializer
-    from nucliadb.writer.settings import back_pressure_settings
-    from nucliadb_utils.settings import is_onprem_nucliadb
-
-    if back_pressure_settings.enabled and not is_onprem_nucliadb():
-        context = get_app_context(application)
-        start_materializer_with_context = functools.partial(start_materializer, context)
-        application.add_event_handler("startup", start_materializer_with_context)
-        application.add_event_handler("shutdown", stop_materializer)

nucliadb/src/nucliadb/writer/lifecycle.py

@@ -25,10 +25,13 @@
 from nucliadb.ingest.processing import start_processing_engine
 from nucliadb.ingest.utils import start_ingest, stop_ingest
 from nucliadb.writer import SERVICE_NAME
+from nucliadb.writer.back_pressure import start_materializer, stop_materializer
+from nucliadb.writer.settings import back_pressure_settings
 from nucliadb.writer.tus import finalize as storage_finalize
 from nucliadb.writer.tus import initialize as storage_initialize
 from nucliadb.writer.utilities import get_processing
 from nucliadb_telemetry.utils import clean_telemetry, setup_telemetry
+from nucliadb_utils.settings import is_onprem_nucliadb
 from nucliadb_utils.utilities import (
     finalize_utilities,
     start_partitioning_utility,
@@ -39,6 +42,8 @@
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
+    back_pressure_enabled = back_pressure_settings.enabled and not is_onprem_nucliadb()
+
     await setup_telemetry(SERVICE_NAME)
     await start_ingest(SERVICE_NAME)
     await start_processing_engine()
@@ -47,9 +52,13 @@ async def lifespan(app: FastAPI):
     await storage_initialize()
 
     # Inject application context into the fastapi app's state
-    async with inject_app_context(app):
+    async with inject_app_context(app) as context:
+        if back_pressure_enabled:
+            await start_materializer(context)
         yield
 
+    if back_pressure_enabled:
+        await stop_materializer()
     await stop_transaction_utility()
     await stop_ingest()
     processing = get_processing()