[core][feat] Use database locks to perform migration (#2184)

Author: aquamatthias

fixcore/fixcore/__main__.py

@@ -37,8 +37,7 @@
     parse_config,
     CoreConfig,
 )
-from fixcore.db import SystemData
-from fixcore.db.db_access import DbAccess, CurrentDatabaseVersion
+from fixcore.db.db_access import DbAccess
 from fixcore.db.system_data_db import EphemeralJwtSigningKey
 from fixcore.dependencies import Dependencies, ServiceNames, TenantDependencies
 from fixcore.dependencies import (
@@ -129,12 +128,7 @@ def run_process(args: Namespace) -> None:
             cert_handler_no_ca = deps.add(ServiceNames.cert_handler, CertificateHandlerNoCA.lookup(config, temp))
             verify: Union[bool, str] = False if args.graphdb_no_ssl_verify else str(cert_handler_no_ca.ca_bundle)
             deps.add(ServiceNames.config, evolve(config, run=RunConfig(temp, verify)))
-            deps.add(ServiceNames.system_data, SystemData("multi-tenant", utc(), CurrentDatabaseVersion))
-            deps.add(
-                ServiceNames.event_sender,
-                PostHogEventSender(deps.system_data) if config.runtime.usage_metrics else NoEventSender(),
-            )
-
+            deps.add(ServiceNames.event_sender, NoEventSender())
             provider: TenantDependencyProvider = deps.add(
                 ServiceNames.tenant_dependency_provider, FromRequestTenantDependencyProvider(deps)
             )
@@ -172,6 +166,9 @@ async def direct_tenant(deps: TenantDependencies) -> None:
     config = deps.config
     event_sender = deps.event_sender
     db = deps.service(ServiceNames.db_access, DbAccess)
+    # migrate the database to latest schema
+    db_change = await db.migrate()
+    deps.add(ServiceNames.system_data, db_change.current)
     message_bus = deps.add(ServiceNames.message_bus, MessageBus())
     scheduler = deps.add(ServiceNames.scheduler, APScheduler() if not config.args.no_scheduling else NoScheduler())
     model = deps.add(ServiceNames.model_handler, ModelHandlerDB(db, config.runtime.plantuml_server))
@@ -202,7 +199,9 @@ async def direct_tenant(deps: TenantDependencies) -> None:
     subscriptions = deps.add(ServiceNames.subscription_handler, SubscriptionHandlerService(message_bus))
     core_config_handler = deps.add(
         ServiceNames.core_config_handler,
-        CoreConfigHandler(config, message_bus, worker_task_queue, config_handler, event_sender, inspector),
+        CoreConfigHandler(
+            config, message_bus, worker_task_queue, config_handler, event_sender, inspector, db_change, db.lock_db
+        ),
     )
     deps.add(ServiceNames.infra_apps_runtime, LocalfixcoreAppRuntime(cli))
     deps.add(

fixcore/fixcore/config/core_config_handler.py

@@ -4,9 +4,9 @@
 from contextlib import suppress
 from functools import partial
 from typing import Optional, List, Callable, Awaitable
-from apscheduler.triggers.cron import CronTrigger
 
 import yaml
+from apscheduler.triggers.cron import CronTrigger
 
 from fixcore.analytics import AnalyticsEventSender, CoreEvent
 from fixcore.config import ConfigHandler, ConfigEntity, ConfigValidation
@@ -25,14 +25,16 @@
     config_model as core_config_model,
     migrate_command_config,
 )
-from fixcore.system_start import empty_config
+from fixcore.db import DatabaseChange
+from fixcore.db.lockdb import LockDB
 from fixcore.ids import SubscriberId, WorkerId, ConfigId
 from fixcore.message_bus import MessageBus, CoreMessage
 from fixcore.model.model import Kind
 from fixcore.model.typed_model import from_js
 from fixcore.report import FixReportBenchmark, FixReportCheck, Inspector, BenchmarkConfigRoot, CheckConfigRoot
 from fixcore.report.report_config import config_model as report_config_model
 from fixcore.service import Service
+from fixcore.system_start import empty_config
 from fixcore.types import Json
 from fixcore.user import config_model as user_config_model, UsersConfigId, FixInventoryUsersConfig
 from fixcore.util import deep_merge, restart_service, value_in_path, value_in_path_get
@@ -50,6 +52,8 @@ def __init__(
         config_handler: ConfigHandler,
         event_sender: AnalyticsEventSender,
         inspector: Inspector,
+        db_change: DatabaseChange,
+        lock_db: LockDB,
         exit_fn: Callable[[], None] = partial(restart_service, "fixcore config changed."),
     ):
         super().__init__()
@@ -61,6 +65,8 @@ def __init__(
         self.config_handler = config_handler
         self.event_sender = event_sender
         self.inspector = inspector
+        self.db_change = db_change
+        self.lock_db = lock_db
         self.exit_fn = exit_fn
         self.config_updated_callbacks: List[Callable[[ConfigId], Awaitable[None]]] = []
 
@@ -250,9 +256,10 @@ async def __update_model(self) -> None:
             log.error(f"Could not update fix core config model: {ex}", exc_info=ex)
 
     async def start(self) -> None:
-        if not self.config.multi_tenant_setup:
-            await self.__update_model()
-            await self.__update_config()
+        if self.db_change.has_changed() and not self.config.multi_tenant_setup:
+            async with self.lock_db.lock("fixcore.config.update"):
+                await self.__update_model()
+                await self.__update_config()
         self.config_updated_listener = asyncio.create_task(self.__handle_events())
         self.config_validator = asyncio.create_task(self.__validate_config())
 

fixcore/fixcore/db/__init__.py

@@ -4,8 +4,12 @@
 from datetime import datetime
 from enum import Enum
 
+from fixcore.core_config import current_git_hash
 from fixcore.types import Json
 
+MigrateAlways = False
+CurrentDatabaseVersion = 2
+
 
 @define
 class SystemData:
@@ -14,6 +18,29 @@ class SystemData:
     db_version: int
     version: Optional[str] = None
 
+    def detect_change(self) -> bool:
+        git_hash = current_git_hash()
+        return (
+            MigrateAlways
+            or self.db_version != CurrentDatabaseVersion
+            or self.version is None
+            or git_hash is None
+            or git_hash != self.version
+        )
+
+
+@define
+class DatabaseChange:
+    previous: Optional[SystemData]
+    current: SystemData
+
+    def has_changed(self) -> bool:
+        return (
+            self.previous is None
+            or self.previous.version != self.current.version
+            or self.previous.db_version != self.current.db_version
+        )
+
 
 class EstimatedQueryCostRating(Enum):
     simple = 1

fixcore/fixcore/db/db_access.py

@@ -10,21 +10,22 @@
 from arango import ArangoServerError
 from arango.client import ArangoClient
 from arango.database import StandardDatabase
-from attr import frozen
+from attr import frozen, evolve
 from dateutil.parser import parse
 from requests.exceptions import RequestException
 
 from fixcore.analytics import AnalyticsEventSender
 from fixcore.async_extensions import run_async
 from fixcore.core_config import CoreConfig, current_git_hash
-from fixcore.db import SystemData
+from fixcore.db import SystemData, DatabaseChange
 from fixcore.db.arangodb_extensions import ArangoHTTPClient
 from fixcore.db.async_arangodb import AsyncArangoDB, AsyncCursor
 from fixcore.db.configdb import config_entity_db, config_validation_entity_db
 from fixcore.db.deferredouteredgedb import deferred_outer_edge_db
 from fixcore.db.entitydb import EventEntityDb
 from fixcore.db.graphdb import ArangoGraphDB, GraphDB, EventGraphDB
 from fixcore.db.jobdb import job_db
+from fixcore.db.lockdb import LockDB
 from fixcore.db.modeldb import ModelDb, model_db
 from fixcore.db.packagedb import app_package_entity_db
 from fixcore.db.reportdb import report_check_db, benchmark_db
@@ -76,6 +77,7 @@ def __init__(
         time_series: str = "ts",
         report_checks: str = "report_checks",
         benchmarks: str = "report_benchmarks",
+        locks: str = "locks",
     ):
         super().__init__()
         self.event_sender = event_sender
@@ -95,34 +97,28 @@ def __init__(
         self.report_check_db = report_check_db(self.db, report_checks)
         self.benchmark_db = benchmark_db(self.db, benchmarks)
         self.time_series_db = TimeSeriesDB(self.db, time_series, config)
+        self.lock_db = LockDB(self.db, locks)
         self.graph_dbs: Dict[str, GraphDB] = {}
         self.config = config
         self.cleaner = Periodic("outdated_updates_cleaner", self.check_outdated_updates, timedelta(seconds=60))
 
     async def start(self) -> None:
-        await self.__migrate()
         await self.cleaner.start()
 
     async def stop(self) -> None:
         await self.cleaner.stop()
 
-    async def __migrate(self) -> None:
-        try:
-            system_data = await self.system_data_db.system_data()
-        except Exception:
-            system_data = None
-            if not await self.db.has_collection("system_data"):  # make sure the system data collection exists
-                await self.db.create_collection("system_data")
-        if system_data is None:  # in case no version is available, create a genesis version
-            system_data = SystemData(uuid_str(), utc(), CurrentDatabaseVersion)
-        git_hash = current_git_hash()
-        if (
-            MigrateAlways
-            or system_data.db_version != CurrentDatabaseVersion
-            or system_data.version is None
-            or git_hash is None
-            or git_hash != system_data.version
-        ):
+    async def migrate(self) -> DatabaseChange:
+        async def sys_data() -> Optional[SystemData]:
+            try:
+                return await self.system_data_db.system_data()
+            except Exception:
+                if not await self.db.has_collection("system_data"):  # make sure the system data collection exists
+                    await self.db.create_collection("system_data")
+                return None
+
+        async def do_migrate(system_data: SystemData) -> SystemData:
+            git_hash = current_git_hash()
             # check if we need to run a migration
             if system_data.db_version < CurrentDatabaseVersion:
                 log.info(f"Database migration required: db={system_data.db_version} -> latest={CurrentDatabaseVersion}")
@@ -131,7 +127,7 @@ async def __migrate(self) -> None:
                 for version in range(system_data.db_version, CurrentDatabaseVersion):
                     log.info(f"Running migration {version} -> {version + 1}")
                     await migrations[version]()
-                system_data.db_version = CurrentDatabaseVersion
+                system_data = evolve(system_data, db_version=CurrentDatabaseVersion)
                 await self.system_data_db.update_system_data(system_data)
 
             # will be executed on every git change
@@ -159,10 +155,22 @@ async def __migrate(self) -> None:
                 await em.create_update_schema()
 
             # update the system data version to the current git hash
-            system_data.version = git_hash
-
+            system_data = evolve(system_data, version=git_hash)
             # update the system data version to not migrate the next time
             await self.system_data_db.update_system_data(system_data)
+            return system_data
+
+        existing = await sys_data()
+        if existing is None or existing.detect_change():
+            # create `lock` collection to be able to create database locks
+            await self.lock_db.create_update_schema()
+            async with self.lock_db.lock("db_access.migrate"):
+                # lookup again with an exclusive database lock
+                existing = await sys_data()
+                if existing is None or existing.detect_change():
+                    updated = await do_migrate(existing or SystemData(uuid_str(), utc(), CurrentDatabaseVersion))
+                    return DatabaseChange(existing, updated)
+        return DatabaseChange(existing, existing)
 
     async def __migrate_v1_to_v2(self) -> None:
         def migrate_config(old_id: str, old_root: str, config: Json) -> Json:

fixcore/fixcore/db/lockdb.py

@@ -0,0 +1,68 @@
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+from datetime import timedelta
+from typing import cast, List, AsyncIterator
+
+from arango import CollectionCreateError
+
+from fixcore.db.async_arangodb import AsyncArangoDB
+from fixcore.types import Json
+from fixcore.util import utc
+
+log = logging.getLogger(__name__)
+
+
+class LockAcquisitionError(Exception):
+    pass
+
+
+class LockDB:
+    def __init__(self, db: AsyncArangoDB, collection_name: str) -> None:
+        self.db = db
+        self.collection_name = collection_name
+
+    @asynccontextmanager
+    async def lock(
+        self,
+        name: str,
+        *,
+        get_lock: timedelta = timedelta(seconds=60),  # how long to try to get the lock
+        lock_for: timedelta = timedelta(seconds=60),  # acquired: how long to hold the lock max
+        retry_interval: timedelta = timedelta(seconds=0.1),  # how long to wait between retries
+    ) -> AsyncIterator[None]:
+        now = utc()
+        ttl = int((now + lock_for).timestamp())
+        deadline = now + get_lock
+        lock_acquired = False
+        while not lock_acquired and utc() < deadline:
+            try:
+                # try to insert a document with key __lock__
+                await self.db.insert(self.collection_name, dict(_key=name, expires=ttl), sync=True)
+            except Exception:
+                # could not insert the document. Wait and try again
+                await asyncio.sleep(retry_interval.total_seconds())
+                continue
+            lock_acquired = True
+            try:
+                yield
+            finally:
+                await self.db.delete(self.collection_name, name, ignore_missing=True)
+        if not lock_acquired:
+            raise LockAcquisitionError(f"Could not acquire lock {name}")
+
+    async def create_update_schema(self) -> None:
+        if not await self.db.has_collection(self.collection_name):
+            try:
+                await self.db.create_collection(self.collection_name)
+            except CollectionCreateError as ex:
+                if ex.error_code != 1207:  # already exists
+                    raise
+        collection = self.db.collection(self.collection_name)
+        indexes = {idx["name"]: idx for idx in cast(List[Json], collection.indexes())}
+        if "ttl" not in indexes:
+            try:
+                # The ttl expiry is only a safeguard - if the process is not able to clean up requested locks
+                collection.add_ttl_index(["expires"], expiry_time=0, name="ttl")
+            except Exception as ex:
+                log.info(f"Could not create TTL index for lock collection: {ex}")

fixcore/fixcore/dependencies.py

@@ -446,6 +446,8 @@ def standard_database() -> StandardDatabase:
         # direct db access
         sdb = deps.add(ServiceNames.system_database, await run_async(standard_database))
         db = deps.add(ServiceNames.db_access, DbAccess(sdb, dp.event_sender, NoAdjust(), config))
+        db_change = await db.migrate()
+        deps.add(ServiceNames.system_data, db_change.current)
         # no scheduler required in multi-tenant mode
         scheduler = deps.add(ServiceNames.scheduler, NoScheduler())
         # all tenants use the same model (derived from code)
@@ -473,7 +475,9 @@ def standard_database() -> StandardDatabase:
         subscriptions = deps.add(ServiceNames.subscription_handler, NoSubscriptionHandler())
         core_config_handler = deps.add(
             ServiceNames.core_config_handler,
-            CoreConfigHandler(config, message_bus, worker_task_queue, config_handler, event_sender, inspector),
+            CoreConfigHandler(
+                config, message_bus, worker_task_queue, config_handler, event_sender, inspector, db_change, db.lock_db
+            ),
         )
         # Enable package manager and runtime for infra apps when required
         # deps.add(ServiceNames.infra_apps_runtime, LocalfixcoreAppRuntime(cli))