fix(api): bound bootstrap lock during startup
Some checks failed
CD Pipeline / workflow-shape (push) Successful in 0s
CD Pipeline / cancel-stale-cd (push) Has been skipped
CD Pipeline / tests (push) Failing after 34s
CD Pipeline / post-deploy-checks (push) Has been cancelled
CD Pipeline / build-and-deploy (push) Has been skipped
Some checks failed
CD Pipeline / workflow-shape (push) Successful in 0s
CD Pipeline / cancel-stale-cd (push) Has been skipped
CD Pipeline / tests (push) Failing after 34s
CD Pipeline / post-deploy-checks (push) Has been cancelled
CD Pipeline / build-and-deploy (push) Has been skipped
This commit is contained in:
@@ -13,6 +13,7 @@ Features:
|
||||
- 所有 Episodic Memory 必須使用 PostgreSQL
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
@@ -183,6 +184,9 @@ async def get_db_context(project_id: str | None = None) -> AsyncGenerator[AsyncS
|
||||
# =============================================================================
|
||||
|
||||
_DB_BOOTSTRAP_LOCK_NAME = "awoooi:init_db:ddl"
|
||||
_DB_BOOTSTRAP_LOCK_WAIT_SECONDS = 45.0
|
||||
_DB_BOOTSTRAP_LOCK_POLL_SECONDS = 2.0
|
||||
_DB_BOOTSTRAP_DDL_WAIT_SECONDS = 120.0
|
||||
|
||||
|
||||
async def init_db() -> None:
|
||||
@@ -196,13 +200,27 @@ async def init_db() -> None:
|
||||
async with engine.connect() as lock_conn:
|
||||
# 2026-05-24 ogt + Codex: 兩個 API replica 同時啟動時,PostgreSQL 會在
|
||||
# ALTER TABLE ... IF NOT EXISTS 上互相等待並 deadlock。整段 bootstrap
|
||||
# DDL 必須序列化,避免 rollout 因一個 pod CrashLoop 變成 1/2 ready。
|
||||
await lock_conn.execute(
|
||||
text("SELECT pg_advisory_lock(hashtext(:lock_name))"),
|
||||
{"lock_name": _DB_BOOTSTRAP_LOCK_NAME},
|
||||
)
|
||||
# DDL 必須序列化。
|
||||
# 2026-06-30 Codex: production rollout 觀察到新 API pod 卡在 init_db()
|
||||
# 超過 startup probe 預算。bootstrap DDL 是冪等補欄流程,不能再用
|
||||
# 無界 pg_advisory_lock 把新版 API route 擋成永久舊版;改為有界等待,
|
||||
# 逾時則 fail-visible skip,讓 serving path 先恢復,再由 verifier 回報。
|
||||
lock_acquired = await _try_acquire_bootstrap_lock(lock_conn)
|
||||
if not lock_acquired:
|
||||
return
|
||||
|
||||
try:
|
||||
await _run_init_db_ddl(engine)
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
_run_init_db_ddl(engine),
|
||||
timeout=_DB_BOOTSTRAP_DDL_WAIT_SECONDS,
|
||||
)
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
"database_bootstrap_ddl_timeout_skipped",
|
||||
timeout_seconds=_DB_BOOTSTRAP_DDL_WAIT_SECONDS,
|
||||
lock_name=_DB_BOOTSTRAP_LOCK_NAME,
|
||||
)
|
||||
finally:
|
||||
await lock_conn.execute(
|
||||
text("SELECT pg_advisory_unlock(hashtext(:lock_name))"),
|
||||
@@ -210,6 +228,34 @@ async def init_db() -> None:
|
||||
)
|
||||
|
||||
|
||||
async def _try_acquire_bootstrap_lock(lock_conn: object) -> bool:
|
||||
loop = asyncio.get_running_loop()
|
||||
deadline = loop.time() + _DB_BOOTSTRAP_LOCK_WAIT_SECONDS
|
||||
|
||||
while True:
|
||||
result = await lock_conn.execute(
|
||||
text("SELECT pg_try_advisory_lock(hashtext(:lock_name))"),
|
||||
{"lock_name": _DB_BOOTSTRAP_LOCK_NAME},
|
||||
)
|
||||
if bool(result.scalar()):
|
||||
logger.info(
|
||||
"database_bootstrap_lock_acquired",
|
||||
lock_name=_DB_BOOTSTRAP_LOCK_NAME,
|
||||
)
|
||||
return True
|
||||
|
||||
remaining = deadline - loop.time()
|
||||
if remaining <= 0:
|
||||
logger.warning(
|
||||
"database_bootstrap_lock_timeout_skipped",
|
||||
timeout_seconds=_DB_BOOTSTRAP_LOCK_WAIT_SECONDS,
|
||||
lock_name=_DB_BOOTSTRAP_LOCK_NAME,
|
||||
)
|
||||
return False
|
||||
|
||||
await asyncio.sleep(min(_DB_BOOTSTRAP_LOCK_POLL_SECONDS, remaining))
|
||||
|
||||
|
||||
async def _run_init_db_ddl(engine: AsyncEngine) -> None:
|
||||
"""
|
||||
Run idempotent DB bootstrap DDL while caller holds the bootstrap advisory lock.
|
||||
|
||||
@@ -14,9 +14,18 @@ from typing import Any
|
||||
import pytest
|
||||
|
||||
|
||||
class _FakeScalarResult:
|
||||
def __init__(self, value: bool) -> None:
|
||||
self.value = value
|
||||
|
||||
def scalar(self) -> bool:
|
||||
return self.value
|
||||
|
||||
|
||||
class _FakeLockConnection:
|
||||
def __init__(self) -> None:
|
||||
def __init__(self, lock_results: list[bool] | None = None) -> None:
|
||||
self.statements: list[str] = []
|
||||
self._lock_results = lock_results or [True]
|
||||
|
||||
async def __aenter__(self) -> _FakeLockConnection:
|
||||
return self
|
||||
@@ -24,14 +33,22 @@ class _FakeLockConnection:
|
||||
async def __aexit__(self, *_exc: object) -> None:
|
||||
return None
|
||||
|
||||
async def execute(self, statement: object, params: dict[str, str] | None = None) -> None:
|
||||
self.statements.append(str(statement))
|
||||
async def execute(
|
||||
self,
|
||||
statement: object,
|
||||
params: dict[str, str] | None = None,
|
||||
) -> _FakeScalarResult:
|
||||
sql = str(statement)
|
||||
self.statements.append(sql)
|
||||
assert params == {"lock_name": "awoooi:init_db:ddl"}
|
||||
if "pg_try_advisory_lock" in sql:
|
||||
return _FakeScalarResult(self._lock_results.pop(0))
|
||||
return _FakeScalarResult(True)
|
||||
|
||||
|
||||
class _FakeEngine:
|
||||
def __init__(self) -> None:
|
||||
self.lock_conn = _FakeLockConnection()
|
||||
def __init__(self, lock_results: list[bool] | None = None) -> None:
|
||||
self.lock_conn = _FakeLockConnection(lock_results=lock_results)
|
||||
|
||||
def connect(self) -> _FakeLockConnection:
|
||||
return self.lock_conn
|
||||
@@ -53,7 +70,7 @@ async def test_init_db_serializes_bootstrap_ddl_with_advisory_lock(monkeypatch):
|
||||
await db_base.init_db()
|
||||
|
||||
assert calls == [fake_engine]
|
||||
assert "pg_advisory_lock" in fake_engine.lock_conn.statements[0]
|
||||
assert "pg_try_advisory_lock" in fake_engine.lock_conn.statements[0]
|
||||
assert "pg_advisory_unlock" in fake_engine.lock_conn.statements[-1]
|
||||
|
||||
|
||||
@@ -72,7 +89,52 @@ async def test_init_db_releases_bootstrap_lock_when_ddl_fails(monkeypatch):
|
||||
with pytest.raises(RuntimeError, match="ddl failed"):
|
||||
await db_base.init_db()
|
||||
|
||||
assert "pg_advisory_lock" in fake_engine.lock_conn.statements[0]
|
||||
assert "pg_try_advisory_lock" in fake_engine.lock_conn.statements[0]
|
||||
assert "pg_advisory_unlock" in fake_engine.lock_conn.statements[-1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_db_skips_bootstrap_when_advisory_lock_times_out(monkeypatch):
|
||||
from src.db import base as db_base
|
||||
|
||||
fake_engine = _FakeEngine(lock_results=[False])
|
||||
calls: list[object] = []
|
||||
|
||||
async def fake_run_init_db_ddl(engine: object) -> None:
|
||||
calls.append(engine)
|
||||
|
||||
monkeypatch.setattr(db_base, "get_engine", lambda: fake_engine)
|
||||
monkeypatch.setattr(db_base, "_run_init_db_ddl", fake_run_init_db_ddl)
|
||||
monkeypatch.setattr(db_base, "_DB_BOOTSTRAP_LOCK_WAIT_SECONDS", 0.0)
|
||||
|
||||
await db_base.init_db()
|
||||
|
||||
assert calls == []
|
||||
assert "pg_try_advisory_lock" in fake_engine.lock_conn.statements[0]
|
||||
assert all("pg_advisory_unlock" not in stmt for stmt in fake_engine.lock_conn.statements)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_db_releases_bootstrap_lock_when_ddl_times_out(monkeypatch):
|
||||
from src.db import base as db_base
|
||||
|
||||
fake_engine = _FakeEngine()
|
||||
|
||||
async def fake_run_init_db_ddl(_engine: object) -> None:
|
||||
raise AssertionError("wait_for should own ddl execution in this test")
|
||||
|
||||
async def fake_wait_for(coro: Awaitable[Any], timeout: float) -> None:
|
||||
assert timeout == 120.0
|
||||
coro.close()
|
||||
raise TimeoutError
|
||||
|
||||
monkeypatch.setattr(db_base, "get_engine", lambda: fake_engine)
|
||||
monkeypatch.setattr(db_base, "_run_init_db_ddl", fake_run_init_db_ddl)
|
||||
monkeypatch.setattr(db_base.asyncio, "wait_for", fake_wait_for)
|
||||
|
||||
await db_base.init_db()
|
||||
|
||||
assert "pg_try_advisory_lock" in fake_engine.lock_conn.statements[0]
|
||||
assert "pg_advisory_unlock" in fake_engine.lock_conn.statements[-1]
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user