Acquire a lock before running initialization steps (#768)

### Description

If multiple workers are detected, and a Redis client, Hyperion will
raise an error.
Otherwise the server won't be able to:
 - run initialization steps
 - send messages over WebSocket reliably.

The logic is extracted from the great #375

---------

Co-authored-by: Thonyk <[email protected]>
Co-authored-by: Timothée Robert <[email protected]>

Author: armanddidierjean

Dockerfile

@@ -16,7 +16,4 @@ COPY migrations migrations/
 COPY assets assets/
 COPY app app/
 
-COPY start.sh .
-RUN chmod +x start.sh
-
-ENTRYPOINT ["./start.sh"]
+ENTRYPOINT ["fastapi", "run", "app/main.py", "--workers", "${NB_WORKERS:-1}"]

app/api.py

@@ -2,14 +2,10 @@
 
 from fastapi import APIRouter
 
-from app.core.core_module_list import core_module_list
-from app.modules.module_list import module_list
+from app.module import all_modules
 
 api_router = APIRouter()
 
 
-for core_module in core_module_list:
+for core_module in all_modules:
     api_router.include_router(core_module.router)
-
-for module in module_list:
-    api_router.include_router(module.router)

app/app.py

@@ -17,8 +17,10 @@
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from fastapi.routing import APIRoute
+from redis import Redis
 from sqlalchemy.engine import Connection, Engine
 from sqlalchemy.exc import IntegrityError
+from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import Session
 
 from app import api
@@ -36,18 +38,23 @@
     get_scheduler,
     get_websocket_connection_manager,
     init_and_get_db_engine,
+    init_websocket_connection_manager,
+)
+from app.module import module_list
+from app.types.exceptions import (
+    ContentHTTPException,
+    GoogleAPIInvalidCredentialsError,
+    MultipleWorkersWithoutRedisInitializationError,
 )
-from app.modules.module_list import module_list
-from app.types.exceptions import ContentHTTPException, GoogleAPIInvalidCredentialsError
 from app.types.sqlalchemy import Base
 from app.utils import initialization
 from app.utils.redis import limiter
 
 if TYPE_CHECKING:
     import redis
 
-    from app.types.scheduler import Scheduler
-    from app.types.websocket import WebsocketConnectionManager
+from app.types.scheduler import Scheduler
+from app.types.websocket import WebsocketConnectionManager
 
 # NOTE: We can not get loggers at the top of this file like we do in other files
 # as the loggers are not yet initialized
@@ -330,6 +337,10 @@ def init_db(
         sync_engine=sync_engine,
         hyperion_error_logger=hyperion_error_logger,
     )
