diff --git a/doc/contributing/5_unit_tests.md b/doc/contributing/5_unit_tests.md index d798f75793..0be9e565c8 100644 --- a/doc/contributing/5_unit_tests.md +++ b/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. diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 671eb27cc8..544741a1c2 100644 --- a/tests/unit/conftest.py +++ b/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]: diff --git a/tests/unit/memory/memory_interface/test_interface_observations.py b/tests/unit/memory/memory_interface/test_interface_observations.py index f8d25e0548..deceb7e3ed 100644 --- a/tests/unit/memory/memory_interface/test_interface_observations.py +++ b/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) == [] diff --git a/tests/unit/memory/test_conversation_branching.py b/tests/unit/memory/test_conversation_branching.py index d201e4319e..c8ccd1b992 100644 --- a/tests/unit/memory/test_conversation_branching.py +++ b/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() == { diff --git a/tests/unit/memory/test_sqlite_fixtures.py b/tests/unit/memory/test_sqlite_fixtures.py index 0f67072a1f..2b073643da 100644 --- a/tests/unit/memory/test_sqlite_fixtures.py +++ b/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()