|
| 1 | +# Copyright (c) Microsoft Corporation. |
| 2 | +# Licensed under the MIT license. |
| 3 | + |
| 4 | +"""Bound native async analytics work through its actual resource cleanup.""" |
| 5 | + |
| 6 | +from __future__ import annotations |
| 7 | + |
| 8 | +import asyncio |
| 9 | +import logging |
| 10 | +from collections import deque |
| 11 | +from dataclasses import dataclass, field |
| 12 | +from sys import float_info |
| 13 | +from time import monotonic |
| 14 | +from typing import TYPE_CHECKING, TypeVar |
| 15 | + |
| 16 | +from pyrit.common.task_utils import gather_with_cleanup_async |
| 17 | +from pyrit.exceptions.analytics_exception import AnalyticsBusyException, AnalyticsTimeoutException |
| 18 | +from pyrit.memory.query_control import QueryControl |
| 19 | + |
| 20 | +if TYPE_CHECKING: |
| 21 | + from collections.abc import Awaitable, Callable, Coroutine |
| 22 | + |
| 23 | +logger = logging.getLogger(__name__) |
| 24 | +T = TypeVar("T") |
| 25 | + |
| 26 | + |
| 27 | +@dataclass(eq=False) |
| 28 | +class _Waiter: |
| 29 | + ready: asyncio.Future[None] |
| 30 | + deadline: float |
| 31 | + |
| 32 | + |
| 33 | +@dataclass(eq=False) |
| 34 | +class _Operation: |
| 35 | + """Keep the task strongly owned even after its requesting coroutine has left.""" |
| 36 | + |
| 37 | + control: QueryControl |
| 38 | + finished: asyncio.Future[None] |
| 39 | + abandoned: bool = False |
| 40 | + task: Awaitable[object] | None = None |
| 41 | + |
| 42 | + |
| 43 | +@dataclass |
| 44 | +class _Lane: |
| 45 | + limit: int |
| 46 | + timeout: float |
| 47 | + active: int = 0 |
| 48 | + queued: deque[_Waiter] = field(default_factory=deque) |
| 49 | + running: set[_Operation] = field(default_factory=set) |
| 50 | + |
| 51 | + |
| 52 | +class AnalyticsExecution: |
| 53 | + """ |
| 54 | + Reserve independent report and quick-query capacity on one event loop. |
| 55 | +
|
| 56 | + Admission is bounded before creating an operation task. A response deadline or |
| 57 | + caller cancellation signals ``QueryControl`` but does not cancel that task: |
| 58 | + native drivers and session cleanup must finish before its slot is released. |
| 59 | + There is no coalescing, result cache, thread pool, or SQL timeout policy here. |
| 60 | + The reader owns database interruption and session cleanup. |
| 61 | + """ |
| 62 | + |
| 63 | + def __init__( |
| 64 | + self, |
| 65 | + *, |
| 66 | + report_workers: int = 5, |
| 67 | + quick_workers: int = 2, |
| 68 | + max_queue: int = 10, |
| 69 | + queue_timeout: float = 1.0, |
| 70 | + report_timeout: float = 5.0, |
| 71 | + quick_timeout: float = 1.0, |
| 72 | + ) -> None: |
| 73 | + """ |
| 74 | + Bind a controller to the running loop without starting database work. |
| 75 | +
|
| 76 | + Args: |
| 77 | + report_workers (int): Concurrent report slots, including cleanup. |
| 78 | + quick_workers (int): Separate slots shared by result pages and facets. |
| 79 | + max_queue (int): Maximum waiting requests per lane, excluding active slots. |
| 80 | + queue_timeout (float): Maximum admission wait in seconds. |
| 81 | + report_timeout (float): Report execution budget after admission, in seconds. |
| 82 | + quick_timeout (float): Result-page/facet execution budget, in seconds. |
| 83 | +
|
| 84 | + Raises: |
| 85 | + ValueError: If a count is not a positive integer or a timeout is not positive and finite. |
| 86 | + RuntimeError: If construction occurs outside a running event loop. |
| 87 | + """ |
| 88 | + for name, value in ( |
| 89 | + ("report_workers", report_workers), |
| 90 | + ("quick_workers", quick_workers), |
| 91 | + ("max_queue", max_queue), |
| 92 | + ): |
| 93 | + if type(value) is not int or value <= 0: |
| 94 | + raise ValueError(f"{name} must be a positive integer.") |
| 95 | + for name, timeout in ( |
| 96 | + ("queue_timeout", queue_timeout), |
| 97 | + ("report_timeout", report_timeout), |
| 98 | + ("quick_timeout", quick_timeout), |
| 99 | + ): |
| 100 | + if isinstance(timeout, bool) or not isinstance(timeout, (int, float)) or not 0 < timeout <= float_info.max: |
| 101 | + raise ValueError(f"{name} must be positive and finite.") |
| 102 | + self._loop = asyncio.get_running_loop() |
| 103 | + self._lanes = { |
| 104 | + True: _Lane(limit=report_workers, timeout=report_timeout), |
| 105 | + False: _Lane(limit=quick_workers, timeout=quick_timeout), |
| 106 | + } |
| 107 | + self._max_queue = max_queue |
| 108 | + self._queue_timeout = queue_timeout |
| 109 | + self._closing = False |
| 110 | + self._closed = False |
| 111 | + self._close_task: asyncio.Task[None] | None = None |
| 112 | + |
| 113 | + @property |
| 114 | + def is_closed(self) -> bool: |
| 115 | + """Whether shutdown has drained all started operations, not just their callers.""" |
| 116 | + return self._closed |
| 117 | + |
| 118 | + async def run_async(self, *, report: bool, task: Callable[[QueryControl], Awaitable[T]]) -> T: |
| 119 | + """ |
| 120 | + Admit one operation and await its response without blocking the owning loop. |
| 121 | +
|
| 122 | + Args: |
| 123 | + report (bool): Use the report lane rather than the reserved quick lane. |
| 124 | + task (Callable[[QueryControl], Awaitable[T]]): Native async work that honors |
| 125 | + the shared control and does not return before releasing its resources. |
| 126 | +
|
| 127 | + Returns: |
| 128 | + T: This operation's result. Separate calls never share response objects. |
| 129 | +
|
| 130 | + Raises: |
| 131 | + AnalyticsBusyException: If the lane is full, admission expires, or shutdown has begun. |
| 132 | + AnalyticsTimeoutException: If execution exceeds its budget. |
| 133 | + CancelledError: If the caller leaves; running work retains its slot until cleanup finishes. |
| 134 | + RuntimeError: If used from a different event loop. |
| 135 | + """ |
| 136 | + self._check_loop() |
| 137 | + lane = self._lanes[report] |
| 138 | + await self._admit_async(lane) |
| 139 | + work = _Operation( |
| 140 | + control=QueryControl(deadline=monotonic() + lane.timeout), |
| 141 | + finished=self._loop.create_future(), |
| 142 | + ) |
| 143 | + lane.running.add(work) |
| 144 | + try: |
| 145 | + operation = self._create_task(self._execute_async(task=task, control=work.control)) |
| 146 | + except BaseException: |
| 147 | + self._finish(lane=lane, work=work) |
| 148 | + raise |
| 149 | + work.task = operation |
| 150 | + operation.add_done_callback(lambda completed: self._complete(lane=lane, work=work, task=completed)) |
| 151 | + try: |
| 152 | + done, _ = await asyncio.wait({operation}, timeout=work.control.remaining) |
| 153 | + if not done: |
| 154 | + self._abandon(work=work, task=operation) |
| 155 | + raise AnalyticsTimeoutException |
| 156 | + return operation.result() |
| 157 | + except asyncio.CancelledError: |
| 158 | + self._abandon(work=work, task=operation) |
| 159 | + raise |
| 160 | + |
| 161 | + async def close_async(self) -> None: |
| 162 | + """ |
| 163 | + Reject admission, signal cancellation, and drain actual work on the owning loop. |
| 164 | +
|
| 165 | + This is idempotent and terminal. It can outlast response deadlines because |
| 166 | + returning early would allow a replacement controller to overlap database |
| 167 | + operations still cleaning up. Cancelling this await, even repeatedly, is |
| 168 | + propagated only after draining. The memory backend itself is not disposed. |
| 169 | + If scheduling the drain fails, admission remains closed and a later |
| 170 | + ``close_async`` call can retry the drain. |
| 171 | + """ |
| 172 | + self._check_loop() |
| 173 | + if self._close_task is None: |
| 174 | + self._closing = True |
| 175 | + for lane in self._lanes.values(): |
| 176 | + while lane.queued: |
| 177 | + waiter = lane.queued.popleft() |
| 178 | + if not waiter.ready.done(): |
| 179 | + waiter.ready.set_exception(AnalyticsBusyException()) |
| 180 | + for work in lane.running: |
| 181 | + work.control.cancel() |
| 182 | + self._close_task = self._create_task(self._drain_async()) |
| 183 | + cancellation: asyncio.CancelledError | None = None |
| 184 | + while not self._close_task.done(): |
| 185 | + try: |
| 186 | + await asyncio.shield(self._close_task) |
| 187 | + except asyncio.CancelledError as error: |
| 188 | + cancellation = error |
| 189 | + self._close_task.result() |
| 190 | + if cancellation is not None: |
| 191 | + raise cancellation |
| 192 | + |
| 193 | + def _check_loop(self) -> None: |
| 194 | + if asyncio.get_running_loop() is not self._loop: |
| 195 | + raise RuntimeError("Analytics must be used and closed on its owning event loop.") |
| 196 | + |
| 197 | + def _create_task(self, coroutine: Coroutine[object, object, T]) -> asyncio.Task[T]: |
| 198 | + try: |
| 199 | + return self._loop.create_task(coroutine) |
| 200 | + except BaseException: |
| 201 | + coroutine.close() |
| 202 | + raise |
| 203 | + |
| 204 | + async def _admit_async(self, lane: _Lane) -> None: |
| 205 | + if self._closing: |
| 206 | + raise AnalyticsBusyException |
| 207 | + if lane.active < lane.limit: |
| 208 | + lane.active += 1 |
| 209 | + return |
| 210 | + if len(lane.queued) >= self._max_queue: |
| 211 | + raise AnalyticsBusyException |
| 212 | + waiter = _Waiter(ready=self._loop.create_future(), deadline=monotonic() + self._queue_timeout) |
| 213 | + lane.queued.append(waiter) |
| 214 | + try: |
| 215 | + try: |
| 216 | + async with asyncio.timeout(self._queue_timeout): |
| 217 | + await waiter.ready |
| 218 | + except TimeoutError as error: |
| 219 | + raise AnalyticsBusyException from error |
| 220 | + if self._closing or monotonic() >= waiter.deadline: |
| 221 | + raise AnalyticsBusyException |
| 222 | + except BaseException: |
| 223 | + if waiter in lane.queued: |
| 224 | + lane.queued.remove(waiter) |
| 225 | + elif not waiter.ready.cancelled() and waiter.ready.exception() is None: |
| 226 | + self._release(lane) |
| 227 | + raise |
| 228 | + |
| 229 | + def _release(self, lane: _Lane) -> None: |
| 230 | + lane.active -= 1 |
| 231 | + while lane.queued and not self._closing: |
| 232 | + waiter = lane.queued.popleft() |
| 233 | + if waiter.ready.done(): |
| 234 | + continue |
| 235 | + if monotonic() >= waiter.deadline: |
| 236 | + waiter.ready.set_exception(AnalyticsBusyException()) |
| 237 | + continue |
| 238 | + lane.active += 1 |
| 239 | + waiter.ready.set_result(None) |
| 240 | + break |
| 241 | + |
| 242 | + def _complete(self, *, lane: _Lane, work: _Operation, task: asyncio.Task[T]) -> None: |
| 243 | + if not task.cancelled(): |
| 244 | + task.exception() |
| 245 | + self._finish(lane=lane, work=work) |
| 246 | + if work.abandoned: |
| 247 | + self._log_abandoned_error(task) |
| 248 | + |
| 249 | + def _finish(self, *, lane: _Lane, work: _Operation) -> None: |
| 250 | + lane.running.remove(work) |
| 251 | + work.task = None |
| 252 | + self._release(lane) |
| 253 | + work.finished.set_result(None) |
| 254 | + |
| 255 | + def _abandon(self, *, work: _Operation, task: asyncio.Task[T]) -> None: |
| 256 | + work.abandoned = True |
| 257 | + work.control.cancel() |
| 258 | + # Completion can precede cancellation delivery, after its callback consumed the exception. |
| 259 | + if work.finished.done(): |
| 260 | + self._log_abandoned_error(task) |
| 261 | + |
| 262 | + async def _drain_async(self) -> None: |
| 263 | + await gather_with_cleanup_async(work.finished for lane in self._lanes.values() for work in lane.running) |
| 264 | + self._closed = True |
| 265 | + |
| 266 | + @staticmethod |
| 267 | + async def _execute_async(*, task: Callable[[QueryControl], Awaitable[T]], control: QueryControl) -> T: |
| 268 | + control.check() |
| 269 | + result = await task(control) |
| 270 | + control.check() |
| 271 | + return result |
| 272 | + |
| 273 | + @staticmethod |
| 274 | + def _log_abandoned_error(task: asyncio.Task[T]) -> None: |
| 275 | + if not task.cancelled(): |
| 276 | + error = task.exception() |
| 277 | + if error is not None and not isinstance(error, AnalyticsTimeoutException): |
| 278 | + logger.error("Analytics operation failed after its caller left.", exc_info=error) |
0 commit comments