From b38012cf97d95cfcf83303a40664e0272190edef Mon Sep 17 00:00:00 2001 From: varunj-msft Date: Tue, 1 Sep 2026 02:47:29 +0000 Subject: [PATCH] FIX Keep the backend responsive while starting scenario runs Scenario initialization loads the default datasets, which takes minutes of mostly synchronous work. Running it on the event loop wedged the backend, so every later scenario failed its health probe. Initialization now runs on a worker thread. - Offload initialization to a single-worker executor so the loop stays free, and serialize preparations because the in-memory backend shares one DBAPI connection. - Hold the concurrency permit until an abandoned preparation thread actually stops, including when it finishes as the cancellation lands, and terminalize the run it already stored instead of leaving it in CREATED. - Drain initialization's own teardown tasks before closing the throwaway loop. If anything outlives the drain, mark the run failed and refuse the start rather than returning a scenario that holds dead async resources. - Do not start a run that was cancelled while it was still initializing. - Serialize in-memory SQLite sessions so a preparation thread and a status poll cannot interleave on the shared connection and lose writes. --- pyrit/backend/routes/scenarios.py | 4 +- .../backend/services/scenario_run_service.py | 258 ++++++++++-- pyrit/memory/sqlite_memory.py | 42 +- .../unit/backend/test_scenario_run_service.py | 384 +++++++++++++++++- tests/unit/memory/test_sqlite_memory.py | 107 ++++- 5 files changed, 752 insertions(+), 43 deletions(-) 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..264a10f583 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,222 @@ 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, + # 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 resumed run is started with an id the caller already knows, so it can be + # cancelled while initialization is still on the worker thread. 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: + 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() - # 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.") + return response - # Track active task - active = _ActiveTask(scenario_result_id=scenario_result_id, scenario=scenario) - self._active_tasks[scenario_result_id] = active + def _release_abandoned_prepare(self, prepare_task: "asyncio.Future[Scenario]") -> None: + """ + Clean up after an abandoned preparation thread has finished. - # 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 + ``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``. - 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 + 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. + scenario_result_id = prepare_task.result()._scenario_result_id + if scenario_result_id: + try: + self._memory.update_scenario_run_state( + scenario_result_id=scenario_result_id, + 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. + scenario_result_id = scenario._scenario_result_id + if scenario_result_id: + try: + self._memory.update_scenario_run_state( + scenario_result_id=scenario_result_id, + 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. + + Raises: + RuntimeError: If any task is still running after the drain timeout. + """ + pending = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()] + if not pending: + return + + done, still_running = await asyncio.wait(pending, timeout=self._INITIALIZATION_DRAIN_TIMEOUT) + 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()}") + + if still_running: + 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 still_running))}" + ) + + 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: """ 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..1ce3951bf5 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,362 @@ 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, "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 + + 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, "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 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 resumed run can be cancelled through its known id before initialization finishes.""" + 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_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, "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" + + 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, "update_scenario_run_state", side_effect=ValueError("gone")): + with pytest.raises(RuntimeError, match="left background tasks"): + service._prepare_run_blocking(request=_make_request()) + + 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.""" diff --git a/tests/unit/memory/test_sqlite_memory.py b/tests/unit/memory/test_sqlite_memory.py index ccdb97bd12..86a51ebf1e 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,104 @@ 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