Files
bytedance--deer-flow/backend/tests/test_persistence_bootstrap_concurrency.py
2026-07-13 11:59:58 +08:00

151 lines
5.2 KiB
Python

"""Concurrency safety tests for ``bootstrap_schema``.
The contract: N concurrent callers against the same DB always converge to
``alembic_version == head`` without exceptions and without duplicate schema
mutations.
We model concurrency at the *async-task* level here (multiple coroutines
inside one process). SQLite is single-node by deployment, so within-process
serialisation -- which is what the per-engine ``_SQLITE_LOCKS`` entry
provides -- is the realistic boundary. Cross-process serialisation falls
through to SQLite's own write lock + ``PRAGMA busy_timeout`` plus the
idempotent revision helpers.
"""
from __future__ import annotations
import asyncio
from pathlib import Path
import pytest
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine
import deerflow.persistence.models # noqa: F401
from deerflow.persistence import bootstrap as bootstrap_mod
from deerflow.persistence.bootstrap import bootstrap_schema
pytestmark = pytest.mark.asyncio
HEAD = "0004_run_ownership"
def _url(tmp_path: Path) -> str:
return f"sqlite+aiosqlite:///{(tmp_path / 'concurrent.db').as_posix()}"
async def _alembic_version(engine) -> str | None:
async with engine.connect() as conn:
row = await conn.execute(sa.text("SELECT version_num FROM alembic_version"))
return row.scalar()
async def _runs_columns(engine) -> set[str]:
async with engine.connect() as conn:
return await conn.run_sync(lambda c: {col["name"] for col in sa.inspect(c).get_columns("runs")})
async def test_two_concurrent_bootstrap_callers_converge(tmp_path: Path) -> None:
engine = create_async_engine(_url(tmp_path))
try:
await asyncio.gather(
bootstrap_schema(engine, backend="sqlite"),
bootstrap_schema(engine, backend="sqlite"),
)
assert await _alembic_version(engine) == HEAD
assert "token_usage_by_model" in await _runs_columns(engine)
finally:
await engine.dispose()
async def test_five_concurrent_bootstrap_callers_converge(tmp_path: Path) -> None:
engine = create_async_engine(_url(tmp_path))
try:
await asyncio.gather(*(bootstrap_schema(engine, backend="sqlite") for _ in range(5)))
assert await _alembic_version(engine) == HEAD
finally:
await engine.dispose()
async def test_cancelled_caller_does_not_block_others(tmp_path: Path) -> None:
"""Cancelling one task mid-bootstrap must not strand the lock or the DB.
After the cancel, a subsequent ``bootstrap_schema`` call must still reach
head.
"""
engine = create_async_engine(_url(tmp_path))
try:
task = asyncio.create_task(bootstrap_schema(engine, backend="sqlite"))
# Give the event loop a turn so the task can start; then cancel.
await asyncio.sleep(0)
task.cancel()
# Cancelled task may have raced past the lock; swallow either outcome.
try:
await task
except (asyncio.CancelledError, Exception): # noqa: BLE001
pass
# Lock must be free for the next caller.
await bootstrap_schema(engine, backend="sqlite")
assert await _alembic_version(engine) == HEAD
finally:
await engine.dispose()
async def test_late_caller_after_head_is_noop(monkeypatch, tmp_path: Path) -> None:
"""When the first caller leaves the DB at head, the second observes
'versioned' and skips create_all / stamp -- it only runs upgrade head,
which is alembic-no-op.
We use a monkeypatched ``_upgrade`` counter to assert the second caller's
upgrade ran but did no real work (no new revision applied).
"""
engine = create_async_engine(_url(tmp_path))
try:
# First caller: empty branch.
await bootstrap_schema(engine, backend="sqlite")
first_version = await _alembic_version(engine)
assert first_version == HEAD
upgrade_calls: list[str] = []
original_upgrade = bootstrap_mod._upgrade
def counting_upgrade(cfg, rev: str) -> None:
upgrade_calls.append(rev)
original_upgrade(cfg, rev)
monkeypatch.setattr(bootstrap_mod, "_upgrade", counting_upgrade)
# Second caller: versioned branch -> calls _upgrade('head').
await bootstrap_schema(engine, backend="sqlite")
assert upgrade_calls == ["head"]
assert await _alembic_version(engine) == HEAD
finally:
await engine.dispose()
async def test_slow_upgrade_does_not_corrupt_concurrent_state(monkeypatch, tmp_path: Path) -> None:
"""Inject a delay into the upgrade path; concurrent callers must still
converge to head with no exceptions."""
engine = create_async_engine(_url(tmp_path))
try:
original_upgrade = bootstrap_mod._upgrade
def slow_upgrade(cfg, rev: str) -> None:
import time # noqa: PLC0415
time.sleep(0.2)
original_upgrade(cfg, rev)
monkeypatch.setattr(bootstrap_mod, "_upgrade", slow_upgrade)
await asyncio.gather(
bootstrap_schema(engine, backend="sqlite"),
bootstrap_schema(engine, backend="sqlite"),
bootstrap_schema(engine, backend="sqlite"),
)
assert await _alembic_version(engine) == HEAD
finally:
await engine.dispose()