diff --git a/pyrit/backend/routes/scenarios.py b/pyrit/backend/routes/scenarios.py index a6e9ea5e00..698fb436a0 100644 --- a/pyrit/backend/routes/scenarios.py +++ b/pyrit/backend/routes/scenarios.py @@ -146,7 +146,9 @@ async def start_scenario_run(request: RunScenarioRequest) -> ScenarioRunSummary: """ Start a new scenario run as a background task. - Returns immediately with a scenario_result_id that can be polled for status. + Initialization runs eagerly so configuration errors surface here, then the run + itself continues in the background. Returns a scenario_result_id that can be + polled for status. Args: request: Scenario run configuration. diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index c635dab06f..b7a9084278 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -11,10 +11,12 @@ import asyncio import base64 import contextlib +import functools import json import logging import uuid from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from datetime import datetime, timezone from typing import Any @@ -162,6 +164,11 @@ class ScenarioRunService: Keeps an in-memory dict only for active asyncio tasks (cancellation support). """ + #: Seconds to let initialization's own background tasks (for example HTTP client teardown + #: scheduled from ``__del__``) finish before the initialization loop is torn down. This is + #: headroom for incidental teardown, not a waiter for real long-running work. + _INITIALIZATION_DRAIN_TIMEOUT = 5.0 + def __init__(self, *, max_concurrent_runs: int = _DEFAULT_MAX_CONCURRENT_RUNS) -> None: """Initialize the scenario run service.""" self._max_concurrent_runs = max_concurrent_runs @@ -170,6 +177,13 @@ def __init__(self, *, max_concurrent_runs: int = _DEFAULT_MAX_CONCURRENT_RUNS) - self._run_semaphore = asyncio.Semaphore(max_concurrent_runs) self._configuration_resolver = ScenarioConfigurationResolver() + # Initialization writes to CentralMemory, and the in-memory SQLite backend shares one + # DBAPI connection across every thread (StaticPool, sqlite_memory.py). Two preparations + # running at once would use that connection concurrently and lose or corrupt writes, so + # they are serialized onto a single worker. The event loop is still free while they run, + # which is the point of the offload. + self._prepare_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="pyrit-scenario-prep") + async def start_run_async(self, *, request: RunScenarioRequest) -> ScenarioRunSummary: """ Start a new scenario run as a background task. @@ -197,46 +211,259 @@ async def start_run_async(self, *, request: RunScenarioRequest) -> ScenarioRunSu await self._run_semaphore.acquire() - # Perform all initialization eagerly — errors propagate to caller + # This frame owns the permit until the background task is created; every exit path + # before that hand-off has to release it, including cancellation, which is a + # BaseException and so is not caught by ``except Exception``. + release_on_exit = True + registered_run_id: str | None = None try: - scenario_class = self._configuration_resolver.resolve_scenario_class(scenario_name=request.scenario_name) - await self._run_initializers_async(request=request) - objective_target = self._configuration_resolver.resolve_target(target_name=request.target_name) - init_kwargs = self._configuration_resolver.resolve_configuration( - scenario_name=request.scenario_name, - scenario_class=scenario_class, - objective_target=objective_target, - techniques=request.techniques, - dataset_names=request.dataset_names, - max_dataset_size=request.max_dataset_size, - dataset_filters=request.dataset_filters, - include_baseline=request.include_baseline, - max_concurrency=request.max_concurrency, - max_retries=request.max_retries, - memory_labels=request.labels, + # A resumed run keeps the state its previous run left behind, so one that was + # cancelled and is now being resumed on purpose is still CANCELLED while it + # initializes. Read that before preparation: the check afterwards otherwise + # cannot tell an intentional resume from a cancellation that landed while the + # worker thread was still initializing, and would refuse to restart it. + resumed_from_cancelled = self._is_run_cancelled(scenario_result_id=request.scenario_result_id) + + # Initialization loads the default datasets, which takes minutes, and is mostly + # synchronous work. Run it on a worker thread so the event loop stays free to + # answer health checks and status polls while a run is starting. + prepare_task = asyncio.get_running_loop().run_in_executor( + self._prepare_executor, functools.partial(self._prepare_run_blocking, request=request) ) - scenario = await self._initialize_scenario_async(request=request, init_kwargs=init_kwargs) - except Exception: - self._run_semaphore.release() - raise + try: + scenario = await asyncio.shield(prepare_task) + except BaseException as exc: + # A worker thread cannot be killed, so it keeps initializing after this frame + # unwinds. Keep holding the permit until it actually finishes, otherwise the + # next caller is admitted while this run is still loading datasets and + # ``max_concurrent_runs`` stops bounding the work that is really running. + if not prepare_task.done(): + prepare_task.add_done_callback(self._release_abandoned_prepare) + release_on_exit = False + elif isinstance(exc, asyncio.CancelledError): + # The thread can finish just as the cancellation lands. A done future never + # calls back, so cleaning up here is the only chance to release the permit + # and terminalize the run that initialization already stored. + release_on_exit = False + try: + self._release_abandoned_prepare(prepare_task) + except Exception as cleanup_error: + # The permit is released first, so it is already back even if the rest + # failed. Never let cleanup replace the cancellation being propagated. + logger.warning(f"Could not clean up after a cancelled scenario preparation: {cleanup_error}") + raise + + # scenario_result_id is set during initialize_async + scenario_result_id = scenario._scenario_result_id + if scenario_result_id is None: + raise ValueError("Scenario did not produce a scenario_result_id during initialization.") + + # Track active task + active = _ActiveTask(scenario_result_id=scenario_result_id, scenario=scenario) + self._active_tasks[scenario_result_id] = active + registered_run_id = scenario_result_id + + # Build the response before spawning the task so that a failure here cannot leave + # a run executing that the caller never received an id for. + response = self.get_run(scenario_result_id=scenario_result_id) + if response is None: + raise RuntimeError( + f"Scenario run {scenario_result_id} was not found in the database after initialization." + ) + + # A run can be cancelled through its id while initialization is still on the worker + # thread: a resume already knows the id, and a fresh run appears in the run list as + # soon as initialization stores it. Nothing has run yet, so honour that instead of + # starting a scenario the caller gave up on. The finally block returns the permit + # and drops the tracking entry. + if response.status == ScenarioRunState.CANCELLED and not resumed_from_cancelled: + logger.info(f"Scenario run {scenario_result_id} was cancelled while it was being initialized.") + return response + + # Spawn background task (only runs scenario.run_async). It releases the permit in + # its own finally, so ownership transfers here and this frame must not release it. + task = asyncio.create_task(self._execute_run_async(scenario_result_id=scenario_result_id)) + active.task = task + release_on_exit = False + registered_run_id = None + finally: + if registered_run_id is not None: + self._active_tasks.pop(registered_run_id, None) + if release_on_exit: + self._run_semaphore.release() + + return response - # scenario_result_id is set during initialize_async - scenario_result_id = scenario._scenario_result_id - if scenario_result_id is None: - raise ValueError("Scenario did not produce a scenario_result_id during initialization.") + def _is_run_cancelled(self, *, scenario_result_id: str | None) -> bool: + """ + Report whether a stored run is already in the CANCELLED state. - # Track active task - active = _ActiveTask(scenario_result_id=scenario_result_id, scenario=scenario) - self._active_tasks[scenario_result_id] = active + Reads the header only. A resumed run can have thousands of linked attack results and + this runs on the event loop, which the rest of this path works to keep free. - # Spawn background task (only runs scenario.run_async) - task = asyncio.create_task(self._execute_run_async(scenario_result_id=scenario_result_id)) - active.task = task + Args: + scenario_result_id: The run being resumed, or None for a fresh run. - response = self.get_run(scenario_result_id=scenario_result_id) - if response is None: - raise RuntimeError(f"Scenario run {scenario_result_id} was not found in the database after initialization.") - return response + Returns: + bool: True when a stored run with this id is CANCELLED. + """ + if not scenario_result_id: + return False + stored = self._memory.get_scenario_result_header(scenario_result_id=scenario_result_id) + return stored is not None and stored.scenario_run_state == ScenarioRunState.CANCELLED + + def _release_abandoned_prepare(self, prepare_task: "asyncio.Future[Scenario]") -> None: + """ + Clean up after an abandoned preparation thread has finished. + + ``start_run_async`` hands ownership of the permit to this callback when it is + cancelled while the worker thread is still initializing, so the permit is only + released after the thread has genuinely stopped using the slot. A preparation that + succeeds anyway leaves behind a scenario result nobody will run, which is marked + cancelled here rather than left waiting in ``CREATED``. + + Args: + prepare_task: The future wrapping the abandoned ``_prepare_run_blocking`` call. + """ + self._run_semaphore.release() + + if prepare_task.cancelled(): + return + error = prepare_task.exception() + if error is not None: + logger.warning(f"Abandoned scenario preparation failed after the request was cancelled: {error}") + return + + # Initialization already stored a CREATED scenario result, and nothing is going to run + # it now, so terminalize it rather than leaving a run that never starts. A run that + # already reached a terminal state keeps it, so a real failure is not relabelled. + scenario_result_id = prepare_task.result()._scenario_result_id + if scenario_result_id: + try: + self._memory.try_update_scenario_run_state( + scenario_result_id=scenario_result_id, + expected_states={ScenarioRunState.CREATED, ScenarioRunState.IN_PROGRESS}, + scenario_run_state=ScenarioRunState.CANCELLED, + error_message="The start request was cancelled while the scenario was being initialized.", + ) + except Exception as update_error: + logger.warning( + f"Could not mark abandoned scenario run {scenario_result_id} as cancelled: {update_error}" + ) + logger.warning("Abandoned scenario preparation completed after the request was cancelled.") + + def _prepare_run_blocking(self, *, request: RunScenarioRequest) -> Scenario: + """ + Run the eager initialization for a scenario run on the calling thread. + + Exists so ``start_run_async`` can offload initialization onto a worker thread. + The scenario is executed later on the caller's event loop, so initialization must not + leave anything bound to the throwaway loop used here. Clients that schedule their own + teardown are given a moment to finish; anything still running after that would be + cancelled when the loop closes, so the start fails rather than handing back a scenario + that holds dead async resources. + + Args: + request: The run request with scenario name, target, and options. + + Returns: + Scenario: The initialized scenario. + + Raises: + RuntimeError: If tasks are still running on the initialization loop after the drain. + """ + + async def prepare_async() -> Scenario: + scenario = await self._prepare_run_async(request=request) + try: + await self._drain_initialization_tasks_async() + except RuntimeError as drain_error: + # Initialization already stored a CREATED row and this start is over, so + # terminalize it here rather than leaving a run that never begins. A cancel + # can land while the drain is running, so keep whatever terminal state won. + scenario_result_id = scenario._scenario_result_id + if scenario_result_id: + try: + self._memory.try_update_scenario_run_state( + scenario_result_id=scenario_result_id, + expected_states={ScenarioRunState.CREATED, ScenarioRunState.IN_PROGRESS}, + scenario_run_state=ScenarioRunState.FAILED, + error_message=str(drain_error), + error_type=type(drain_error).__name__, + ) + except Exception as update_error: + logger.warning(f"Could not mark scenario run {scenario_result_id} as failed: {update_error}") + raise + return scenario + + return asyncio.run(prepare_async()) + + async def _drain_initialization_tasks_async(self) -> None: + """ + Let initialization's background tasks finish before the initialization loop closes. + + Initialization builds throwaway async clients, and some of them schedule their own + teardown from ``__del__``, so a task can appear purely because a garbage collection + landed late. Waiting for those is the difference between a scenario that starts and + one that fails at random. A draining task can also start another one, so the set is + rebuilt after every wait and the whole drain shares a single deadline. + + Raises: + RuntimeError: If any task is still running after the drain timeout. + """ + loop = asyncio.get_running_loop() + current_task = asyncio.current_task() + deadline = loop.time() + self._INITIALIZATION_DRAIN_TIMEOUT + + while True: + pending = [task for task in asyncio.all_tasks() if task is not current_task] + if not pending: + return + + remaining = deadline - loop.time() + if remaining <= 0: + raise RuntimeError( + "Scenario initialization left background tasks on the initialization loop, which is " + "about to close. They would be cancelled and the scenario would hold dead async " + f"resources: {', '.join(sorted(task.get_name() for task in pending))}" + ) + + done, _ = await asyncio.wait(pending, timeout=remaining) + for task in done: + # Retrieve outcomes so a failed teardown task does not log "never retrieved" noise. + if not task.cancelled() and task.exception() is not None: + logger.debug(f"A scenario initialization task failed during teardown: {task.exception()}") + + async def _prepare_run_async(self, *, request: RunScenarioRequest) -> Scenario: + """ + Resolve and initialize the scenario for a run request. + + Args: + request: The run request with scenario name, target, and options. + + Returns: + Scenario: The initialized scenario. + + Raises: + ValueError: If scenario, target, initializer, or technique cannot be found. + """ + scenario_class = self._configuration_resolver.resolve_scenario_class(scenario_name=request.scenario_name) + await self._run_initializers_async(request=request) + objective_target = self._configuration_resolver.resolve_target(target_name=request.target_name) + init_kwargs = self._configuration_resolver.resolve_configuration( + scenario_name=request.scenario_name, + scenario_class=scenario_class, + objective_target=objective_target, + techniques=request.techniques, + dataset_names=request.dataset_names, + max_dataset_size=request.max_dataset_size, + dataset_filters=request.dataset_filters, + include_baseline=request.include_baseline, + max_concurrency=request.max_concurrency, + max_retries=request.max_retries, + memory_labels=request.labels, + ) + return await self._initialize_scenario_async(request=request, init_kwargs=init_kwargs) def get_run(self, *, scenario_result_id: str) -> ScenarioRunSummary | None: """ @@ -353,9 +580,11 @@ async def cancel_run_async(self, *, scenario_result_id: str) -> ScenarioRunSumma with contextlib.suppress(asyncio.CancelledError, asyncio.TimeoutError): await asyncio.wait_for(active.task, timeout=5.0) - # Persist cancelled state to DB - self._memory.update_scenario_run_state( + # The run can reach a terminal state during the await above, so only cancel a run that + # is still going. The re-read below reports whichever state actually won. + self._memory.try_update_scenario_run_state( scenario_result_id=scenario_result_id, + expected_states={ScenarioRunState.CREATED, ScenarioRunState.IN_PROGRESS}, scenario_run_state=ScenarioRunState.CANCELLED, error_message="Run was cancelled by user", error_type="CancelledError", diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 758c03bdcd..9943dfa7a7 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -10,7 +10,7 @@ import re import uuid import weakref -from collections.abc import Iterator, Mapping, MutableSequence, Sequence +from collections.abc import Collection, Iterator, Mapping, MutableSequence, Sequence from contextlib import closing from dataclasses import dataclass from datetime import datetime, timezone @@ -3886,6 +3886,66 @@ def update_scenario_run_state( logger.info(f"Updated scenario {scenario_result_id} state to '{scenario_run_state.value}'") + def try_update_scenario_run_state( + self, + *, + scenario_result_id: str, + expected_states: Collection[ScenarioRunState], + scenario_run_state: ScenarioRunState, + error_message: str | None = None, + error_type: str | None = None, + ) -> bool: + """ + Update the run state only when the stored state is one of ``expected_states``. + + The compare and the write are a single UPDATE so a run that reached a terminal state + on another thread is not overwritten. A read followed by + ``update_scenario_run_state`` cannot give that guarantee because scenario + preparation and cancellation run on different threads. + + Args: + scenario_result_id (str): The ID of the scenario result to update. + expected_states (Collection[ScenarioRunState]): States the row may currently be in. + scenario_run_state (ScenarioRunState): The new state for the scenario. + error_message (str | None): Optional scenario-level error message. + error_type (str | None): Optional exception class name. + + Returns: + bool: True if the row was updated, False if it was missing or in another state. + + Raises: + ValueError: If ``expected_states`` is empty. + """ + if not expected_states: + raise ValueError("expected_states must not be empty") + + values: dict[Any, Any] = { + "scenario_run_state": scenario_run_state.value, + "error_message": error_message, + "error_type": error_type, + } + if scenario_run_state in ( + ScenarioRunState.COMPLETED, + ScenarioRunState.FAILED, + ScenarioRunState.CANCELLED, + ): + values["completion_time"] = datetime.now(tz=timezone.utc) + + with closing(self.get_session()) as session: + updated_rows = ( + session.query(ScenarioResultEntry) + .filter( + ScenarioResultEntry.id == scenario_result_id, + ScenarioResultEntry.scenario_run_state.in_([state.value for state in expected_states]), + ) + .update(values, synchronize_session=False) + ) + session.commit() + + if updated_rows: + logger.info(f"Updated scenario {scenario_result_id} state to '{scenario_run_state.value}'") + return bool(updated_rows) + def update_scenario_metadata( self, *, diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index 210c3f1da6..d6beae537b 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -2,6 +2,8 @@ # Licensed under the MIT license. import logging +import threading +import weakref from collections.abc import Sequence from contextlib import closing from datetime import datetime @@ -72,6 +74,11 @@ def __init__( self.db_path = Path(db_path or Path(DB_DATA_PATH, self.DEFAULT_DB_FILE_NAME)).resolve() self.results_path = str(DB_DATA_PATH) + # An in-memory database shares a single DBAPI connection across every thread (see + # ``_create_engine``), so concurrent sessions would interleave on it. Serialize session + # lifetimes for that backend only; file-backed databases get a connection per checkout. + self._connection_lock: threading.RLock | None = threading.RLock() if self.db_path == ":memory:" else None + self.engine = self._create_engine(has_echo=verbose) self.SessionFactory = sessionmaker(bind=self.engine) if not skip_schema_migration: @@ -285,10 +292,43 @@ def get_session(self) -> Session: """ Provide a SQLAlchemy session for transactional operations. + For an in-memory database every session borrows the same DBAPI connection, so the + session is handed out under a lock that is only released when it is closed. That keeps + a whole transaction, not just a single statement, isolated from the other threads. + Returns: Session: A SQLAlchemy session bound to the engine. """ - return self.SessionFactory() + session = self.SessionFactory() + connection_lock = self._connection_lock + if connection_lock is None: + return session + + connection_lock.acquire() + close_session = session.close + released = False + + def release_once() -> None: + # Also runs if the session is discarded without being closed, so one caller that + # forgets cannot leave the lock held and stall every other thread forever. + nonlocal released + if released: + return + released = True + try: + connection_lock.release() + except RuntimeError: + logger.warning("An in-memory session was discarded by a thread that did not open it.") + + def close_and_release() -> None: + try: + close_session() + finally: + release_once() + + session.close = close_and_release # type: ignore[ty:invalid-assignment] + weakref.finalize(session, release_once) + return session def print_schema(self) -> None: """ diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index 83eb1b0995..469268be39 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -6,6 +6,9 @@ """ import asyncio +import logging +import threading +import time import uuid from datetime import datetime, timezone from typing import Any @@ -548,6 +551,16 @@ async def test_start_run_exceeds_concurrent_limit(self, mock_all_registries) -> scenario_instance = mock_all_registries["scenario_instance"] mock_sr = mock_all_registries["scenario_registry"] + # A real run holds its permit until it finishes, so the background task has to stay + # in flight for the limit to be reachable. The default AsyncMock returns immediately + # and would hand every permit straight back. + still_running = asyncio.Event() + + async def _block_until_released() -> None: + await still_running.wait() + + scenario_instance.run_async = _block_until_released + # Each call needs a unique scenario_result_id call_count = 0 @@ -559,13 +572,16 @@ async def _set_unique_id(*args: object, **kwargs: object) -> object: mock_sr.create_and_initialize_async = AsyncMock(side_effect=_set_unique_id) - # Fill up to the limit - for _ in range(_DEFAULT_MAX_CONCURRENT_RUNS): - await service.start_run_async(request=_make_request()) + try: + # Fill up to the limit + for _ in range(_DEFAULT_MAX_CONCURRENT_RUNS): + await service.start_run_async(request=_make_request()) - # Next one should fail - with pytest.raises(ValueError, match="Maximum concurrent runs"): - await service.start_run_async(request=_make_request()) + # Next one should fail + with pytest.raises(ValueError, match="Maximum concurrent runs"): + await service.start_run_async(request=_make_request()) + finally: + still_running.set() async def test_start_run_runs_initializers(self, mock_all_registries) -> None: """Test that initializers are run during start_run_async.""" @@ -603,6 +619,500 @@ async def test_start_run_omits_scenario_result_id_when_none(self, mock_all_regis assert call.args[0] == "foundry.red_team_agent" assert call.kwargs["scenario_result_id"] is None + async def test_start_run_keeps_event_loop_responsive(self, mock_all_registries) -> None: + """Initialization is offloaded, so the loop keeps running while a run starts.""" + service = ScenarioRunService() + + def _slow_prepare(*, request: Any) -> Any: + time.sleep(0.5) + return mock_all_registries["scenario_instance"] + + beats = 0 + + async def _heartbeat() -> None: + nonlocal beats + while True: + await asyncio.sleep(0.01) + beats += 1 + + with patch.object(service, "_prepare_run_blocking", _slow_prepare): + heartbeat = asyncio.create_task(_heartbeat()) + try: + await service.start_run_async(request=_make_request()) + finally: + heartbeat.cancel() + + # A blocked event loop yields zero heartbeats over the same window. + assert beats > 10 + + async def test_start_run_background_task_survives_handoff(self, mock_all_registries) -> None: + """The background task must outlive start_run_async and actually execute the run.""" + service = ScenarioRunService() + scenario_instance = mock_all_registries["scenario_instance"] + executed = asyncio.Event() + + async def _run_async() -> None: + executed.set() + + scenario_instance.run_async = _run_async + + response = await service.start_run_async(request=_make_request()) + + await asyncio.wait_for(executed.wait(), timeout=5) + assert executed.is_set() + assert response.status == ScenarioRunState.IN_PROGRESS + + async def test_start_run_marks_abandoned_prepare_cancelled(self, mock_all_registries) -> None: + """A preparation that finishes after its caller left must not sit in CREATED forever.""" + service = ScenarioRunService() + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._scenario_result_id = "abandoned-id" + finished = threading.Event() + + def _slow_prepare(*, request: Any) -> Any: + time.sleep(0.5) + finished.set() + return scenario_instance + + with patch.object(service, "_prepare_run_blocking", _slow_prepare): + with patch.object(service._memory, "try_update_scenario_run_state") as update_state: + task = asyncio.create_task(service.start_run_async(request=_make_request())) + await asyncio.sleep(0.1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + await asyncio.sleep(1.0) + assert finished.is_set() + update_state.assert_called_once() + assert update_state.call_args.kwargs["scenario_result_id"] == "abandoned-id" + assert update_state.call_args.kwargs["scenario_run_state"] == ScenarioRunState.CANCELLED + assert update_state.call_args.kwargs["expected_states"] == { + ScenarioRunState.CREATED, + ScenarioRunState.IN_PROGRESS, + } + + async def test_start_run_marks_prepare_cancelled_when_it_finishes_before_cancellation_lands( + self, mock_all_registries + ) -> None: + """A done future never calls back, so this race used to leave the run stuck in CREATED.""" + service = ScenarioRunService() + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._scenario_result_id = "raced-id" + + def _instant_prepare(*, request: Any) -> Any: + return scenario_instance + + async def _complete_then_cancel(awaitable): + # The preparation finishes, then the cancellation lands: the exact ordering that + # leaves ``prepare_task.done()`` True inside the handler. + await awaitable + raise asyncio.CancelledError + + with patch.object(service, "_prepare_run_blocking", _instant_prepare): + with patch("asyncio.shield", _complete_then_cancel): + with patch.object(service._memory, "try_update_scenario_run_state") as update_state: + with pytest.raises(asyncio.CancelledError): + await service.start_run_async(request=_make_request()) + + update_state.assert_called_once() + assert update_state.call_args.kwargs["scenario_result_id"] == "raced-id" + assert update_state.call_args.kwargs["scenario_run_state"] == ScenarioRunState.CANCELLED + assert update_state.call_args.kwargs["expected_states"] == { + ScenarioRunState.CREATED, + ScenarioRunState.IN_PROGRESS, + } + assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS + + async def test_start_run_does_not_run_a_scenario_cancelled_during_initialization(self, mock_all_registries) -> None: + """A run appears in the run list as soon as it is stored, so it can be cancelled mid-init.""" + service = ScenarioRunService() + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._scenario_result_id = "cancelled-during-init" + + def _prepare(*, request: Any) -> Any: + return scenario_instance + + cancelled = _make_db_scenario_result(result_id="cancelled-during-init", run_state=ScenarioRunState.CANCELLED) + mock_all_registries["memory"].get_scenario_results.return_value = [cancelled] + with patch.object(service, "_prepare_run_blocking", _prepare): + with patch.object(service, "_execute_run_async") as execute: + response = await service.start_run_async(request=_make_request()) + + assert response.status == ScenarioRunState.CANCELLED + execute.assert_not_called() + assert "cancelled-during-init" not in service._active_tasks + assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS + + async def test_start_run_resumes_a_run_that_was_already_cancelled(self, mock_all_registries) -> None: + """Resuming a cancelled run is deliberate, so its starting state must not look like a cancel.""" + service = ScenarioRunService() + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._scenario_result_id = "resumed-cancelled" + + def _prepare(*, request: Any) -> Any: + return scenario_instance + + cancelled = _make_db_scenario_result(result_id="resumed-cancelled", run_state=ScenarioRunState.CANCELLED) + mock_all_registries["memory"].get_scenario_results.return_value = [cancelled] + mock_all_registries["memory"].get_scenario_result_header.return_value = cancelled + + with patch.object(service, "_prepare_run_blocking", _prepare): + with patch.object(service, "_execute_run_async") as execute: + response = await service.start_run_async(request=_make_request(scenario_result_id="resumed-cancelled")) + + execute.assert_called_once() + assert "resumed-cancelled" in service._active_tasks + # The stored row is still CANCELLED until the scenario moves it on; what matters is + # that the resume was not mistaken for a cancellation and actually started. + assert response.scenario_result_id == "resumed-cancelled" + + async def test_start_run_honours_a_cancel_that_lands_while_a_live_run_initializes( + self, mock_all_registries + ) -> None: + """A run that was not already cancelled must still respect a cancel during preparation.""" + service = ScenarioRunService() + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._scenario_result_id = "cancelled-mid-init" + + def _prepare(*, request: Any) -> Any: + return scenario_instance + + cancelled = _make_db_scenario_result(result_id="cancelled-mid-init", run_state=ScenarioRunState.CANCELLED) + in_progress = _make_db_scenario_result(result_id="cancelled-mid-init", run_state=ScenarioRunState.IN_PROGRESS) + mock_all_registries["memory"].get_scenario_results.return_value = [cancelled] + # The pre-preparation read sees a live run; the cancel lands while the worker prepares. + mock_all_registries["memory"].get_scenario_result_header.return_value = in_progress + + with patch.object(service, "_prepare_run_blocking", _prepare): + with patch.object(service, "_execute_run_async") as execute: + response = await service.start_run_async(request=_make_request(scenario_result_id="cancelled-mid-init")) + + assert response.status == ScenarioRunState.CANCELLED + execute.assert_not_called() + assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS + + async def test_start_run_does_not_read_a_header_for_a_fresh_run(self, mock_all_registries) -> None: + """A fresh run has no stored state, so it must not pay for an extra query.""" + service = ScenarioRunService() + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._scenario_result_id = "fresh-id" + + def _prepare(*, request: Any) -> Any: + return scenario_instance + + with patch.object(service, "_prepare_run_blocking", _prepare): + with patch.object(service, "_execute_run_async"): + await service.start_run_async(request=_make_request()) + + mock_all_registries["memory"].get_scenario_result_header.assert_not_called() + + async def test_start_run_failure_does_not_report_a_cancellation(self, mock_all_registries, caplog) -> None: + service = ScenarioRunService() + + def _failing_prepare(*, request: Any) -> Any: + raise ValueError("Scenario 'nope' not found") + + with patch.object(service, "_prepare_run_blocking", _failing_prepare): + with caplog.at_level(logging.WARNING): + with pytest.raises(ValueError, match="not found"): + await service.start_run_async(request=_make_request()) + + assert "cancelled" not in caplog.text.lower() + assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS + + async def test_start_run_cleanup_failure_still_propagates_cancellation(self, mock_all_registries) -> None: + """Cleanup runs inline on this path, so it must not replace the CancelledError.""" + service = ScenarioRunService() + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._scenario_result_id = "raced-id" + + def _instant_prepare(*, request: Any) -> Any: + return scenario_instance + + async def _complete_then_cancel(awaitable): + await awaitable + raise asyncio.CancelledError + + with patch.object(service, "_prepare_run_blocking", _instant_prepare): + with patch("asyncio.shield", _complete_then_cancel): + with patch.object(service, "_release_abandoned_prepare", side_effect=RuntimeError("cleanup exploded")): + with pytest.raises(asyncio.CancelledError): + await service.start_run_async(request=_make_request()) + + def test_prepare_executor_serializes_preparations(self, mock_all_registries) -> None: + """In-memory SQLite shares one connection across threads, so preparations must not overlap.""" + service = ScenarioRunService() + overlap = [] + active = 0 + lock = threading.Lock() + + def _prepare(*, request: Any) -> Any: + nonlocal active + with lock: + active += 1 + overlap.append(active) + time.sleep(0.05) + with lock: + active -= 1 + return mock_all_registries["scenario_instance"] + + with patch.object(service, "_prepare_run_blocking", _prepare): + futures = [ + service._prepare_executor.submit(lambda: service._prepare_run_blocking(request=_make_request())) + for _ in range(4) + ] + for future in futures: + future.result() + + assert max(overlap) == 1 + + def test_prepare_run_blocking_waits_for_initialization_teardown_tasks(self, mock_all_registries) -> None: + """Async clients schedule their own teardown, so a benign task must not fail the start.""" + service = ScenarioRunService() + + async def _prepare_with_teardown(*, request: Any) -> Any: + task = asyncio.create_task(asyncio.sleep(0.05)) + task.set_name("client-teardown-task") + return mock_all_registries["scenario_instance"] + + with patch.object(service, "_prepare_run_async", _prepare_with_teardown): + assert service._prepare_run_blocking(request=_make_request()) is mock_all_registries["scenario_instance"] + + def test_prepare_run_blocking_fails_when_a_task_outlives_the_drain(self, mock_all_registries) -> None: + """A task still running when the loop closes is cancelled, so the scenario is unusable.""" + service = ScenarioRunService() + + async def _leaky_prepare(*, request: Any) -> Any: + task = asyncio.create_task(asyncio.sleep(3600)) + task.set_name("stray-initializer-task") + await asyncio.sleep(0) + return mock_all_registries["scenario_instance"] + + with patch.object(service, "_prepare_run_async", _leaky_prepare): + with patch.object(ScenarioRunService, "_INITIALIZATION_DRAIN_TIMEOUT", 0.05): + with pytest.raises(RuntimeError, match="left background tasks on the initialization loop") as exc_info: + service._prepare_run_blocking(request=_make_request()) + + assert "stray-initializer-task" in str(exc_info.value) + + def test_prepare_run_blocking_marks_the_run_failed_when_the_drain_fails(self, mock_all_registries) -> None: + """Initialization already stored the run, so a failed drain must not leave it in CREATED.""" + service = ScenarioRunService() + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._scenario_result_id = "drained-id" + + async def _leaky_prepare(*, request: Any) -> Any: + task = asyncio.create_task(asyncio.sleep(3600)) + task.set_name("stray-initializer-task") + await asyncio.sleep(0) + return scenario_instance + + with patch.object(service, "_prepare_run_async", _leaky_prepare): + with patch.object(ScenarioRunService, "_INITIALIZATION_DRAIN_TIMEOUT", 0.05): + with patch.object(service._memory, "try_update_scenario_run_state") as update_state: + with pytest.raises(RuntimeError, match="left background tasks"): + service._prepare_run_blocking(request=_make_request()) + + update_state.assert_called_once() + assert update_state.call_args.kwargs["scenario_result_id"] == "drained-id" + assert update_state.call_args.kwargs["scenario_run_state"] == ScenarioRunState.FAILED + assert update_state.call_args.kwargs["error_type"] == "RuntimeError" + assert update_state.call_args.kwargs["expected_states"] == { + ScenarioRunState.CREATED, + ScenarioRunState.IN_PROGRESS, + } + + def test_prepare_run_blocking_reports_the_drain_error_when_marking_failed_fails(self, mock_all_registries) -> None: + """A bookkeeping failure must not replace the error that explains the failed start.""" + service = ScenarioRunService() + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._scenario_result_id = "drained-id" + + async def _leaky_prepare(*, request: Any) -> Any: + task = asyncio.create_task(asyncio.sleep(3600)) + task.set_name("stray-initializer-task") + await asyncio.sleep(0) + return scenario_instance + + with patch.object(service, "_prepare_run_async", _leaky_prepare): + with patch.object(ScenarioRunService, "_INITIALIZATION_DRAIN_TIMEOUT", 0.05): + with patch.object(service._memory, "try_update_scenario_run_state", side_effect=ValueError("gone")): + with pytest.raises(RuntimeError, match="left background tasks"): + service._prepare_run_blocking(request=_make_request()) + + def test_prepare_run_blocking_waits_for_a_task_spawned_during_the_drain(self, mock_all_registries) -> None: + """A draining task can start another one, which the first snapshot never saw.""" + service = ScenarioRunService() + scenario_instance = mock_all_registries["scenario_instance"] + child_finished = threading.Event() + + async def _prepare_spawning_a_child(*, request: Any) -> Any: + async def _child() -> None: + await asyncio.sleep(0.05) + child_finished.set() + + async def _parent() -> None: + await asyncio.sleep(0.01) + asyncio.create_task(_child(), name="child-teardown-task") + + asyncio.create_task(_parent(), name="parent-teardown-task") + await asyncio.sleep(0) + return scenario_instance + + with patch.object(service, "_prepare_run_async", _prepare_spawning_a_child): + assert service._prepare_run_blocking(request=_make_request()) is scenario_instance + + assert child_finished.is_set() + + def test_prepare_run_blocking_fails_when_a_spawned_child_outlives_the_drain(self, mock_all_registries) -> None: + """The child is the one that would be cancelled by the closing loop, so it must be named.""" + service = ScenarioRunService() + + async def _prepare_spawning_a_slow_child(*, request: Any) -> Any: + async def _parent() -> None: + await asyncio.sleep(0.01) + asyncio.create_task(asyncio.sleep(3600), name="child-teardown-task") + + asyncio.create_task(_parent(), name="parent-teardown-task") + await asyncio.sleep(0) + return mock_all_registries["scenario_instance"] + + with patch.object(service, "_prepare_run_async", _prepare_spawning_a_slow_child): + with patch.object(ScenarioRunService, "_INITIALIZATION_DRAIN_TIMEOUT", 0.3): + with pytest.raises(RuntimeError, match="child-teardown-task"): + service._prepare_run_blocking(request=_make_request()) + + def test_prepare_run_blocking_drain_uses_one_deadline_across_generations(self, mock_all_registries) -> None: + """Re-scanning must not restart the budget, or a chain of tasks could stall a start forever.""" + service = ScenarioRunService() + + async def _prepare_spawning_a_chain(*, request: Any) -> Any: + async def _link(depth: int) -> None: + await asyncio.sleep(0.05) + asyncio.create_task(_link(depth + 1), name=f"chain-task-{depth + 1}") + + asyncio.create_task(_link(0), name="chain-task-0") + await asyncio.sleep(0) + return mock_all_registries["scenario_instance"] + + started = time.monotonic() + with patch.object(service, "_prepare_run_async", _prepare_spawning_a_chain): + with patch.object(ScenarioRunService, "_INITIALIZATION_DRAIN_TIMEOUT", 0.3): + with pytest.raises(RuntimeError, match="chain-task"): + service._prepare_run_blocking(request=_make_request()) + + assert time.monotonic() - started < 3 + + async def test_start_run_releases_semaphore_when_initialization_leaks_a_task(self, mock_all_registries) -> None: + """Failing the preparation must not strand the permit it was holding.""" + service = ScenarioRunService() + + async def _leaky_prepare(*, request: Any) -> Any: + task = asyncio.create_task(asyncio.sleep(3600)) + task.set_name("stray-initializer-task") + await asyncio.sleep(0) + return mock_all_registries["scenario_instance"] + + with patch.object(service, "_prepare_run_async", _leaky_prepare): + with patch.object(ScenarioRunService, "_INITIALIZATION_DRAIN_TIMEOUT", 0.05): + with pytest.raises(RuntimeError, match="left background tasks"): + await service.start_run_async(request=_make_request()) + + assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS + + def test_prepare_run_blocking_is_quiet_when_initialization_is_self_contained( + self, mock_all_registries, caplog + ) -> None: + """The happy path must not warn, otherwise the signal is worthless.""" + service = ScenarioRunService() + + async def _clean_prepare(*, request: Any) -> Any: + return mock_all_registries["scenario_instance"] + + with patch.object(service, "_prepare_run_async", _clean_prepare): + with caplog.at_level(logging.WARNING): + service._prepare_run_blocking(request=_make_request()) + + assert "left background tasks" not in caplog.text + + async def test_start_run_holds_semaphore_until_abandoned_prepare_finishes(self, mock_all_registries) -> None: + """A cancelled start must not free capacity while its worker thread is still initializing.""" + service = ScenarioRunService() + finished = threading.Event() + + def _slow_prepare(*, request: Any) -> Any: + time.sleep(0.5) + finished.set() + return mock_all_registries["scenario_instance"] + + with patch.object(service, "_prepare_run_blocking", _slow_prepare): + task = asyncio.create_task(service.start_run_async(request=_make_request())) + await asyncio.sleep(0.1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + # The worker thread cannot be killed, so admitting another run here would let + # two initializations share one permit. + assert not finished.is_set() + assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS - 1 + + await asyncio.sleep(1.0) + assert finished.is_set() + assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS + + async def test_start_run_releases_semaphore_when_prepare_fails(self, mock_all_registries) -> None: + """CancelledError is a BaseException, so it needs an explicit release path.""" + service = ScenarioRunService() + + def _failing_prepare(*, request: Any) -> Any: + raise ValueError("boom") + + with patch.object(service, "_prepare_run_blocking", _failing_prepare): + with pytest.raises(ValueError, match="boom"): + await service.start_run_async(request=_make_request()) + + assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS + + async def test_start_run_releases_semaphore_when_result_id_missing(self, mock_all_registries) -> None: + """The missing scenario_result_id check used to sit outside the try block.""" + service = ScenarioRunService() + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._scenario_result_id = None + + with pytest.raises(ValueError, match="did not produce a scenario_result_id"): + await service.start_run_async(request=_make_request()) + + assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS + + async def test_start_run_cleans_up_when_response_lookup_fails(self, mock_all_registries) -> None: + """A response failure must not strand a permit or leave an active-task entry.""" + service = ScenarioRunService() + + with patch.object(service, "get_run", return_value=None): + with pytest.raises(RuntimeError, match="not found in the database"): + await service.start_run_async(request=_make_request()) + + assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS + assert service._active_tasks == {} + + async def test_start_run_releases_semaphore_exactly_once_on_success(self, mock_all_registries) -> None: + """The background task owns the permit after handoff, so it is not double-released.""" + service = ScenarioRunService() + released = asyncio.Event() + + async def _run_async() -> None: + released.set() + + mock_all_registries["scenario_instance"].run_async = _run_async + + await service.start_run_async(request=_make_request()) + await asyncio.wait_for(released.wait(), timeout=5) + await asyncio.sleep(0) + + assert service._run_semaphore._value == _DEFAULT_MAX_CONCURRENT_RUNS + class TestScenarioRunServiceGetRun: """Tests for ScenarioRunService.get_run.""" @@ -741,8 +1251,9 @@ async def test_cancel_run_sets_cancelled_status(self, mock_all_registries) -> No result = await service.cancel_run_async(scenario_result_id=response.scenario_result_id) - mock_memory.update_scenario_run_state.assert_called_once_with( + mock_memory.try_update_scenario_run_state.assert_called_once_with( scenario_result_id=response.scenario_result_id, + expected_states={ScenarioRunState.CREATED, ScenarioRunState.IN_PROGRESS}, scenario_run_state=ScenarioRunState.CANCELLED, error_message="Run was cancelled by user", error_type="CancelledError", diff --git a/tests/unit/memory/memory_interface/test_interface_scenario_results.py b/tests/unit/memory/memory_interface/test_interface_scenario_results.py index ad75bc73e7..eab0248c74 100644 --- a/tests/unit/memory/memory_interface/test_interface_scenario_results.py +++ b/tests/unit/memory/memory_interface/test_interface_scenario_results.py @@ -2,6 +2,7 @@ # Licensed under the MIT license. from datetime import datetime, timedelta, timezone +from uuid import uuid4 import pytest from unit.mocks import get_mock_scorer_identifier, make_scenario_result @@ -905,6 +906,85 @@ def test_update_scenario_run_state_updates_state_and_error_fields( assert hydrated.error_type is None +def test_try_update_scenario_run_state_updates_when_the_state_matches( + sqlite_instance: MemoryInterface, +): + """The compare-and-set applies when the row is still in one of the expected states.""" + scenario_result = create_scenario_result(name="CAS Match", attack_results={"a": []}) + sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) + sid = str(scenario_result.id) + sqlite_instance.update_scenario_run_state(scenario_result_id=sid, scenario_run_state=ScenarioRunState.CREATED) + + updated = sqlite_instance.try_update_scenario_run_state( + scenario_result_id=sid, + expected_states={ScenarioRunState.CREATED, ScenarioRunState.IN_PROGRESS}, + scenario_run_state=ScenarioRunState.FAILED, + error_message="boom", + error_type="RuntimeError", + ) + + assert updated is True + [hydrated] = sqlite_instance.get_scenario_results(scenario_result_ids=[sid]) + assert hydrated.scenario_run_state == ScenarioRunState.FAILED + assert hydrated.error_message == "boom" + assert hydrated.error_type == "RuntimeError" + + +def test_try_update_scenario_run_state_preserves_a_terminal_state( + sqlite_instance: MemoryInterface, +): + """A run cancelled while preparation was draining must keep its cancellation.""" + scenario_result = create_scenario_result(name="CAS Mismatch", attack_results={"a": []}) + sqlite_instance.add_scenario_results_to_memory(scenario_results=[scenario_result]) + sid = str(scenario_result.id) + sqlite_instance.update_scenario_run_state( + scenario_result_id=sid, + scenario_run_state=ScenarioRunState.CANCELLED, + error_message="Run was cancelled by user", + error_type="CancelledError", + ) + + updated = sqlite_instance.try_update_scenario_run_state( + scenario_result_id=sid, + expected_states={ScenarioRunState.CREATED, ScenarioRunState.IN_PROGRESS}, + scenario_run_state=ScenarioRunState.FAILED, + error_message="boom", + error_type="RuntimeError", + ) + + assert updated is False + [hydrated] = sqlite_instance.get_scenario_results(scenario_result_ids=[sid]) + assert hydrated.scenario_run_state == ScenarioRunState.CANCELLED + assert hydrated.error_message == "Run was cancelled by user" + assert hydrated.error_type == "CancelledError" + + +def test_try_update_scenario_run_state_reports_a_missing_row( + sqlite_instance: MemoryInterface, +): + """A run that is gone is not an error here; the caller only needs to know nothing changed.""" + assert ( + sqlite_instance.try_update_scenario_run_state( + scenario_result_id=str(uuid4()), + expected_states={ScenarioRunState.CREATED}, + scenario_run_state=ScenarioRunState.FAILED, + ) + is False + ) + + +def test_try_update_scenario_run_state_rejects_empty_expected_states( + sqlite_instance: MemoryInterface, +): + """An empty set would silently never match, which would hide the bug it exists to prevent.""" + with pytest.raises(ValueError, match="expected_states"): + sqlite_instance.try_update_scenario_run_state( + scenario_result_id=str(uuid4()), + expected_states=set(), + scenario_run_state=ScenarioRunState.FAILED, + ) + + def test_get_scenario_results_by_target_identifier_filter_hash( sqlite_instance: MemoryInterface, ): diff --git a/tests/unit/memory/test_sqlite_memory.py b/tests/unit/memory/test_sqlite_memory.py index ccdb97bd12..8af4c17614 100644 --- a/tests/unit/memory/test_sqlite_memory.py +++ b/tests/unit/memory/test_sqlite_memory.py @@ -1,12 +1,15 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import gc import io import logging import os import tempfile +import threading import uuid from collections.abc import Sequence +from contextlib import closing from unittest.mock import MagicMock import pytest @@ -16,10 +19,12 @@ from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.sql.sqltypes import NullType +from pyrit.common.singleton import Singleton from pyrit.converter.base64_converter import Base64Converter from pyrit.memory.alembic.versions.ab8f2c1a9d07_pre_alembic_release_schema import INITIAL_METADATA from pyrit.memory.memory_models import EmbeddingDataEntry, PromptMemoryEntry from pyrit.memory.migration import run_schema_migrations +from pyrit.memory.sqlite_memory import SQLiteMemory from pyrit.memory.storage.serializers import set_message_piece_sha256_async from pyrit.models import Conversation, MessagePiece, flatten_to_message_pieces from pyrit.prompt_target.text_target import TextTarget @@ -332,7 +337,6 @@ def test_reset_database_keeps_foreign_alembic_version_table(sqlite_instance): async def test_insert_entry(sqlite_instance): - session = sqlite_instance.get_session() message_piece_entry = MessagePiece( id=uuid.uuid4(), conversation_id="123", @@ -999,3 +1003,106 @@ def test_run_schema_migrations_no_memory_tables(): }.issubset(table_names) finally: engine.dispose() + + +@pytest.fixture +def isolated_memory_factory(): + """Build SQLiteMemory instances that are not the shared process-wide singleton.""" + saved = Singleton._instances.copy() + Singleton._instances.clear() + created = [] + + def _factory(**kwargs): + Singleton._instances.pop(SQLiteMemory, None) + memory = SQLiteMemory(**kwargs) + created.append(memory) + return memory + + try: + yield _factory + finally: + for memory in created: + memory.dispose_engine() + Singleton._instances.clear() + Singleton._instances.update(saved) + + +def test_in_memory_database_serializes_sessions_across_threads(isolated_memory_factory): + """ + An in-memory database shares one DBAPI connection, so overlapping sessions corrupt writes. + Without serialization this loses rows and raises sqlite3.InterfaceError. + """ + memory = isolated_memory_factory(db_path=":memory:") + with closing(memory.get_session()) as session: + session.execute(text("CREATE TABLE lock_probe (id INTEGER PRIMARY KEY, value TEXT)")) + session.commit() + + errors: list[str] = [] + + def _writer(worker: int) -> None: + try: + for index in range(30): + with closing(memory.get_session()) as session: + session.execute( + text("INSERT INTO lock_probe (value) VALUES (:value)"), + {"value": f"{worker}-{index}"}, + ) + session.commit() + except Exception as exc: # pragma: no cover - only runs when serialization breaks + errors.append(f"{type(exc).__name__}: {exc}") + + threads = [threading.Thread(target=_writer, args=(worker,)) for worker in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=30) + + assert not any(thread.is_alive() for thread in threads), "session lock deadlocked" + assert errors == [] + with closing(memory.get_session()) as session: + assert session.execute(text("SELECT COUNT(*) FROM lock_probe")).scalar() == 120 + + +def test_in_memory_database_allows_nested_sessions_on_one_thread(isolated_memory_factory): + """The lock is re-entrant so a caller that opens a second session cannot deadlock itself.""" + memory = isolated_memory_factory(db_path=":memory:") + with closing(memory.get_session()) as outer: + with closing(memory.get_session()) as inner: + assert inner.execute(text("SELECT 1")).scalar() == 1 + assert outer.execute(text("SELECT 1")).scalar() == 1 + + +def test_in_memory_session_close_is_idempotent(isolated_memory_factory): + """A double close must not release the lock twice and free it for another thread.""" + memory = isolated_memory_factory(db_path=":memory:") + session = memory.get_session() + session.close() + session.close() + + assert not memory._connection_lock._is_owned() + with closing(memory.get_session()) as session: + assert session.execute(text("SELECT 1")).scalar() == 1 + + +def test_in_memory_session_discarded_without_close_frees_the_lock(isolated_memory_factory): + """One caller that forgets to close must not stall every other thread forever.""" + memory = isolated_memory_factory(db_path=":memory:") + + def _leak_a_session() -> None: + memory.get_session() + + _leak_a_session() + gc.collect() + + assert not memory._connection_lock._is_owned() + with closing(memory.get_session()) as session: + assert session.execute(text("SELECT 1")).scalar() == 1 + + +def test_file_backed_database_is_not_serialized(isolated_memory_factory): + """File-backed databases get a connection per checkout, so they must not pay for the lock.""" + with tempfile.TemporaryDirectory() as temp_dir: + memory = isolated_memory_factory(db_path=os.path.join(temp_dir, "locking.db")) + assert memory._connection_lock is None + # Windows cannot remove the temp directory while the engine still holds the file open. + memory.dispose_engine()