From e545c28ac301e6f9feae247c513d8ea8f83df294 Mon Sep 17 00:00:00 2001 From: NanPan <111261006+poijygfdyy@users.noreply.github.com> Date: Sun, 20 Sep 2026 22:42:23 +0800 Subject: [PATCH] fix(persistence): drain postgres bootstrap unlock across cancellation (#5573) --- .../harness/deerflow/persistence/AGENTS.md | 3 + .../harness/deerflow/persistence/bootstrap.py | 4 +- backend/tests/test_agent_guidance_check.py | 1 + .../test_persistence_bootstrap_pg_lock.py | 62 +++++++++++++++++++ 4 files changed, 69 insertions(+), 1 deletion(-) create mode 100644 backend/packages/harness/deerflow/persistence/AGENTS.md diff --git a/backend/packages/harness/deerflow/persistence/AGENTS.md b/backend/packages/harness/deerflow/persistence/AGENTS.md new file mode 100644 index 000000000..74bd638e3 --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/AGENTS.md @@ -0,0 +1,3 @@ +# Persistence lifecycle + +Postgres bootstrap owns its session-scoped advisory lock until `pg_advisory_unlock` completes. Drain that unlock across host cancellation before leaving the SQLAlchemy connection context; repeated cancellation must not return a pooled session while it still holds the bootstrap mutex. Ordinary database errors remain best-effort and are logged. diff --git a/backend/packages/harness/deerflow/persistence/bootstrap.py b/backend/packages/harness/deerflow/persistence/bootstrap.py index de4862a50..b66f7acfb 100644 --- a/backend/packages/harness/deerflow/persistence/bootstrap.py +++ b/backend/packages/harness/deerflow/persistence/bootstrap.py @@ -99,6 +99,8 @@ from sqlalchemy import inspect as sa_inspect from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncEngine +from deerflow.utils.file_io import await_drained + logger = logging.getLogger(__name__) @@ -556,7 +558,7 @@ async def _postgres_lock(engine: AsyncEngine): yield finally: try: - await conn.execute(text("SELECT pg_advisory_unlock(:k)"), {"k": _PG_LOCK_KEY}) + await await_drained(conn.execute(text("SELECT pg_advisory_unlock(:k)"), {"k": _PG_LOCK_KEY})) except Exception: # noqa: BLE001 logger.warning("bootstrap: pg_advisory_unlock raised; session close will release", exc_info=True) diff --git a/backend/tests/test_agent_guidance_check.py b/backend/tests/test_agent_guidance_check.py index 7119dc7fa..bc85b1d74 100644 --- a/backend/tests/test_agent_guidance_check.py +++ b/backend/tests/test_agent_guidance_check.py @@ -27,6 +27,7 @@ EXPECTED_GUIDANCE_PATHS = { "backend/packages/harness/deerflow/sandbox/AGENTS.md", "backend/packages/harness/deerflow/mcp/AGENTS.md", "backend/packages/harness/deerflow/models/AGENTS.md", + "backend/packages/harness/deerflow/persistence/AGENTS.md", "backend/packages/harness/deerflow/persistence/migrations/AGENTS.md", "backend/packages/harness/deerflow/persistence/user/AGENTS.md", "backend/packages/harness/deerflow/reflection/AGENTS.md", diff --git a/backend/tests/test_persistence_bootstrap_pg_lock.py b/backend/tests/test_persistence_bootstrap_pg_lock.py index 5755d3b2e..c4087fbf6 100644 --- a/backend/tests/test_persistence_bootstrap_pg_lock.py +++ b/backend/tests/test_persistence_bootstrap_pg_lock.py @@ -25,6 +25,8 @@ runtime effect is Postgres's contract, not ours. from __future__ import annotations +import asyncio + import pytest from deerflow.persistence import bootstrap as bootstrap_mod @@ -96,3 +98,63 @@ async def test_postgres_lock_releases_even_if_body_raises() -> None: sqls = [stmt for stmt, _ in engine.conn.executed] assert any("pg_advisory_unlock" in s for s in sqls), f"unlock missing after body error; saw: {sqls}" + + +class _BlockingUnlockConn(_FakeAsyncConn): + def __init__(self) -> None: + super().__init__() + self.unlock_started = asyncio.Event() + self.allow_unlock = asyncio.Event() + self.unlock_finished = asyncio.Event() + + async def execute(self, stmt, params=None): + sql = str(stmt) + self.executed.append((sql, params)) + if "pg_advisory_unlock" in sql: + self.unlock_started.set() + await self.allow_unlock.wait() + self.unlock_finished.set() + return None + + +class _BlockingUnlockEngine: + def __init__(self) -> None: + self.conn = _BlockingUnlockConn() + + def connect(self) -> _BlockingUnlockConn: + return self.conn + + +@pytest.mark.asyncio +async def test_postgres_lock_drains_unlock_across_repeated_cancellation() -> None: + engine = _BlockingUnlockEngine() + entered = asyncio.Event() + hold_body = asyncio.Event() + + async def owner() -> None: + async with bootstrap_mod._postgres_lock(engine): # type: ignore[arg-type] + entered.set() + await hold_body.wait() + + task = asyncio.create_task(owner()) + await asyncio.wait_for(entered.wait(), timeout=1) + + task.cancel() + await asyncio.wait_for(engine.conn.unlock_started.wait(), timeout=1) + + try: + task.cancel() + for _ in range(5): + await asyncio.sleep(0) + assert not task.done(), "bootstrap returned before advisory unlock finished" + assert not engine.conn.unlock_finished.is_set() + + engine.conn.allow_unlock.set() + with pytest.raises(asyncio.CancelledError): + await task + assert engine.conn.unlock_finished.is_set() + finally: + engine.conn.allow_unlock.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True)