deer-flow/backend/tests/test_persistence_bootstrap_concurrency.py

"""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
import threading
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 _get_head_revision, bootstrap_schema

pytestmark = pytest.mark.asyncio


HEAD = _get_head_revision()


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()


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()