From d21e7dea11c3343e8b12d2cd9633674f29c091ea Mon Sep 17 00:00:00 2001 From: sjj <626650687@qq.com> Date: Sun, 19 Jul 2026 20:20:24 +0800 Subject: [PATCH] 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.audit.schemas.request_context import RequestContextMiddlewareCommon 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: