TEST: Stop SQLite lock flakes on slow CI disks

doc/contributing/5_unit_tests.md

@@ -72,3 +72,5 @@ Run `uv sync --extra all` before validating fixture changes with `make unit-test
 To avoid replaying the migration history for every ordinary test, `sqlite_template` runs the real Alembic migrations and schema check once per pytest session (once in each xdist worker). SQLite's backup API copies that private, read-only template into each test's database. Rows, schema changes, temporary tables, and result files are not shared between tests.
 
 Tests of initialization, upgrades, downgrades, schema checks, and `reset_database()` must still call those production paths explicitly. The fixture does not replace or patch migration or reset APIs.
+
+`tests/unit/conftest.py` turns off fsync (`PRAGMA synchronous=OFF`) when a unit test opens a SQLite connection, including file-backed databases in `tmp_path`. A commit then does not wait for the disk, so a loaded CI disk cannot keep SQLite's write lock for seconds and push concurrent writers past the busy timeout. Locking and transactions do not change. If another connection holds an exclusive lock at that time, the new connection does not wait and keeps fsync. When a statement fails with `SQLITE_BUSY`, the error has a note with the time that SQLite waited for the lock. A wait near the busy timeout shows lock contention. A wait near 0 s shows that SQLite did not wait.

tests/unit/conftest.py

@@ -5,12 +5,18 @@
 import sqlite3
 import sys
 import tempfile
+import time
+import weakref
 from collections.abc import AsyncGenerator, Generator
 from contextlib import closing, contextmanager
+from typing import Any
 from unittest.mock import AsyncMock, patch
 
+import aiosqlite
 import pytest
-from sqlalchemy import Engine, create_engine
+from sqlalchemy import Engine, create_engine, event
+from sqlalchemy.engine import ExceptionContext, ExecutionContext
+from sqlalchemy.pool import ConnectionPoolEntry, Pool
 
 from pyrit.common.path import DATASETS_PATH
 from pyrit.datasets import SeedDatasetProvider
@@ -24,6 +30,55 @@
 os.environ["RETRY_WAIT_MIN_SECONDS"] = "0"
 os.environ["RETRY_WAIT_MAX_SECONDS"] = "1"
 
+_sqlite_statement_starts: weakref.WeakKeyDictionary[ExecutionContext, float] = weakref.WeakKeyDictionary()
+
+
+@event.listens_for(Pool, "connect", named=True)
+def _skip_sqlite_fsync(*, dbapi_connection: Any, connection_record: ConnectionPoolEntry, **kwargs: Any) -> None:
+    """
+    Turn off fsync when a unit test opens a SQLite connection.
+
+    On loaded CI disks, each fsync can take seconds. The writer holds SQLite's lock during fsync, so
+    concurrent writers can exceed the busy timeout. Locking and transactions do not change, and test
+    databases do not need durability after an OS crash. The PRAGMA reads the schema, so it does not
+    wait for an exclusive lock of another connection. In that case, the connection keeps fsync.
+    """
+    if not isinstance(connection_record.driver_connection, (sqlite3.Connection, aiosqlite.Connection)):
+        return
+    cursor = dbapi_connection.cursor()
+    try:
+        cursor.execute("PRAGMA busy_timeout")
+        busy_timeout = cursor.fetchone()[0]
+        cursor.execute("PRAGMA busy_timeout = 0")
+        try:
+            cursor.execute("PRAGMA synchronous = OFF")
+        except sqlite3.OperationalError as error:
+            if error.sqlite_errorcode != sqlite3.SQLITE_BUSY:
+                raise
+        finally:
+            cursor.execute(f"PRAGMA busy_timeout = {busy_timeout}")
+    finally:
+        cursor.close()
+
+
+@event.listens_for(Engine, "before_cursor_execute", named=True)
+def _record_sqlite_statement_start(*, context: ExecutionContext | None, **kwargs: Any) -> None:
+    if context is not None and context.dialect.name == "sqlite":
+        _sqlite_statement_starts[context] = time.monotonic()
+
+
+@event.listens_for(Engine, "handle_error")
+def _note_sqlite_busy_wait(context: ExceptionContext) -> None:
+    """Show how long a statement waited for the SQLite lock, to tell contention from an immediate conflict."""
+    if getattr(context.original_exception, "sqlite_errorcode", None) != sqlite3.SQLITE_BUSY:
+        return
+    started = _sqlite_statement_starts.get(context.execution_context) if context.execution_context else None
+    if started is not None and context.sqlalchemy_exception is not None:
+        context.sqlalchemy_exception.add_note(
+            f"SQLite returned SQLITE_BUSY after {time.monotonic() - started:.2f} s. A wait near the busy "
+            "timeout means that another connection kept the lock. A wait near 0 s means that SQLite did not wait."
+        )
+
 
 @pytest.fixture
 def block_dataset_provider_lookup() -> Generator[AsyncMock, None, None]:

tests/unit/memory/memory_interface/test_interface_observations.py

@@ -14,6 +14,7 @@
 from sqlalchemy.dialects import mssql
 from sqlalchemy.exc import OperationalError, SQLAlchemyError
 from sqlalchemy.orm import Session
+from unit.async_utils import settle_tasks_async, wait_for_completion_async
 from unit.mocks import run_memory_session_async
 
 from pyrit.memory import MemoryInterface, SQLiteMemory
@@ -640,11 +641,20 @@ async def test_concurrent_orm_score_deletions_clean_final_observation(file_memor
     async def delete_async(score_id: uuid.UUID) -> None:
         async with await file_memory.get_session_async() as session:
             entry = await session.get(ScoreEntry, score_id)
-            await asyncio.wait_for(ready.wait(), timeout=5)
+            await ready.wait()
             await session.delete(entry)
             await session.commit()
 
-    await asyncio.wait_for(asyncio.gather(*(delete_async(value.id) for value in (score, replay))), timeout=10)
+    tasks = [asyncio.create_task(delete_async(value.id)) for value in (score, replay)]
+    completion = asyncio.gather(*tasks)
+    try:
+        await wait_for_completion_async(future=completion)
+    finally:
+        for task in tasks:
+            task.cancel()
+        await settle_tasks_async(*tasks)
+        if not completion.cancelled():
+            completion.exception()
     assert file_memory._query_entries(ScoreEntry) == []
     assert (await file_memory.get_observations_async(observation_ids=[observation.id])) == []
     assert file_memory._query_entries(ObservationMessagePieceEntry) == []

tests/unit/memory/test_conversation_branching.py

@@ -36,7 +36,7 @@
     MessagePiece,
     TargetIdentifier,
 )
-from unit.async_utils import wait_for_completion_async
+from unit.async_utils import settle_tasks_async, wait_for_completion_async
 from unit.mocks import run_memory_session_async
 
 
@@ -385,7 +385,7 @@ async def test_concurrent_branch_creation_and_promotions_preserve_membership(
         start = asyncio.Barrier(10)
 
         async def register_async(*, conversation: Conversation) -> bool:
-            await asyncio.wait_for(start.wait(), timeout=10)
+            await start.wait()
             return await file_memory.add_conversation_branches_to_attack_async(
                 attack_result_id=attack.attack_result_id,
                 source_conversation=source,
@@ -394,18 +394,24 @@ async def register_async(*, conversation: Conversation) -> bool:
             )
 
         async def promote_async(*, conversation_id: str) -> bool:
-            await asyncio.wait_for(start.wait(), timeout=10)
+            await start.wait()
             return await file_memory.promote_attack_conversation_async(
                 attack_result_id=attack.attack_result_id, conversation_id=conversation_id
             )
 
-        results = await asyncio.wait_for(
-            asyncio.gather(
-                *(register_async(conversation=branch) for branch in branches),
-                *(promote_async(conversation_id=branch.conversation_id) for branch in initial),
-            ),
-            timeout=60,
-        )
+        tasks = [
+            *(asyncio.create_task(register_async(conversation=branch)) for branch in branches),
+            *(asyncio.create_task(promote_async(conversation_id=branch.conversation_id)) for branch in initial),
+        ]
+        completion = asyncio.gather(*tasks)
+        try:
+            results = await wait_for_completion_async(future=completion)
+        finally:
+            for task in tasks:
+                task.cancel()
+            await settle_tasks_async(*tasks)
+            if not completion.cancelled():
+                completion.exception()
         assert all(results)
         current = (await file_memory.get_attack_results_async(attack_result_ids=[attack.attack_result_id]))[0]
         assert current.get_active_conversation_ids() == {

tests/unit/memory/test_sqlite_fixtures.py

@@ -3,7 +3,9 @@
 
 import asyncio
 import gc
+import re
 import sqlite3
+import time
 import weakref
 from collections.abc import Generator
 from contextlib import asynccontextmanager, closing, contextmanager
@@ -12,18 +14,76 @@
 
 import pytest
 from sqlalchemy import Engine, create_engine, event, func, inspect, select, text
+from sqlalchemy.engine import AdaptedConnection
+from sqlalchemy.exc import OperationalError
+from sqlalchemy.ext.asyncio import create_async_engine
 from sqlalchemy.pool import ConnectionPoolEntry, StaticPool
 
 from pyrit.common.singleton import Singleton
 from pyrit.memory import CentralMemory, MemoryInterface, SQLiteMemory, migration
 from pyrit.memory.memory_models import Base, PromptMemoryEntry
+from pyrit.memory.memory_session import MemorySession
 from pyrit.models import MessagePiece
 from unit import conftest as unit_fixtures
+from unit.async_utils import settle_tasks_async, wait_for_event_async
+from unit.memory import test_conversation_branching
+from unit.memory.memory_interface import test_interface_observations
 
 _sqlite_instance = asynccontextmanager(unit_fixtures.sqlite_instance.__wrapped__)
 _sqlite_template = contextmanager(unit_fixtures.sqlite_template.__wrapped__)
 
 
+@pytest.mark.parametrize("operation", ["branch", "delete"])
+async def test_concurrent_sqlite_tests_release_workers_after_watchdog_async(tmp_path: Path, operation: str) -> None:
+    memory = SQLiteMemory.__new__(SQLiteMemory)
+    with patch.object(memory, "cleanup"):
+        memory.__init__(db_path=tmp_path / "watchdog.db", skip_schema_migration=True)
+    started, release = asyncio.Event(), asyncio.Event()
+    before = asyncio.all_tasks()
+    module = test_conversation_branching if operation == "branch" else test_interface_observations
+
+    async def hold_write_async(driver: object) -> None:
+        started.set()
+        await release.wait()
+
+    def pause_after_flush(session: MemorySession, context: object) -> None:
+        connection = session.connection().connection.dbapi_connection
+        assert isinstance(connection, AdaptedConnection)
+        connection.run_async(hold_write_async)
+
+    async def expire_async(*, future: asyncio.Future[object]) -> None:
+        event.listen(MemorySession, "after_flush", pause_after_flush)
+        try:
+            await wait_for_event_async(event=started, tasks=[future])
+            raise TimeoutError("forced watchdog expiry")
+        finally:
+            event.remove(MemorySession, "after_flush", pause_after_flush)
+
+    try:
+        await asyncio.to_thread(Base.metadata.create_all, memory.engine)
+        with patch.object(module, "wait_for_completion_async", side_effect=expire_async):
+            with pytest.raises(TimeoutError, match="forced watchdog expiry"):
+                if operation == "branch":
+                    test = test_conversation_branching.TestAtomicConversationBranching()
+                    await test.test_concurrent_branch_creation_and_promotions_preserve_membership(file_memory=memory)
+                else:
+                    await test_interface_observations.test_concurrent_orm_score_deletions_clean_final_observation(
+                        file_memory=memory
+                    )
+        assert not [task for task in asyncio.all_tasks() - before if not task.done()]
+        async with memory._get_async_engine().connect() as connection:
+            await connection.exec_driver_sql("PRAGMA busy_timeout = 0")
+            await connection.exec_driver_sql("BEGIN IMMEDIATE")
+            await connection.rollback()
+    finally:
+        release.set()
+        remaining = asyncio.all_tasks() - before
+        for task in remaining:
+            task.cancel()
+        await settle_tasks_async(*remaining)
+        await memory.dispose_engine_async()
+
+
 @pytest.fixture
 def sqlite_connections() -> Generator[list[sqlite3.Connection], None, None]:
     connections: list[sqlite3.Connection] = []
@@ -289,3 +349,51 @@ async def write_rows_async(worker: int) -> None:
     await asyncio.gather(*(write_rows_async(worker) for worker in range(4)))
     async with await sqlite_instance.get_session_async() as session:
         assert (await session.execute(text("SELECT COUNT(DISTINCT value) FROM fixture_threads"))).scalar_one() == 40
+
+
+async def test_file_sqlite_connections_skip_fsync_async(tmp_path: Path) -> None:
+    memory = SQLiteMemory.__new__(SQLiteMemory)
+    with patch.object(memory, "cleanup"):
+        memory.__init__(db_path=tmp_path / "fsync.db", skip_schema_migration=True)
+    try:
+        with memory.engine.connect() as connection:
+            assert connection.exec_driver_sql("PRAGMA synchronous").scalar_one() == 0
+        async with await memory.get_session_async() as session:
+            assert (await session.execute(text("PRAGMA synchronous"))).scalar_one() == 0
+    finally:
+        await memory.dispose_engine_async()
+
+
+def test_file_sqlite_connect_does_not_wait_for_an_exclusive_lock(tmp_path: Path) -> None:
+    database = tmp_path / "exclusive.db"
+    engine = create_engine(f"sqlite:///{database}", connect_args={"timeout": 30})
+    try:
+        with closing(sqlite3.connect(database, isolation_level=None)) as holder:
+            holder.execute("CREATE TABLE lock_probe (id INTEGER)")
+            holder.execute("BEGIN EXCLUSIVE")
+            started = time.monotonic()
+            with engine.connect() as connection:
+                assert time.monotonic() - started < 15
+                holder.execute("ROLLBACK")
+                assert connection.exec_driver_sql("PRAGMA busy_timeout").scalar_one() == 30_000
+                assert connection.exec_driver_sql("PRAGMA synchronous").scalar_one() == 2
+    finally:
+        engine.dispose()
+
+
+async def test_sqlite_busy_error_reports_lock_wait_async(tmp_path: Path) -> None:
+    database = tmp_path / "busy.db"
+    engine = create_async_engine(f"sqlite+aiosqlite:///{database}", connect_args={"timeout": 0.05})
+    try:
+        with closing(sqlite3.connect(database, isolation_level=None)) as holder:
+            holder.execute("CREATE TABLE lock_probe (id INTEGER)")
+            holder.execute("BEGIN IMMEDIATE")
+            async with engine.connect() as connection:
+                with pytest.raises(OperationalError, match="database is locked") as locked:
+                    await connection.exec_driver_sql("INSERT INTO lock_probe VALUES (1)")
+        [note] = locked.value.__notes__
+        waited = re.match(r"SQLite returned SQLITE_BUSY after (\d+\.\d\d) s\.", note)
+        assert waited is not None
+        assert float(waited.group(1)) >= 0.05
+    finally:
+        await engine.dispose()