diff --git a/backend/packages/harness/deerflow/persistence/AGENTS.md b/backend/packages/harness/deerflow/persistence/AGENTS.md index 850696418..cb697fc59 100644 --- a/backend/packages/harness/deerflow/persistence/AGENTS.md +++ b/backend/packages/harness/deerflow/persistence/AGENTS.md @@ -3,3 +3,5 @@ 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. When `database.postgres_schema` is configured, both async ORM connections and the synchronous SQLAlchemy connections used by DB-backed custom agents and managed subagents must use the same `search_path`; preserve this invariant when adding another persistence entry point. + +Alembic stamp/upgrade workers started inside `bootstrap_schema()` remain owned by the bootstrap critical section until the worker finishes. Drain those `asyncio.to_thread()` calls across host cancellation before releasing the in-process SQLite bootstrap lock or PostgreSQL advisory lock; otherwise another bootstrap can overlap a still-running migration worker. diff --git a/backend/packages/harness/deerflow/persistence/bootstrap.py b/backend/packages/harness/deerflow/persistence/bootstrap.py index b66f7acfb..df066a470 100644 --- a/backend/packages/harness/deerflow/persistence/bootstrap.py +++ b/backend/packages/harness/deerflow/persistence/bootstrap.py @@ -507,12 +507,12 @@ def _run_baseline_create_all_sync(sync_conn: Any) -> None: def _stamp(cfg: AlembicConfig, revision: str) -> None: - """Synchronous alembic stamp; callers must wrap in ``asyncio.to_thread``.""" + """Synchronous alembic stamp; callers must drain via ``await_drained(asyncio.to_thread(...))``.""" alembic_command.stamp(cfg, revision) def _upgrade(cfg: AlembicConfig, revision: str) -> None: - """Synchronous alembic upgrade; callers must wrap in ``asyncio.to_thread``.""" + """Synchronous alembic upgrade; callers must drain via ``await_drained(asyncio.to_thread(...))``.""" alembic_command.upgrade(cfg, revision) @@ -611,7 +611,8 @@ async def bootstrap_schema(engine: AsyncEngine, *, backend: str, postgres_schema Branch dispatch is documented at module top. ``alembic.command.stamp`` and ``alembic.command.upgrade`` are synchronous and would block the event - loop; both are wrapped in ``asyncio.to_thread``. + loop; both are wrapped in ``await_drained(asyncio.to_thread(...))`` so the + worker finishes before the bootstrap lock is released on cancellation. *postgres_schema*, when set, is forwarded to the alembic config so the alembic-spawned engine pins its ``search_path`` to that schema. The target @@ -631,7 +632,7 @@ async def bootstrap_schema(engine: AsyncEngine, *, backend: str, postgres_schema logger.info("bootstrap: branch=empty -> create_all + stamp head (%s)", head) async with engine.begin() as conn: await conn.run_sync(_run_create_all_sync) - await asyncio.to_thread(_stamp, cfg, head) + await await_drained(asyncio.to_thread(_stamp, cfg, head)) elif decision == "legacy": logger.info( @@ -651,8 +652,8 @@ async def bootstrap_schema(engine: AsyncEngine, *, backend: str, postgres_schema # columns those revisions would add. async with engine.begin() as conn: await conn.run_sync(_run_baseline_create_all_sync) - await asyncio.to_thread(_stamp, cfg, _BASELINE_REVISION) - await asyncio.to_thread(_upgrade, cfg, "head") + await await_drained(asyncio.to_thread(_stamp, cfg, _BASELINE_REVISION)) + await await_drained(asyncio.to_thread(_upgrade, cfg, "head")) elif decision == "versioned": # The same revision id once named a different out-of-tree schema. @@ -669,7 +670,7 @@ async def bootstrap_schema(engine: AsyncEngine, *, backend: str, postgres_schema head, ) try: - await asyncio.to_thread(_upgrade, cfg, "head") + await await_drained(asyncio.to_thread(_upgrade, cfg, "head")) except CommandError: # SQLite has no cross-process bootstrap mutex. Another # process may advance 0018 to the reviewed 0019 after this diff --git a/backend/tests/test_persistence_bootstrap_concurrency.py b/backend/tests/test_persistence_bootstrap_concurrency.py index e74510990..292b6a675 100644 --- a/backend/tests/test_persistence_bootstrap_concurrency.py +++ b/backend/tests/test_persistence_bootstrap_concurrency.py @@ -15,6 +15,7 @@ idempotent revision helpers. from __future__ import annotations import asyncio +import threading from pathlib import Path import pytest @@ -148,3 +149,98 @@ async def test_slow_upgrade_does_not_corrupt_concurrent_state(monkeypatch, tmp_p assert await _alembic_version(engine) == HEAD finally: await engine.dispose() + + +async def test_cancelled_bootstrap_keeps_sqlite_lock_until_alembic_worker_finishes(monkeypatch, tmp_path: Path) -> None: + """Cancellation must not detach an Alembic worker from its bootstrap lock.""" + engine = create_async_engine(_url(tmp_path)) + stamp_started = threading.Event() + allow_stamp = threading.Event() + second_reflect_started = asyncio.Event() + reflect_calls = 0 + original_reflect = bootstrap_mod._reflect_state + + def blocking_stamp(_cfg, _revision: str) -> None: + stamp_started.set() + assert allow_stamp.wait(5), "test did not release the blocked stamp worker" + + def recording_reflect(sync_conn): + nonlocal reflect_calls + reflect_calls += 1 + if reflect_calls >= 2: + second_reflect_started.set() + return original_reflect(sync_conn) + + monkeypatch.setattr(bootstrap_mod, "_stamp", blocking_stamp) + monkeypatch.setattr(bootstrap_mod, "_upgrade", lambda *_args, **_kwargs: None) + monkeypatch.setattr(bootstrap_mod, "_reflect_state", recording_reflect) + + first = asyncio.create_task(bootstrap_schema(engine, backend="sqlite")) + second: asyncio.Task[None] | None = None + try: + assert await asyncio.to_thread(stamp_started.wait, 2), "bootstrap did not reach the Alembic stamp worker" + + first.cancel() + await asyncio.sleep(0) + first.cancel() + for _ in range(10): + await asyncio.sleep(0) + + assert not first.done(), "cancelled bootstrap returned while its Alembic worker was still running" + + second = asyncio.create_task(bootstrap_schema(engine, backend="sqlite")) + for _ in range(10): + await asyncio.sleep(0) + assert not second_reflect_started.is_set(), "cancelled bootstrap released the SQLite mutex before its worker finished" + finally: + allow_stamp.set() + await asyncio.gather(first, *(task for task in (second,) if task is not None), return_exceptions=True) + await engine.dispose() + + +async def test_cancelled_versioned_bootstrap_keeps_sqlite_lock_until_upgrade_worker_finishes(monkeypatch, tmp_path: Path) -> None: + """The versioned upgrade path must retain the same lock while its worker drains.""" + engine = create_async_engine(_url(tmp_path)) + await bootstrap_schema(engine, backend="sqlite") + + upgrade_started = threading.Event() + allow_upgrade = threading.Event() + second_reflect_started = asyncio.Event() + reflect_calls = 0 + original_reflect = bootstrap_mod._reflect_state + + def blocking_upgrade(_cfg, _revision: str) -> None: + upgrade_started.set() + assert allow_upgrade.wait(5), "test did not release the blocked upgrade worker" + + def recording_reflect(sync_conn): + nonlocal reflect_calls + reflect_calls += 1 + if reflect_calls >= 2: + second_reflect_started.set() + return original_reflect(sync_conn) + + monkeypatch.setattr(bootstrap_mod, "_upgrade", blocking_upgrade) + monkeypatch.setattr(bootstrap_mod, "_reflect_state", recording_reflect) + + first = asyncio.create_task(bootstrap_schema(engine, backend="sqlite")) + second: asyncio.Task[None] | None = None + try: + assert await asyncio.to_thread(upgrade_started.wait, 2), "bootstrap did not reach the Alembic upgrade worker" + + first.cancel() + await asyncio.sleep(0) + first.cancel() + for _ in range(10): + await asyncio.sleep(0) + + assert not first.done(), "cancelled bootstrap returned while its Alembic upgrade worker was still running" + + second = asyncio.create_task(bootstrap_schema(engine, backend="sqlite")) + for _ in range(10): + await asyncio.sleep(0) + assert not second_reflect_started.is_set(), "cancelled bootstrap released the SQLite mutex before upgrade finished" + finally: + allow_upgrade.set() + await asyncio.gather(first, *(task for task in (second,) if task is not None), return_exceptions=True) + await engine.dispose()