mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-25 05:56:18 +00:00
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:
parent
0ab3b227af
commit
4117a50676
@ -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(
|
||||
|
||||
@ -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)
|
||||
Loading…
x
Reference in New Issue
Block a user