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