From 1363e384eda0184ba62454afbd9059d4435e9643 Mon Sep 17 00:00:00 2001 From: sjj <626650687@qq.com> Date: Sun, 19 Jul 2026 20:20:24 +0800 Subject: [PATCH 1/2] fix: avoid duplicate startup embedding jobs --- backend/common/utils/distributed_lock.py | 133 +++++++++++++++++++++++ backend/main.py | 25 ++++- 2 files changed, 154 insertions(+), 4 deletions(-) create mode 100644 backend/common/utils/distributed_lock.py diff --git a/backend/common/utils/distributed_lock.py b/backend/common/utils/distributed_lock.py new file mode 100644 index 00000000..7c25418d --- /dev/null +++ b/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 diff --git a/backend/main.py b/backend/main.py index a8bab974..4718702c 100644 --- a/backend/main.py +++ b/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: From 5b0cb82ed99f0e125132f2c97ed24bcdca5160c5 Mon Sep 17 00:00:00 2001 From: shaojunjie <626650687@qq.com> Date: Wed, 22 Jul 2026 10:27:47 +0800 Subject: [PATCH 2/2] test: cover startup distributed lock behavior --- tests/test_distributed_lock.py | 170 +++++++++++++++++++++++++++++++++ 1 file changed, 170 insertions(+) create mode 100644 tests/test_distributed_lock.py diff --git a/tests/test_distributed_lock.py b/tests/test_distributed_lock.py new file mode 100644 index 00000000..475fc295 --- /dev/null +++ b/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()