+
+    # TODO: we may allow the following steps to be run by other workers
+    # and may not need to wait for them
+    # These two steps could use an async database connection
     initialize_schools(
         sync_engine=sync_engine,
         hyperion_error_logger=hyperion_error_logger,
@@ -340,15 +351,30 @@ def init_db(
     )
 
 
-# We wrap the application in a function to be able to pass the settings and drop_db parameters
-# The drop_db parameter is used to drop the database tables before creating them again
-def get_application(settings: Settings, drop_db: bool = False) -> FastAPI:
-    # Initialize loggers
-    LogConfig().initialize_loggers(settings=settings)
+async def init_google_API(
+    db: AsyncSession,
+    settings: Settings,
+) -> None:
+    # Init Google API credentials
 
-    hyperion_access_logger = logging.getLogger("hyperion.access")
-    hyperion_security_logger = logging.getLogger("hyperion.security")
-    hyperion_error_logger = logging.getLogger("hyperion.error")
+    google_api = GoogleAPI()
+
+    if google_api.is_google_api_configured(settings):
+        try:
+            await google_api.get_credentials(db, settings)
+
+        except GoogleAPIInvalidCredentialsError:
+            # We expect this error to be raised if the credentials were never set before
+            pass
+
+
+def test_configuration(
+    settings: Settings,
+    hyperion_error_logger: logging.Logger,
+) -> None:
+    """
+    Test configuration and log warnings if some settings are not configured correctly.
+    """
 
     # We use warning level so that the message is not sent to matrix again
     if not settings.MATRIX_TOKEN:
@@ -369,42 +395,108 @@ def get_application(settings: Settings, drop_db: bool = False) -> FastAPI:
     Path("data/ics/").mkdir(parents=True, exist_ok=True)
     Path("data/core/").mkdir(parents=True, exist_ok=True)
 
+
+async def init_lifespan(
+    app: FastAPI,
+    settings: Settings,
+    hyperion_error_logger: logging.Logger,
+    drop_db: bool,
+) -> tuple[Scheduler, WebsocketConnectionManager]:
+    hyperion_error_logger.info("Startup: Initializing application")
+
+    # Init the Redis client
+    redis_client: redis.Redis | bool | None = app.dependency_overrides.get(
+        get_redis_client,
+        get_redis_client,
+    )(settings=settings)
+
+    # Initialization steps should only be run once across all workers
+    # We use Redis locks to ensure that the initialization steps are only run once
+    if initialization.get_number_of_workers() > 1 and not isinstance(
+        redis_client,
+        Redis,
+    ):
+        raise MultipleWorkersWithoutRedisInitializationError
+
+    # We need to run the database initialization only once across all the workers
+    # Other workers have to wait for the db to be initialized
+    await initialization.use_lock_for_workers(
+        init_db,
+        "init_db",
+        redis_client,
+        hyperion_error_logger,
+        unlock_key="db_initialized",
+        settings=settings,
+        hyperion_error_logger=hyperion_error_logger,
+        drop_db=drop_db,
+    )
+
+    await initialization.use_lock_for_workers(
+        test_configuration,
+        "test_configuration",
+        redis_client,
+        hyperion_error_logger,
+        settings=settings,
+        hyperion_error_logger=hyperion_error_logger,
+    )
+
+    async for db in app.dependency_overrides.get(
+        get_db,
+        get_db,
+    )():
+        await initialization.use_lock_for_workers(
+            init_google_API,
+            "init_google_API",
+            redis_client,
+            hyperion_error_logger,
+            db=db,
+            settings=settings,
+        )
+
+    init_websocket_connection_manager(WebsocketConnectionManager(settings=settings))
+    ws_manager: WebsocketConnectionManager = app.dependency_overrides.get(
+        get_websocket_connection_manager,
+        get_websocket_connection_manager,
+    )(settings=settings)
+
+    arq_scheduler: Scheduler = app.dependency_overrides.get(
+        get_scheduler,
+        get_scheduler,
+    )(settings=settings)
+
+    await ws_manager.connect_broadcaster()
+    await arq_scheduler.start(
+        redis_host=settings.REDIS_HOST,
+        redis_port=settings.REDIS_PORT,
+        redis_password=settings.REDIS_PASSWORD,
+        _dependency_overrides=app.dependency_overrides,
+    )
+    return arq_scheduler, ws_manager
+
+
+# We wrap the application in a function to be able to pass the settings and drop_db parameters
+# The drop_db parameter is used to drop the database tables before creating them again
+def get_application(settings: Settings, drop_db: bool = False) -> FastAPI:
+    # Initialize loggers
+    LogConfig().initialize_loggers(settings=settings)
+
+    hyperion_access_logger = logging.getLogger("hyperion.access")
+    hyperion_security_logger = logging.getLogger("hyperion.security")
+    hyperion_error_logger = logging.getLogger("hyperion.error")
+
     # Creating a lifespan which will be called when the application starts then shuts down
     # https://fastapi.tiangolo.com/advanced/events/
     @asynccontextmanager
     async def lifespan(app: FastAPI) -> AsyncGenerator:
-        # Init Google API credentials
-        google_api = GoogleAPI()
-        if google_api.is_google_api_configured(settings):
-            async for db in app.dependency_overrides.get(
-                get_db,
-                get_db,
-            )():
-                try:
-                    await google_api.get_credentials(db, settings)
-                except GoogleAPIInvalidCredentialsError:
-                    # We expect this error to be raised if the credentials were never set before
-                    pass
-
-        ws_manager: WebsocketConnectionManager = app.dependency_overrides.get(
-            get_websocket_connection_manager,
-            get_websocket_connection_manager,
-        )(settings=settings)
-
-        arq_scheduler: Scheduler = app.dependency_overrides.get(
-            get_scheduler,
-            get_scheduler,
-        )(settings=settings)
-
-        await ws_manager.connect_broadcaster()
-        await arq_scheduler.start(
-            redis_host=settings.REDIS_HOST,
-            redis_port=settings.REDIS_PORT,
-            redis_password=settings.REDIS_PASSWORD,
-            _dependency_overrides=app.dependency_overrides,
+        arq_scheduler, ws_manager = await init_lifespan(
+            app=app,
+            settings=settings,
+            hyperion_error_logger=hyperion_error_logger,
+            drop_db=drop_db,
         )
 
         yield
+
         hyperion_error_logger.info("Shutting down")
         await arq_scheduler.close()
         await ws_manager.disconnect_broadcaster()
@@ -429,21 +521,6 @@ async def lifespan(app: FastAPI) -> AsyncGenerator:
     calypsso = get_calypsso_app()
     app.mount("/calypsso", calypsso, "Calypsso")
 
-    if settings.HYPERION_INIT_DB:
-        init_db(
-            settings=settings,
-            hyperion_error_logger=hyperion_error_logger,
-            drop_db=drop_db,
-        )
-    else:
-        hyperion_error_logger.info("Database initialization skipped")
-
-    # Initialize Redis
-    if not app.dependency_overrides.get(get_redis_client, get_redis_client)(
-        settings=settings,
-    ):
-        hyperion_error_logger.info("Redis client not configured")
-
     # We need to init the database engine to be able to use it in dependencies
     init_and_get_db_engine(settings)
 

app/core/core_endpoints/endpoints_core.py

@@ -16,7 +16,7 @@
     is_user,
     is_user_in,
 )
-from app.modules.module_list import module_list
+from app.module import module_list
 from app.types.module import CoreModule
 from app.utils.tools import is_group_id_valid
 

app/core/core_module_list.py

@@ -9,13 +9,12 @@
 from pydantic import TypeAdapter, ValidationError
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from app.core.core_module_list import core_module_list
 from app.core.payment import cruds_payment, models_payment, schemas_payment
 from app.core.payment.types_payment import (
     NotificationResultContent,
 )
 from app.dependencies import get_db
-from app.modules.module_list import module_list
+from app.module import all_modules
 from app.types.module import CoreModule
 
 router = APIRouter(tags=["Payments"])
@@ -133,7 +132,7 @@ async def webhook(
 
         # If a callback is defined for the module, we want to call it
         try:
-            for module in module_list + core_module_list:
+            for module in all_modules:
                 if module.root == checkout.module:
                     if module.payment_callback is not None:
                         hyperion_error_logger.info(

app/core/payment/endpoints_payment.py

@@ -309,11 +309,6 @@ def MINIMAL_TITAN_VERSION_CODE(cls) -> str:
     # Automatically generated parameters #
     ######################################
 
-    # If Hyperion should initialize the database on startup
-    # This environment variable is set by our init Python file to tell the workers to avoid initializing the database
-    # You don't want to set this variable manually
-    HYPERION_INIT_DB: bool = True
-
     # The following properties can not be instantiated as class variables as them need to be computed using another property from the class,
     # which won't be available before the .env file parsing.
     # We thus decide to use the decorator `@property` to make these methods usable as properties and not functions: as properties: Settings.RSA_PRIVATE_KEY, Settings.RSA_PUBLIC_KEY and Settings.RSA_PUBLIC_JWK

app/core/utils/config.py

@@ -205,14 +205,16 @@ def get_scheduler(settings: Settings = Depends(get_settings)) -> Scheduler:
     return scheduler
 
 
-def get_websocket_connection_manager(
-    settings: Settings = Depends(get_settings),
-):
+def init_websocket_connection_manager(
+    wscm: WebsocketConnectionManager,
+) -> None:
     global websocket_connection_manager
+    websocket_connection_manager = wscm
 
-    if websocket_connection_manager is None:
-        websocket_connection_manager = WebsocketConnectionManager(settings=settings)
 
+def get_websocket_connection_manager(
+    settings: Settings = Depends(get_settings),
+):
     return websocket_connection_manager
 
 

app/dependencies.py

@@ -2,11 +2,13 @@
 import logging
 from pathlib import Path
 
-from app.types.module import Module
+from app.types.module import CoreModule, Module
 
 hyperion_error_logger = logging.getLogger("hyperion.error")
 
 module_list: list[Module] = []
+core_module_list: list[CoreModule] = []
+all_modules: list[CoreModule] = []
 
 for endpoints_file in Path().glob("app/modules/*/endpoints_*.py"):
     endpoint_module = importlib.import_module(
@@ -19,3 +21,19 @@
         hyperion_error_logger.error(
             f"Module {endpoints_file} does not declare a module. It won't be enabled.",
         )
+
+for endpoints_file in Path().glob("app/core/*/endpoints_*.py"):
+    endpoint_module = importlib.import_module(
+        ".".join(endpoints_file.with_suffix("").parts),
+    )
+
+    if hasattr(endpoint_module, "core_module"):
+        core_module: CoreModule = endpoint_module.core_module
+        core_module_list.append(core_module)
+    else:
+        hyperion_error_logger.error(
+            f"Core module {endpoints_file} does not declare a core module. It won't be enabled.",
+        )
+
+
+all_modules = module_list + core_module_list