fix: avoid duplicate startup embedding jobs

Author: tarjan1

backend/common/utils/distributed_lock.py

@@ -0,0 +1,133 @@
+"""Distributed-lock abstraction for process-wide singleton tasks.
+
+The current provider uses PostgreSQL advisory locks. Business callers use
+``SingleWorkerGuard``, which delegates lock operations to ``DistributedLock``.
+"""
+
+import hashlib
+from collections.abc import Callable
+from enum import Enum, auto
+from functools import wraps
+from threading import Lock
+from typing import ClassVar, Self
+
+from sqlalchemy import func, select
+from sqlalchemy.engine import Connection
+from sqlalchemy.exc import SQLAlchemyError
+
+from common.core.db import engine
+from common.utils.utils import SQLBotLogUtil
+
+
+class LockStatus(Enum):
+    """Process-local state of the single-worker lock acquisition attempt."""
+
+    NOT_ATTEMPTED = auto()
+    ACQUIRED = auto()
+    NOT_ACQUIRED = auto()
+
+
+class DistributedLock:
+    """A lock held by this process until ``release`` or process shutdown."""
+
+    def __init__(
+        self,
+        advisory_key: int,
+        connection: Connection,
+    ) -> None:
+        self._connection = connection
+        self._advisory_key = advisory_key
+
+    @staticmethod
+    def _postgres_advisory_key(key: str) -> int:
+        """Convert a readable lock key into PostgreSQL's signed 64-bit key."""
+        digest = hashlib.sha256(key.encode("utf-8")).digest()
+        return int.from_bytes(digest[:8], byteorder="big", signed=True)
+
+    @classmethod
+    def try_acquire(cls, key: str) -> Self | None:
+        """Try to acquire a named distributed lock without waiting."""
+        advisory_key = cls._postgres_advisory_key(key)
+        connection: Connection | None = None
+        try:
+            connection = engine.connect()
+            acquired = connection.execute(
+                select(func.pg_try_advisory_lock(advisory_key))
+            ).scalar_one()
+            # End the implicit transaction; the session-level lock remains held.
+            connection.commit()
+        except SQLAlchemyError:
+            SQLBotLogUtil.exception("Failed to acquire distributed lock: %s", key)
+            if connection is not None:
+                connection.close()
+            return None
+
+        if not acquired:
+            connection.close()
+            return None
+
+        return cls(advisory_key, connection)
+
+    def release(self) -> None:
+        """Release the lock and return its dedicated connection to the pool."""
+        if self._connection.closed:
+            return
+
+        try:
+            self._connection.execute(
+                select(func.pg_advisory_unlock(self._advisory_key))
+            )
+            self._connection.commit()
+        finally:
+            self._connection.close()
+
+
+class SingleWorkerGuard:
+    """Run startup tasks in only one worker process."""
+
+    _LOCK_KEY = "sqlbot:startup:embedding"
+    _lock_status: ClassVar[LockStatus] = LockStatus.NOT_ATTEMPTED
+    # Keep the acquired lock and its dedicated connection alive.
+    _acquired_lock: ClassVar[DistributedLock | None] = None
+    _state_mutex = Lock()
+
+    @classmethod
+    def _has_acquired_lock(cls) -> bool:
+        with cls._state_mutex:
+            if cls._lock_status is LockStatus.NOT_ATTEMPTED:
+                lock = DistributedLock.try_acquire(cls._LOCK_KEY)
+                if lock is None:
+                    cls._lock_status = LockStatus.NOT_ACQUIRED
+                    SQLBotLogUtil.info(
+                        "Current process FAILED to acquire the single-worker lock"
+                    )
+                else:
+                    cls._acquired_lock = lock
+                    cls._lock_status = LockStatus.ACQUIRED
+                    SQLBotLogUtil.info(
+                        "Current process acquired the single-worker lock"
+                    )
+            return cls._lock_status is LockStatus.ACQUIRED
+
+    @classmethod
+    def once(cls, func: Callable[[], None]) -> Callable[[], None]:
+        """Run the decorated startup task in only one worker."""
+
+        @wraps(func)
+        def wrapper() -> None:
+            if not cls._has_acquired_lock():
+                return
+            func()
+
+        return wrapper
+
+    @classmethod
+    def release(cls) -> None:
+        """Release the single-worker lock and reset this process's state."""
+        with cls._state_mutex:
+            try:
+                if cls._acquired_lock is not None:
+                    cls._acquired_lock.release()
+            finally:
+                cls._acquired_lock = None
+                cls._lock_status = LockStatus.NOT_ATTEMPTED

backend/main.py

@@ -27,8 +27,12 @@
 from common.core.config import settings
 from common.core.response_middleware import ResponseMiddleware, exception_handler
 from common.core.sqlbot_cache import init_sqlbot_cache
-from common.utils.embedding_threads import fill_empty_terminology_embeddings, fill_empty_data_training_embeddings, \
-    fill_empty_table_and_ds_embeddings
+from common.utils.distributed_lock import SingleWorkerGuard
+from common.utils.embedding_threads import (
+    fill_empty_data_training_embeddings,
+    fill_empty_table_and_ds_embeddings,
+    fill_empty_terminology_embeddings,
+)
 from common.utils.utils import SQLBotLogUtil
 
 
@@ -37,18 +41,28 @@ def run_migrations():
     command.upgrade(alembic_cfg, "head")
 
 
+@SingleWorkerGuard.once
 def init_terminology_embedding_data():
     fill_empty_terminology_embeddings()
 
 
+@SingleWorkerGuard.once
 def init_data_training_embedding_data():
     fill_empty_data_training_embeddings()
 
 
+@SingleWorkerGuard.once
 def init_table_and_ds_embedding():
     fill_empty_table_and_ds_embeddings()
 
 
+def shutdown_resources() -> None:
+    try:
+        SingleWorkerGuard.release()
+    except Exception:
+        SQLBotLogUtil.exception("SQLBot shutdown failed")
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     run_migrations()
@@ -61,8 +75,11 @@ async def lifespan(app: FastAPI):
     await sqlbot_xpack.core.clean_xpack_cache()
     await async_model_info()  # 异步加密已有模型的密钥和地址
     await sqlbot_xpack.core.monitor_app(app)
-    yield
-    SQLBotLogUtil.info("SQLBot 应用关闭")
+    try:
+        yield
+    finally:
+        shutdown_resources()
+        SQLBotLogUtil.info("SQLBot 应用关闭")
 
 
 def custom_generate_unique_id(route: APIRoute) -> str: