fix: prevent duplicate startup embedding jobs across workers

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:

tests/test_distributed_lock.py

@@ -0,0 +1,170 @@
+"""Focused tests for the startup embedding distributed lock."""
+
+import sys
+import unittest
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from sqlalchemy.engine import Connection
+from sqlalchemy.exc import SQLAlchemyError
+
+
+BACKEND_DIR = Path(__file__).resolve().parents[1] / "backend"
+sys.path.insert(0, str(BACKEND_DIR))
+
+from common.utils import distributed_lock as lock_module  # noqa: E402
+from common.utils.distributed_lock import (  # noqa: E402
+    DistributedLock,
+    LockStatus,
+    SingleWorkerGuard,
+)
+
+
+class DistributedLockTestCase(unittest.TestCase):
+    """Verify PostgreSQL advisory-lock acquisition and cleanup behavior."""
+
+    @staticmethod
+    def _connection(acquired: bool = True) -> MagicMock:
+        connection = MagicMock(spec=Connection)
+        connection.closed = False
+        connection.execute.return_value.scalar_one.return_value = acquired
+        return connection
+
+    def test_advisory_key_is_stable_across_processes(self) -> None:
+        key = DistributedLock._postgres_advisory_key("sqlbot:startup:embedding")
+
+        self.assertEqual(6735965423478359195, key)
+        self.assertGreaterEqual(key, -(2**63))
+        self.assertLessEqual(key, 2**63 - 1)
+
+    def test_try_acquire_keeps_successful_connection_open(self) -> None:
+        connection = self._connection(acquired=True)
+
+        with patch.object(lock_module, "engine") as engine:
+            engine.connect.return_value = connection
+
+            acquired_lock = DistributedLock.try_acquire("startup-task")
+
+        self.assertIsInstance(acquired_lock, DistributedLock)
+        connection.commit.assert_called_once_with()
+        connection.close.assert_not_called()
+
+    def test_try_acquire_closes_connection_when_lock_is_busy(self) -> None:
+        connection = self._connection(acquired=False)
+
+        with patch.object(lock_module, "engine") as engine:
+            engine.connect.return_value = connection
+
+            acquired_lock = DistributedLock.try_acquire("startup-task")
+
+        self.assertIsNone(acquired_lock)
+        connection.commit.assert_called_once_with()
+        connection.close.assert_called_once_with()
+
+    def test_try_acquire_closes_connection_on_sqlalchemy_error(self) -> None:
+        connection = self._connection()
+        connection.execute.side_effect = SQLAlchemyError("database error")
+
+        with (
+            patch.object(lock_module, "engine") as engine,
+            patch.object(lock_module.SQLBotLogUtil, "exception") as log_error,
+        ):
+            engine.connect.return_value = connection
+
+            acquired_lock = DistributedLock.try_acquire("startup-task")
+
+        self.assertIsNone(acquired_lock)
+        connection.close.assert_called_once_with()
+        log_error.assert_called_once()
+
+    def test_release_unlocks_commits_and_closes_connection(self) -> None:
+        connection = self._connection()
+        acquired_lock = DistributedLock(1234, connection)
+
+        acquired_lock.release()
+
+        connection.execute.assert_called_once()
+        connection.commit.assert_called_once_with()
+        connection.close.assert_called_once_with()
+
+    def test_release_still_closes_connection_when_unlock_fails(self) -> None:
+        connection = self._connection()
+        connection.execute.side_effect = SQLAlchemyError("unlock error")
+        acquired_lock = DistributedLock(1234, connection)
+
+        with self.assertRaises(SQLAlchemyError):
+            acquired_lock.release()
+
+        connection.commit.assert_not_called()
+        connection.close.assert_called_once_with()
+
+
+class SingleWorkerGuardTestCase(unittest.TestCase):
+    """Verify that only the process holding the lock runs startup tasks."""
+
+    def setUp(self) -> None:
+        SingleWorkerGuard._lock_status = LockStatus.NOT_ATTEMPTED
+        SingleWorkerGuard._acquired_lock = None
+
+    def tearDown(self) -> None:
+        SingleWorkerGuard._lock_status = LockStatus.NOT_ATTEMPTED
+        SingleWorkerGuard._acquired_lock = None
+
+    def test_acquired_worker_runs_all_tasks_with_one_lock_attempt(self) -> None:
+        acquired_lock = MagicMock(spec=DistributedLock)
+        executed_tasks: list[str] = []
+
+        @SingleWorkerGuard.once
+        def first_task() -> None:
+            executed_tasks.append("first")
+
+        @SingleWorkerGuard.once
+        def second_task() -> None:
+            executed_tasks.append("second")
+
+        with patch.object(
+            DistributedLock,
+            "try_acquire",
+            return_value=acquired_lock,
+        ) as try_acquire:
+            first_task()
+            second_task()
+
+        self.assertEqual(["first", "second"], executed_tasks)
+        try_acquire.assert_called_once_with(SingleWorkerGuard._LOCK_KEY)
+
+    def test_non_acquired_worker_skips_tasks_without_retrying(self) -> None:
+        task = MagicMock()
+        guarded_task = SingleWorkerGuard.once(task)
+
+        with (
+            patch.object(
+                DistributedLock,
+                "try_acquire",
+                return_value=None,
+            ) as try_acquire,
+            patch.object(lock_module.SQLBotLogUtil, "info"),
+        ):
+            guarded_task()
+            guarded_task()
+
+        task.assert_not_called()
+        try_acquire.assert_called_once_with(SingleWorkerGuard._LOCK_KEY)
+
+    def test_release_releases_lock_and_resets_process_state(self) -> None:
+        acquired_lock = MagicMock(spec=DistributedLock)
+        SingleWorkerGuard._lock_status = LockStatus.ACQUIRED
+        SingleWorkerGuard._acquired_lock = acquired_lock
+
+        SingleWorkerGuard.release()
+
+        acquired_lock.release.assert_called_once_with()
+        self.assertIsNone(SingleWorkerGuard._acquired_lock)
+        self.assertIs(
+            LockStatus.NOT_ATTEMPTED,
+            SingleWorkerGuard._lock_status,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()