learningcircuit--local-deep-research/tests/database/test_concurrent_user_db_open.py

"""Concurrency tests for ``DatabaseManager.open_user_database`` cold-open.

Regression test for the race where two simultaneous first-opens of the same
user's database both ran Alembic migrations against one database file at once,
with the loser failing (Alembic's non-thread-safe module-level proxy / the
``alembic_version`` row update). The per-user init lock must serialize the
cold-open so the engine build + migration runs exactly once.
"""

import threading
import time

import pytest

import local_deep_research.database.initialize as init_mod
from local_deep_research.database.encrypted_db import DatabaseManager


@pytest.fixture
def db_manager(tmp_path, monkeypatch):
    monkeypatch.setattr(
        "local_deep_research.database.encrypted_db.get_data_directory",
        lambda: tmp_path,
    )
    manager = DatabaseManager()
    yield manager
    for username in list(manager.connections.keys()):
        manager.close_user_database(username)


def test_get_init_lock_is_per_user_and_stable(db_manager):
    """Same user -> same lock (serializes); different users -> different locks."""
    alice_1 = db_manager._get_init_lock("alice")
    alice_2 = db_manager._get_init_lock("alice")
    bob = db_manager._get_init_lock("bob")
    assert alice_1 is alice_2
    assert alice_1 is not bob


def test_concurrent_cold_open_runs_init_once(db_manager, monkeypatch):
    """N simultaneous first-opens of one user trigger exactly one cold-open."""
    username, password = "raceuser", "TestPassword123!"
    db_manager.create_user_database(username, password)
    # Evict (and dispose) the cached engine so the next opens are cold-opens
    # that hit the build + migrate path concurrently. The DB file remains.
    db_manager.close_user_database(username)

    n_threads = 8
    calls = []
    barrier = threading.Barrier(n_threads)

    def counting_init(engine, *args, **kwargs):
        # create_user_database already migrated the DB to head, so a no-op is
        # correct here; the sleep widens the window so all threads pile up on
        # the per-user lock while the first thread is inside the cold-open.
        calls.append(threading.current_thread().name)
        # Real wall-clock sleep widens the race window so the other threads
        # pile up on the per-user lock; freezegun cannot model concurrent
        # thread timing, so the sleep is intentional here.
        time.sleep(0.25)  # allow: unmarked-sleep

    monkeypatch.setattr(init_mod, "initialize_database", counting_init)

    results = []
    errors = []

    def worker():
        barrier.wait()
        try:
            results.append(db_manager.open_user_database(username, password))
        except Exception as exc:  # noqa: BLE001 - record any failure
            errors.append(exc)

    threads = [
        threading.Thread(target=worker, name=f"open-{i}")
        for i in range(n_threads)
    ]
    for t in threads:
        t.start()
    for t in threads:
        t.join(timeout=30)

    assert not errors, f"cold-open raced and failed: {errors!r}"
    assert len(results) == n_threads
    # The cold-open (hence the migration) ran exactly once despite n_threads
    # simultaneous first-opens.
    assert len(calls) == 1, (
        f"init ran {len(calls)}x; expected 1 (cold-open not serialized)"
    )
    # Every caller received the single engine the cold-open built and cached.
    assert results[0] is not None
    assert all(r is results[0] for r in results)


def test_close_user_database_keeps_init_lock(db_manager):
    """close_user_database intentionally retains the per-user init lock.

    Dropping it on close would let a concurrent open that already holds a
    reference to the old lock race a later open that creates a fresh one --
    two cold-opens migrating one DB file at once, the race the lock prevents.
    The lock is kept (bounded, one small Lock per username); only
    close_all_databases clears the dict wholesale, at shutdown.
    """
    username, password = "lockcleanup", "TestPassword123!"
    db_manager.create_user_database(username, password)
    db_manager._get_init_lock(username)  # ensure the lock exists
    assert username in db_manager._init_locks

    db_manager.close_user_database(username)

    # Retained on purpose -- see docstring; only close_all clears it.
    assert username in db_manager._init_locks


def test_close_all_databases_clears_init_locks(db_manager):
    """close_all_databases clears the per-user init-lock dict too."""
    for name in ("user_a", "user_b"):
        db_manager.create_user_database(name, "TestPassword123!")
        db_manager._get_init_lock(name)
    assert db_manager._init_locks

    db_manager.close_all_databases()

    assert db_manager._init_locks == {}


def test_concurrent_opens_of_different_users_run_in_parallel(
    db_manager, monkeypatch
):
    """Cold-opens of *different* users must not serialize against each other.

    The lock is deliberately per-user, not a single global init lock, so two
    users opening at once proceed in parallel. This is asserted
    deterministically (no timing): a 2-party barrier inside the patched init
    only releases when *both* users' cold-opens are inside it simultaneously.
    If the opens serialized (e.g. a regression to one global lock), the first
    thread would block at the barrier while holding the lock and the second
    could never reach it -> the barrier times out with BrokenBarrierError,
    surfaced as an error and failing the test.
    """
    users = [("alice_par", "TestPassword123!"), ("bob_par", "TestPassword123!")]
    for name, pw in users:
        db_manager.create_user_database(name, pw)
        # Evict the cached engine so the next open is a cold-open.
        db_manager.close_user_database(name)

    both_inside_init = threading.Barrier(len(users))

    def barrier_init(engine, *args, **kwargs):
        # Both users' cold-opens must be in here at once for this to return;
        # a global lock would deadlock one of them out and trip the timeout.
        both_inside_init.wait(timeout=10)

    monkeypatch.setattr(init_mod, "initialize_database", barrier_init)

    results = []
    errors = []

    def worker(name, pw):
        try:
            results.append(db_manager.open_user_database(name, pw))
        except Exception as exc:  # noqa: BLE001 - record any failure
            errors.append(exc)

    threads = [
        threading.Thread(target=worker, args=(n, p), name=f"open-{n}")
        for n, p in users
    ]
    for t in threads:
        t.start()
    for t in threads:
        t.join(timeout=30)

    assert not errors, f"cross-user opens serialized / failed: {errors!r}"
    assert len(results) == len(users)
    assert all(r is not None for r in results)