fix(persistence): drain the auto-create maintenance dispose across cancellation (#5649)

_auto_create_postgres_db opens a throwaway engine against the server's
`postgres` database and disposes it in its finally, but that await was bare.
Startup cancellation landing there abandoned a half-disposed pool that nothing
else owned — the engine is a local, and the CREATE DATABASE has already
committed under AUTOCOMMIT, so the next boot does not take this path again.

Drain the dispose through the cancellation-safe await_drained() helper the
engine's own close path already uses, and pin it with a repeated-cancellation
regression.
This commit is contained in:
lihongyuan99 2026-09-23 10:14:35 +08:00 committed by GitHub
parent 0ab3b227af
commit 4117a50676
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 90 additions and 1 deletions

View File

@ -81,7 +81,11 @@ async def _auto_create_postgres_db(url: str) -> None:
await conn.execute(text(f'CREATE DATABASE "{db_name}"'))
logger.info("Auto-created PostgreSQL database: %s", db_name)
finally:
await maint_engine.dispose()
# Drain: the maintenance engine is local to this helper, so an
# abandoned dispose leaves its pool to the ``postgres`` database with
# nothing left to close it. Host cancellation must therefore be
# delivered only after disposal actually finishes.
await await_drained(maint_engine.dispose())
async def init_engine(

View File

@ -0,0 +1,85 @@
from __future__ import annotations
import asyncio
import pytest
from deerflow.persistence import engine as engine_mod
_TARGET_URL = "postgresql+asyncpg://deerflow:secret@127.0.0.1:5432/deerflow_target"
class _RecordingConnection:
def __init__(self) -> None:
self.executed = asyncio.Event()
self.statements: list[str] = []
async def __aenter__(self) -> _RecordingConnection:
return self
async def __aexit__(self, *exc_info: object) -> None:
return None
async def execute(self, statement: object) -> None:
self.statements.append(str(statement))
self.executed.set()
class _BlockingMaintenanceEngine:
"""Fake maintenance engine whose ``dispose()`` can be held open."""
def __init__(self) -> None:
self.connection = _RecordingConnection()
self.dispose_started = asyncio.Event()
self.allow_dispose = asyncio.Event()
self.dispose_finished = asyncio.Event()
def connect(self) -> _RecordingConnection:
return self.connection
async def dispose(self) -> None:
self.dispose_started.set()
await self.allow_dispose.wait()
self.dispose_finished.set()
@pytest.mark.asyncio
async def test_auto_create_db_drains_maintenance_dispose_across_repeated_cancellation(monkeypatch: pytest.MonkeyPatch) -> None:
"""A cancelled first boot must still dispose the engine it opened.
``_auto_create_postgres_db`` owns a throwaway engine against the server's
``postgres`` database. Its ``finally`` dispose was a bare await, so host
cancellation could return from the helper with that engine's pool still
live — and because the engine is a local, nothing else ever disposes it.
"""
maintenance_engine = _BlockingMaintenanceEngine()
task: asyncio.Task[None] | None = None
monkeypatch.setattr(engine_mod, "create_async_engine", lambda *args, **kwargs: maintenance_engine)
try:
task = asyncio.create_task(engine_mod._auto_create_postgres_db(_TARGET_URL))
await asyncio.wait_for(maintenance_engine.connection.executed.wait(), timeout=1)
assert 'CREATE DATABASE "deerflow_target"' in maintenance_engine.connection.statements[0]
await asyncio.wait_for(maintenance_engine.dispose_started.wait(), timeout=1)
task.cancel()
for _ in range(5):
await asyncio.sleep(0)
assert not task.done(), "database auto-creation returned before the maintenance engine was disposed"
task.cancel()
for _ in range(5):
await asyncio.sleep(0)
assert not task.done(), "repeated cancellation interrupted the maintenance engine dispose"
maintenance_engine.allow_dispose.set()
with pytest.raises(asyncio.CancelledError):
await task
assert maintenance_engine.dispose_finished.is_set()
finally:
maintenance_engine.allow_dispose.set()
if task is not None and not task.done():
task.cancel()
await asyncio.gather(task, return_exceptions=True)