diff --git a/CHANGELOG.md b/CHANGELOG.md index a067f9f..6d30d86 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/). ## [Unreleased] +### Added + +- Flexible request body types on call_runner and selectors. + ### Changed - `discover_orchestrators` is now asynchronous and must be awaited, including @@ -13,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/). ### Removed +- The unused `label` option from `register_runner` and `LiveRunnerRegistration`. - Synchronous HTTP helpers `request_json_sync`, `get_json_sync`, and `post_json_sync`. All HTTP requests now use the asynchronous transport. - TLS certificate-verification bypasses in HTTP requests, streaming responses, diff --git a/README b/README index 6e32fcf..312f2fe 100644 --- a/README +++ b/README @@ -80,8 +80,8 @@ asyncio.run(main("")) ``` For single-shot apps, use `runner_selector()` and `await cursor.next()`, or -`call_runner()` for a known runner URL. Paid runners support `hour`, `seconds`, -and `fixed` pricing through a remote signer. +`call_runner()` for a known runner URL. See [Calling runners](docs/runner-payloads.md) +for request options, streaming uploads and responses, payment, and retries. Read demuxed packets without decoding, using the output channel URL returned by the app: diff --git a/docs/runner-payloads.md b/docs/runner-payloads.md new file mode 100644 index 0000000..347f831 --- /dev/null +++ b/docs/runner-payloads.md @@ -0,0 +1,193 @@ +# Calling runners + +Use `call_runner()` for a known runner URL, or a [selector](#selectors) to discover +one. By default, calls return a complete `LiveRunnerCallResult`: JSON objects in +`result.data`, other formats as bytes in `result.content`, and the response media +type in `result.content_type`. + +## `call_runner` parameters + +All parameters except `runner_url` must be passed by keyword. + +| Parameter | Default | Meaning | +|---|---|---| +| `runner_url` | `""` | Runner endpoint URL. Required unless supplied by `runner`. | +| `runner` | `None` | A `LiveRunnerInstance` with URL and app/pricing metadata. An explicit `runner_url` takes precedence. | +| `payload` | `None` | Request body; see [body types](#body-types). | +| `content_type` | `None` | Override the request's `Content-Type` header without changing its encoding. | +| `method` | `"POST"` | HTTP request method. | +| `signer_url` | `None` | Remote signer URL for automatic payment. Required for paid calls. | +| `signer_headers` | `None` | Dictionary of additional HTTP headers for signer requests. | +| `payment_unit` | `None` | `"hour"`, `"seconds"`, or `"fixed"`. Defaults to the runner's pricing unit; an explicit value must agree with it. With neither, payment uses ongoing funding. | +| `timeout` | `5.0` | Seconds per request attempt; see [timeouts](#timeouts) for streaming behavior. | +| `max_payment_challenge_retries` | `3` | Payment retry budget, with additional allowance to refresh expired payment parameters. | +| `stream` | `False` | Set to `True` for an incremental [response](#streaming-responses) as a `LiveRunnerCallStream`. | + +### Body types + +`dict` values default to `application/json`; other bodies default to +`application/octet-stream`. Set `content_type` for another media type, such as +`"audio/wav"`. + +| `payload` | Sent as | +|---|---| +| `dict` | JSON. | +| `str` | UTF-8. | +| `bytes`, `bytearray`, or `memoryview` | Raw bytes. | +| `BinaryIO` or `AsyncIterable[bytes]` | Incremental upload using `StreamPayload` defaults. | +| `StreamPayload` | Incremental upload with explicit settings. | +| Omitted or `None` | Empty body. | + +## Streaming uploads + +`StreamPayload` is an immutable configuration for a file or async iterable +payload. Bare files and async iterables use the same defaults. + +| Field | Default | Meaning | +|---|---|---| +| `source` | Required | A binary file with a synchronous `read()` method, or an async iterable yielding `bytes`. | +| `expect100` | `True` | Send headers first and wait for a continuation response before reading or sending input; see [when input is read](#when-input-is-read). Set to `False` for endpoints without this support. | +| `retry_buffer` | `True` | Retain input from non-seekable files and async iterables so it can be sent again on a retry. | +| `retry_buffer_memory_limit` | `25_000_000` | Bytes of retained input allowed in memory before using temporary disk storage (25 MB). | +| `retry_buffer_limit` | `250_000_000` | Total bytes of input retained for retries (250 MB). | + +Limits use decimal bytes and require +`0 < retry_buffer_memory_limit <= retry_buffer_limit`. They bound the retained +copy, not upload size or memory used by your code. Exceeding the total limit or a +storage failure lets the current upload continue but prevents later retries. +Seekable files rewind for retries and do not use these limits. + +### When input is read + +Creating a `StreamPayload` or discovering runners does not read the input. +Reading begins during `call_runner()` or `cursor.next()`: with `expect100=True`, +headers are sent first so the runner can request payment or reject the call +before consuming input. The SDK uses aiohttp's native handling: `100 Continue` +normally starts the upload, but other informational `1xx` responses except +`101 Switching Protocols`, such as `102 Processing` or `103 Early Hints`, can +also start it (unexpected informational responses like these are rare in +practice). If waiting for this response times out, the request fails before +reading or sending the body. + +With `expect100=False`, reading starts after the headers are sent. In either +case, input is read as it uploads, starting at the file's current position. +An early response or failure can leave it partly consumed. + +You are responsible for closing files and managing generator resources. Keep +inputs usable until the call finishes or the response/selector context closes, +and do not read them elsewhere during the operation. The examples nest contexts +to ensure this. + +For an endpoint without `100 Continue` support: + +```python +from livepeer_gateway import StreamPayload, call_runner + + +async def upload(runner_url, path): + with open(path, "rb") as audio: + return await call_runner( + runner_url, + payload=StreamPayload(source=audio, expect100=False), + content_type="audio/wav", + timeout=60, + ) +``` + +## Retries + +A retry sends the same body in a new request: + +- After `402 Payment Required`, `call_runner()` obtains payment through the + configured signer and submits the request again. +- If a runner call fails, a selector's `next()` tries another candidate. + +For streamed uploads, `call_runner()` does not automatically resend after +connection failures or follow redirects. A selector can try another candidate. + +Failover can repeat work if the original runner executed it but its response was +lost. + +Dictionaries, strings, and byte values can be sent again directly. Streamed input +must still be usable: + +- Seekable files rewind to their position at the start of the operation, even + with `retry_buffer=False`. +- Non-seekable files and async iterables can resend the full body, including + after a partial upload, if retry buffering remains available. +- With `retry_buffer=False`, non-seekable input can only be retried before reading + starts—for example, after a payment challenge received before upload permission. + +Retry eligibility follows actual input consumption. If an informational response +starts an unbuffered upload before a payment challenge, the consumed input can +prevent an automatic retry. + +A required retry raises `LivepeerGatewayError` if the body cannot be sent again. +This retry support lasts for one call or selector cursor. For a separate call, +supply a fresh input or rewind the file yourself. + +## Streaming responses + +`stream=True` controls the **response**, independently of the upload type. It +returns after payment challenges are handled and final response headers arrive; +the upload may still be running. Read bytes with `aiter_bytes()` or decoded text +lines with `aiter_lines()`. JSON is not parsed automatically. + +For an app that accepts audio and returns text incrementally: + +```python +from livepeer_gateway import call_runner + + +async def transcribe(runner_url, path): + with open(path, "rb") as audio: + async with await call_runner( + runner_url, + payload=audio, + content_type="audio/wav", + stream=True, + ) as response: + async for line in response.aiter_lines(): + print(line) +``` + +Use the response as an async context manager or call `await response.aclose()`. +Closing it closes the connection and stops any unfinished upload and ongoing +payment funding. Errors while reading a returned response are not retried. + +### Timeouts + +| Response mode | What `timeout` bounds per attempt | +|---|---| +| `stream=False` | The whole HTTP request, including upload and response reading. Increase it for longer uploads or jobs. | +| `stream=True` | Connection establishment and, with `expect100=True`, [waiting for a continuation response](#when-input-is-read). Uploads and response reading have no deadline after that response; with `expect100=False`, only connection establishment is bounded. | + +Retries can make the overall operation take longer. Use `asyncio.timeout()` +around the operation if you need an overall deadline. + +## Selectors + +`runner_selector()` discovers candidates and returns a cursor. Pass input as +`body` instead of `payload`, with the same optional `content_type`. Each +`cursor.next()` returns the next successful result; another call sends the same +body to subsequent candidates. Responses are always read in full. +Await each `cursor.next()` before calling it again. + +Use an async context manager or `await cursor.aclose()` when finished with the +cursor. Closing it ends selection but leaves your file open. + +```python +from livepeer_gateway import runner_selector + + +async def upload_to_available_runner(path, discovery_url): + with open(path, "rb") as audio: + async with await runner_selector( + discovery_url=discovery_url, + app="example/transcribe", + body=audio, + content_type="audio/wav", + timeout=60, + ) as cursor: + return await cursor.next() +``` diff --git a/examples/text/runners.toml b/examples/text/runners.toml index 0ebb404..90e623c 100644 --- a/examples/text/runners.toml +++ b/examples/text/runners.toml @@ -1,6 +1,5 @@ [[Runners]] ID = "story-runner" -Label = "story-runner" App = "livepeer/read-story" RunnerURL = "http://127.0.0.1:8990" HealthURL = "http://127.0.0.1:8990/healthz" diff --git a/pyproject.toml b/pyproject.toml index 3c19135..c405c6e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,6 +59,7 @@ include = [ "/CHANGELOG.md", "/LICENSE", "/README", + "/docs", "/examples", "/pyproject.toml", "/src", diff --git a/src/livepeer_gateway/__init__.py b/src/livepeer_gateway/__init__.py index 9e4c284..da6d233 100644 --- a/src/livepeer_gateway/__init__.py +++ b/src/livepeer_gateway/__init__.py @@ -46,6 +46,7 @@ ) from .discovery import discover_orchestrators, discover_runners from .remote_signer import LivePaymentChallenge, LivePaymentSession +from .runner_payload import RunnerPayload, StreamPayload from .selection import ( RunnerSelectionCursor, runner_selector, @@ -79,6 +80,8 @@ "LivePaymentChallenge", "LivePaymentSession", "LiveRunnerProxy", + "RunnerPayload", + "StreamPayload", "LivepeerGatewayError", "LivepeerHTTPError", "NoRunnerAvailableError", diff --git a/src/livepeer_gateway/_http_continue.py b/src/livepeer_gateway/_http_continue.py new file mode 100644 index 0000000..d07fc31 --- /dev/null +++ b/src/livepeer_gateway/_http_continue.py @@ -0,0 +1,54 @@ +"""Bound aiohttp's native continue handshake without replacing response parsing.""" + +from __future__ import annotations + +import asyncio +import contextlib +from typing import Any + +import aiohttp + + +class _ContinueResponse(aiohttp.ClientResponse): + def __init__(self, *args: Any, continue_timeout: float, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._continue_timeout = continue_timeout + + async def start(self, connection: Any) -> aiohttp.ClientResponse: + continued = self._continue + if continued is None: + return await super().start(connection) + + # aiohttp releases this waiter on any informational response except 101. + # Only the wait is bounded; uploads/final headers can take longer once + # permission arrives. Final headers also end the wait through start(). + waiting = True + async with asyncio.timeout(self._continue_timeout) as handshake: + def allow_upload(future: asyncio.Future[bool]) -> None: + if waiting and not future.cancelled() and not handshake.expired(): + handshake.reschedule(None) + + continued.add_done_callback(allow_upload) + try: + return await super().start(connection) + finally: + waiting = False + continued.remove_done_callback(allow_upload) + + +async def _close_response(response: aiohttp.ClientResponse) -> None: + """Stop and join an attempt's writer before another attempt borrows its state.""" + writer = response._writer + response.close() + await _cancel_writer(writer) + + +async def _cancel_writer(writer: asyncio.Task[None] | None) -> None: + if writer is not None: + writer.cancel() + # aiohttp 3.9 tries write_eof() on cancellation. Its failure after the + # peer closes must not mask a final response that was already received. + with contextlib.suppress( + asyncio.CancelledError, ConnectionError, aiohttp.ClientError + ): + await writer diff --git a/src/livepeer_gateway/http.py b/src/livepeer_gateway/http.py index 964a32e..9bfec69 100644 --- a/src/livepeer_gateway/http.py +++ b/src/livepeer_gateway/http.py @@ -1,11 +1,14 @@ from __future__ import annotations import json +from functools import partial from typing import Any from urllib.parse import ParseResult, urlparse import aiohttp +from ._http_continue import _ContinueResponse, _close_response +from .runner_payload import RunnerPayload, _StreamState, _UploadBody from .errors import ( LivepeerHTTPError, LivepeerGatewayError, @@ -48,21 +51,32 @@ def _extract_error_message_from_body(body: str) -> str: return _truncate(body) -def _json_request_parts( +def _request_parts( url: str, *, method: str | None = None, - payload: dict[str, Any] | None = None, + payload: RunnerPayload | _StreamState = None, headers: dict[str, str] | None = None, -) -> tuple[str, dict[str, str], bytes | None]: +) -> tuple[str, dict[str, str], bytes | _UploadBody | None]: req_headers: dict[str, str] = { "Accept": "application/json", "User-Agent": "livepeer-python-gateway/0.1", } - body: bytes | None = None + body: bytes | _UploadBody | None = None if payload is not None: - req_headers["Content-Type"] = "application/json" + req_headers["Content-Type"] = "application/json" if isinstance(payload, dict) else "application/octet-stream" + if isinstance(payload, dict): body = json.dumps(payload).encode("utf-8") + elif isinstance(payload, _StreamState): + payload.require_replayable() + body = _UploadBody(payload) + elif payload is not None: + if isinstance(payload, str): + body = payload.encode("utf-8") + elif isinstance(payload, (bytes, bytearray, memoryview)): + body = bytes(payload) + else: + raise TypeError(f"Unsupported runner payload: {type(payload).__name__}") if headers: req_headers.update(headers) @@ -70,6 +84,30 @@ def _json_request_parts( return resolved_method, req_headers, body +def _request_session( + payload: RunnerPayload | _StreamState, + timeout: aiohttp.ClientTimeout, + handshake_timeout: float, +) -> tuple[aiohttp.ClientSession, dict[str, Any]]: + if not isinstance(payload, _StreamState): + return aiohttp.ClientSession(timeout=timeout), {} + session = aiohttp.ClientSession( + timeout=timeout, + response_class=partial(_ContinueResponse, continue_timeout=handshake_timeout), + ) + # aiohttp >= 3.10 may silently retry idempotent requests after disconnect. + # Stream resends are controlled by the SDK's replay state instead. + if hasattr(session, "_retry_connection"): + session._retry_connection = False + return session, {"expect100": payload.config.expect100, "allow_redirects": False} + + +def _request_failed(status: int, payload: RunnerPayload | _StreamState) -> bool: + return status >= 400 or ( + isinstance(payload, _StreamState) and status in (301, 302, 303, 307, 308) + ) + + def _raise_http_json_error( status: int, url: str, @@ -105,12 +143,12 @@ async def _request_body( url: str, *, method: str | None = None, - payload: dict[str, Any] | None = None, + payload: RunnerPayload | _StreamState = None, headers: dict[str, str] | None = None, timeout: float = 5.0, ) -> tuple[bytes, str]: """ - Make an async JSON-payload HTTP request and return the raw response body. + Make an HTTP request and return the raw response body. Returns ``(body, content_type)`` without assuming the response is JSON; request semantics and error mapping match request_json. @@ -119,23 +157,17 @@ async def _request_body( Raises LivepeerGatewayError on HTTP/network errors. """ - resolved_method, req_headers, body = _json_request_parts( - url, - method=method, - payload=payload, - headers=headers, - ) - try: - client_timeout = aiohttp.ClientTimeout(total=timeout) - async with aiohttp.ClientSession(timeout=client_timeout) as session: - async with session.request(resolved_method, url, data=body, headers=req_headers) as resp: - raw = await resp.read() + session, resp = await _open_request( + url, method=method, payload=payload, headers=headers, + timeout=aiohttp.ClientTimeout(total=timeout), handshake_timeout=timeout, + ) + async with session: + try: + raw = await resp.content.read() if isinstance(payload, _StreamState) else await resp.read() content_type = resp.content_type or "" - if resp.status >= 400: - _raise_http_json_error( - resp.status, url, raw.decode(errors="replace") - ) + finally: + await _close_response(resp) except (SignerRefreshRequired, SkipPaymentCycle, LivepeerGatewayError): raise except ConnectionRefusedError as e: @@ -159,6 +191,9 @@ async def _request_body( raise LivepeerGatewayError( f"HTTP JSON error: unexpected error: {e.__class__.__name__}: {e} (url={url})" ) from e + finally: + if isinstance(payload, _StreamState): + await payload.finish_attempt() return raw, content_type @@ -197,7 +232,7 @@ async def open_stream( url: str, *, method: str | None = None, - payload: dict[str, Any] | None = None, + payload: RunnerPayload | _StreamState = None, headers: dict[str, str] | None = None, connect_timeout: float = 10.0, ) -> tuple[aiohttp.ClientSession, aiohttp.ClientResponse]: @@ -206,31 +241,53 @@ async def open_stream( body, for streaming responses (SSE, chunked). The caller owns both and must close them. - No total timeout (streams run indefinitely) only connect/first-byte are bounded. + No total/read timeout; bounds connection and optional 100-continue handshake. Raises LivepeerHTTPError on >= 400 (e.g. the 402 payment retry). """ - resolved_method, req_headers, body = _json_request_parts( - url, - method=method, - payload=payload, - headers=headers, - ) - - timeout = aiohttp.ClientTimeout(total=None, sock_connect=connect_timeout, sock_read=None) - session = aiohttp.ClientSession(timeout=timeout) try: - resp = await session.request(resolved_method, url, data=body, headers=req_headers) + return await _open_request( + url, method=method, payload=payload, headers=headers, + timeout=aiohttp.ClientTimeout(total=None, sock_connect=connect_timeout, sock_read=None), + handshake_timeout=connect_timeout, + ) except (TimeoutError, aiohttp.ClientError) as e: - await session.close() raise LivepeerGatewayError( f"HTTP stream error: failed to reach endpoint: {getattr(e, 'message', e)} (url={url})" ) from e - if resp.status >= 400: - raw = await resp.text() - resp.release() - await session.close() - _raise_http_json_error(resp.status, url, raw) - return session, resp + + +async def _open_request( + url: str, + *, + method: str | None, + payload: RunnerPayload | _StreamState, + headers: dict[str, str] | None, + timeout: aiohttp.ClientTimeout, + handshake_timeout: float, +) -> tuple[aiohttp.ClientSession, aiohttp.ClientResponse]: + method, headers, body = _request_parts(url, method=method, payload=payload, headers=headers) + session, options = _request_session(payload, timeout, handshake_timeout) + if payload is None: + # aiohttp otherwise adds application/octet-stream even for an empty POST. + options["skip_auto_headers"] = {"Content-Type"} + resp = None + try: + resp = await session.request(method, url, data=body, headers=headers, **options) + if _request_failed(resp.status, payload): + # Read before cancelling the writer: aiohttp 3.13 closes the socket + # on cancellation, which can discard a fragmented payment challenge. + raw = (await resp.content.read()).decode(errors="replace") + _raise_http_json_error(resp.status, url, raw) + return session, resp + except BaseException: + try: + if resp is not None: + await _close_response(resp) + finally: + await session.close() + if isinstance(payload, _StreamState): + await payload.finish_attempt() + raise async def post_json( diff --git a/src/livepeer_gateway/live_runner.py b/src/livepeer_gateway/live_runner.py index 8d80b3b..6bcb280 100644 --- a/src/livepeer_gateway/live_runner.py +++ b/src/livepeer_gateway/live_runner.py @@ -26,6 +26,7 @@ import aiohttp from aiohttp.helpers import parse_mimetype +from ._http_continue import _close_response from .channel_reader import ChannelReader from .errors import LivepeerGatewayError, LivepeerHTTPError, SignerRefreshRequired from .http import _post_empty, _request_body, open_stream, post_json, request_json @@ -36,6 +37,7 @@ _freeze_headers, get_signer_info, ) +from .runner_payload import RunnerPayload, _StreamState, _stream_state _LOG = logging.getLogger(__name__) @@ -188,6 +190,9 @@ class LiveRunnerCallStream: payment_session: LivePaymentSession | None _session: aiohttp.ClientSession = field(repr=False, compare=False) _response: aiohttp.ClientResponse = field(repr=False, compare=False) + _payload_state: _StreamState | None = field( + default=None, init=False, repr=False, compare=False + ) # True once the orchestrator reported the backing session gone. released: bool = False _payment_task: asyncio.Task[None] | None = field( @@ -219,10 +224,16 @@ def _start_payments( async def aclose(self) -> None: # Stop funding first: don't pay for a stream we are about to drop. - await _stop_funding(self._payment_task) - self._payment_task = None - self._response.release() - await self._session.close() + try: + await _stop_funding(self._payment_task) + finally: + self._payment_task = None + try: + await _close_response(self._response) + finally: + await self._session.close() + if self._payload_state is not None: + await self._payload_state.aclose() async def __aenter__(self) -> LiveRunnerCallStream: return self @@ -274,7 +285,6 @@ def __init__( runner_id: str = "", mode: str = "persistent", proxy: bool = False, - label: str = "", version: str = "", metadata: str = "", status: str = "ready", @@ -298,7 +308,6 @@ def __init__( self._mode = _normalize_runner_mode(mode) self._proxy = proxy self._price_info = price_info - self._label = label self._version = version self._metadata = metadata self._status = status @@ -445,8 +454,6 @@ def _payload(self) -> dict[str, Any]: payload["runner_id"] = self.runner_id if self._proxy: payload["proxy"] = True - if self._label: - payload["label"] = self._label if self._version: payload["version"] = self._version if self._metadata: @@ -609,7 +616,6 @@ async def register_runner( runner_id: str = "", mode: str = "persistent", proxy: bool = False, - label: str = "", version: str = "", metadata: str = "", status: str = "ready", @@ -634,7 +640,6 @@ async def register_runner( runner_id=runner_id, mode=mode, proxy=proxy, - label=label, version=version, metadata=metadata, status=status, @@ -755,7 +760,8 @@ async def call_runner( runner_url: str = ..., *, runner: LiveRunnerInstance | None = ..., - payload: dict[str, Any] | None = ..., + payload: RunnerPayload = ..., + content_type: str | None = ..., method: str = ..., signer_url: str | None = ..., signer_headers: dict[str, str] | None = ..., @@ -772,7 +778,8 @@ async def call_runner( runner_url: str = ..., *, runner: LiveRunnerInstance | None = ..., - payload: dict[str, Any] | None = ..., + payload: RunnerPayload = ..., + content_type: str | None = ..., method: str = ..., signer_url: str | None = ..., signer_headers: dict[str, str] | None = ..., @@ -788,7 +795,8 @@ async def call_runner( runner_url: str = "", *, runner: LiveRunnerInstance | None = None, - payload: dict[str, Any] | None = None, + payload: RunnerPayload | _StreamState = None, + content_type: str | None = None, method: str = "POST", signer_url: str | None = None, signer_headers: dict[str, str] | None = None, @@ -807,137 +815,153 @@ async def call_runner( (an image, ndjson) comes back unparsed in ``result.content`` + ``result.content_type``. The request asks for no particular format, so the app picks what it returns. + + Dictionaries serialize to JSON; strings encode as UTF-8; bytes and streams + send raw data. ``content_type`` overrides the default ``application/json`` + for dictionaries or ``application/octet-stream`` for other bodies. + Omitted/None payloads send an empty body. Bare files and async byte iterables + use StreamPayload defaults (100-continue and bounded replay buffering). """ - runner_url = runner_url.strip() or (runner.url.strip() if runner is not None else "") - if not runner_url: - raise LivepeerGatewayError("Live runner call requires runner_url") - request_payload = payload or {} - payment_type = _runner_payment_type(runner, payment_unit) if signer_url else "" - max_price: LiveRunnerPriceInfo | None = None - if signer_url and runner is not None and runner.price_info is not None: - max_price = _pad_runner_price(runner.price_info) - payer_address = "" - if signer_url: - signer = await get_signer_info(signer_url, _freeze_headers(signer_headers)) - payer_address = cast(str, signer.address) - challenge: LivePaymentChallenge | None = None - attempts = (max(0, int(max_payment_challenge_retries)) + 1) * 2 - for attempt in range(attempts): - payment_session: LivePaymentSession | None = None - session_id = "" - needs_ongoing_funding = False - # No preferred format: the app, or the upstream it fronts, picks. Only - # control-plane calls ask for JSON. - request_headers: dict[str, str] = {"Accept": "*/*"} + state = _stream_state(payload) + owned = state is not None and not isinstance(payload, _StreamState) + payload = state if state is not None else payload + try: + runner_url = runner_url.strip() or (runner.url.strip() if runner is not None else "") + if not runner_url: + raise LivepeerGatewayError("Live runner call requires runner_url") + payment_type = _runner_payment_type(runner, payment_unit) if signer_url else "" + max_price: LiveRunnerPriceInfo | None = None + if signer_url and runner is not None and runner.price_info is not None: + max_price = _pad_runner_price(runner.price_info) + payer_address = "" if signer_url: - request_headers[_LIVE_RUNNER_PAYER_ADDRESS_HEADER] = payer_address - # Pending challenge means payment is needed. - if challenge is not None: - try: - payment_session, payment = await _get_runner_payment( - challenge, - payment_type=payment_type, - app=runner.app if runner is not None else None, - signer_url=signer_url or "", - signer_headers=signer_headers, - max_price=max_price, - ) - except SignerRefreshRequired as e: - if attempt + 1 >= attempts: - raise - # Could happen if embedded payment params expire; just retry in this case. - _LOG.info( - "Live runner reservation payment challenge needs refresh; retrying with a fresh challenge: %s", - e, - ) - challenge = None - continue - request_headers["Livepeer-Payment"] = payment.payment - request_headers["Livepeer-Segment"] = payment.seg_creds or "" - session_id = challenge.manifest_id + signer = await get_signer_info(signer_url, _freeze_headers(signer_headers)) + payer_address = cast(str, signer.address) + challenge: LivePaymentChallenge | None = None + attempts = (max(0, int(max_payment_challenge_retries)) + 1) * 2 + for attempt in range(attempts): + payment_session: LivePaymentSession | None = None + session_id = "" + needs_ongoing_funding = False + # No preferred format: the app, or the upstream it fronts, picks. Only + # control-plane calls ask for JSON. + request_headers: dict[str, str] = {"Accept": "*/*"} + if content_type is not None: + request_headers["Content-Type"] = content_type + if signer_url: + request_headers[_LIVE_RUNNER_PAYER_ADDRESS_HEADER] = payer_address + # Pending challenge means payment is needed. + if challenge is not None: + try: + payment_session, payment = await _get_runner_payment( + challenge, + payment_type=payment_type, + app=runner.app if runner is not None else None, + signer_url=signer_url or "", + signer_headers=signer_headers, + max_price=max_price, + ) + except SignerRefreshRequired as e: + if attempt + 1 >= attempts: + raise + # Could happen if embedded payment params expire; just retry in this case. + _LOG.info( + "Live runner reservation payment challenge needs refresh; retrying with a fresh challenge: %s", + e, + ) + challenge = None + continue + request_headers["Livepeer-Payment"] = payment.payment + request_headers["Livepeer-Segment"] = payment.seg_creds or "" + session_id = challenge.manifest_id - # Metered pricing bills for as long as the work runs. - needs_ongoing_funding = payment_type == "live" + # Metered pricing bills for as long as the work runs. + needs_ongoing_funding = payment_type == "live" - try: - request_kwargs: dict[str, Any] = {"timeout": timeout} - if request_headers: - request_kwargs["headers"] = request_headers - - if stream: - # Hand back the live response unbuffered. open_stream raises on a 402 - # before any body, so the payment retry below still catches it. - session, resp = await open_stream( - runner_url, - method=method, - payload=request_payload, - headers=request_headers or None, - ) - call_stream = LiveRunnerCallStream( - resp.status, - resp.headers, - runner_url, - runner, - None if payment_type == "fixed" else payment_session, - session, - resp, - ) - # The stream outlives this call, so it owns the funding. - if needs_ongoing_funding: - call_stream._start_payments(cast(LivePaymentSession, payment_session)) - return call_stream - - # The request ends with this call, so the funding ends with it. - pay_task = ( - _start_funding(cast(LivePaymentSession, payment_session)) - if needs_ongoing_funding - else None - ) try: - body, content_type = await _request_body( - runner_url, - method=method, - payload=request_payload, - **request_kwargs, + if stream: + # Payment challenges are handled before returning a live response. + session, resp = await open_stream( + runner_url, + method=method, + payload=payload, + headers=request_headers or None, + connect_timeout=timeout, + ) + call_stream = LiveRunnerCallStream( + resp.status, + resp.headers, + runner_url, + runner, + None if payment_type == "fixed" else payment_session, + session, + resp, + ) + # The stream outlives this call, so it owns the funding. + if needs_ongoing_funding: + call_stream._start_payments(cast(LivePaymentSession, payment_session)) + if owned: + call_stream._payload_state = state + owned = False + return call_stream + + # The request ends with this call, so the funding ends with it. + pay_task = ( + _start_funding(cast(LivePaymentSession, payment_session)) + if needs_ongoing_funding + else None ) - finally: - await _stop_funding(pay_task) - # Non-JSON bodies (an image, ndjson) are handed back unparsed in `content`. - is_json = _is_json_content_type(content_type) - data: dict[str, Any] = {} - if is_json: try: - data = json.loads(body) - except (UnicodeDecodeError, json.JSONDecodeError) as e: - raise LivepeerGatewayError( - f"HTTP JSON error: endpoint did not return valid JSON: {e} " - f"(url={runner_url}, content_type={content_type})" - ) from e - if not isinstance(data, dict): - raise LivepeerGatewayError( - f"Live runner call expected JSON object, got {type(data).__name__}" + body, response_content_type = await _request_body( + runner_url, + method=method, + payload=payload, + headers=request_headers, + timeout=timeout, ) - return LiveRunnerCallResult( - data, - runner_url=runner_url, - runner=runner, - session_id=( - session_id - or (data["session_id"].strip() if isinstance(data.get("session_id"), str) else "") - ), - payment_session=None if payment_type == "fixed" else payment_session, - content=None if is_json else body, - content_type=content_type, - ) - except LivepeerHTTPError as e: - if e.status_code != 402: - raise - if not signer_url: - raise LivepeerGatewayError("Live runner paid call requires signer_url") from e - challenge = _parse_runner_payment_challenge(e) - continue + finally: + await _stop_funding(pay_task) + # Non-JSON bodies (an image, ndjson) are handed back unparsed in `content`. + is_json = _is_json_content_type(response_content_type) + data: dict[str, Any] = {} + if is_json: + try: + data = json.loads(body) + except (UnicodeDecodeError, json.JSONDecodeError) as e: + raise LivepeerGatewayError( + f"HTTP JSON error: endpoint did not return valid JSON: {e} " + f"(url={runner_url}, content_type={response_content_type})" + ) from e + if not isinstance(data, dict): + raise LivepeerGatewayError( + f"Live runner call expected JSON object, got {type(data).__name__}" + ) + return LiveRunnerCallResult( + data, + runner_url=runner_url, + runner=runner, + session_id=( + session_id + or (data["session_id"].strip() if isinstance(data.get("session_id"), str) else "") + ), + payment_session=None if payment_type == "fixed" else payment_session, + content=None if is_json else body, + content_type=response_content_type, + ) + except LivepeerHTTPError as e: + if e.status_code != 402: + raise + if not signer_url: + raise LivepeerGatewayError("Live runner paid call requires signer_url") from e + if state is not None: + state.require_replayable(e) + challenge = _parse_runner_payment_challenge(e) + continue - raise LivepeerGatewayError("Live runner call exhausted payment challenge retries") + raise LivepeerGatewayError("Live runner call exhausted payment challenge retries") + finally: + if owned: + await state.aclose() def _parse_runner_payment_challenge(error: LivepeerHTTPError) -> LivePaymentChallenge: diff --git a/src/livepeer_gateway/runner_payload.py b/src/livepeer_gateway/runner_payload.py new file mode 100644 index 0000000..f127cd4 --- /dev/null +++ b/src/livepeer_gateway/runner_payload.py @@ -0,0 +1,272 @@ +"""Runner upload configuration and operation-scoped replay state.""" + +from __future__ import annotations + +import asyncio +import contextlib +import inspect +import io +import logging +import tempfile +from collections.abc import AsyncIterable, AsyncIterator, Callable +from dataclasses import dataclass +from typing import Any, BinaryIO + +from aiohttp.payload import Payload + +from ._http_continue import _cancel_writer +from .errors import LivepeerGatewayError + +_LOG = logging.getLogger(__name__) +_CHUNK_SIZE = 64 * 1024 + + +@dataclass(frozen=True, kw_only=True) +class StreamPayload: + """A caller-owned source with 100-continue and bounded replay buffering. + + Limits are decimal bytes. Overflow or storage failure disables later retries + without interrupting the upload. Seekable files rewind instead of buffering, + even with ``retry_buffer=False``. Bare files and iterables use these defaults. + """ + + source: AsyncIterable[bytes] | BinaryIO + expect100: bool = True + retry_buffer: bool = True + retry_buffer_memory_limit: int = 25_000_000 + retry_buffer_limit: int = 250_000_000 + + def __post_init__(self) -> None: + if not (0 < self.retry_buffer_memory_limit <= self.retry_buffer_limit): + raise ValueError("Replay buffer limits must satisfy 0 < memory limit <= total limit") + read = getattr(self.source, "read", None) + if isinstance(self.source, io.TextIOBase) or ( + not isinstance(self.source, AsyncIterable) + and (not callable(read) or inspect.iscoroutinefunction(read)) + ): + raise TypeError("StreamPayload requires a synchronous binary file or async byte iterable") + + +type RunnerPayload = ( + StreamPayload | AsyncIterable[bytes] | BinaryIO | dict[str, Any] + | bytes | bytearray | memoryview | str | None +) + + +class _ReplayUnavailableError(LivepeerGatewayError): + pass + + +async def _file_io[T](operation: Callable[[], T], *, threaded: bool = True) -> T: + """Join worker I/O before cancellation can hand its file to another attempt.""" + if not threaded: + return operation() + task = asyncio.create_task(asyncio.to_thread(operation)) + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + with contextlib.suppress(Exception): + await task + raise + + +class _StreamState: + def __init__(self, config: StreamPayload) -> None: + self.config = config + self.source = config.source + self.started = self.closed = self.eof = self.failed = False + self.cache_size = 0 + self.cache: tempfile.SpooledTemporaryFile[bytes] | None = None + self._iterator: AsyncIterator[bytes] | None = None + self._pending: asyncio.Task[bytes] | None = None + self._writer: asyncio.Task[None] | None = None + self._start_position: int | None = None + if not isinstance(self.source, AsyncIterable): + with contextlib.suppress(AttributeError, OSError, ValueError): + if self.source.seekable(): + self._start_position = self.source.tell() + self._buffering = config.retry_buffer and self._start_position is None + + @property + def replayable(self) -> bool: + if self.closed: + return False + source_open = not ( + getattr(self.source, "closed", False) + or (hasattr(self.source, "ag_frame") and self.source.ag_frame is None) + ) + if self._start_position is not None: + return source_open + return not self.failed and ( + (self.eof and self._buffering) + or (source_open and (not self.started or self._buffering)) + ) + + def require_replayable(self, cause: Exception | None = None) -> None: + if not self.replayable: + raise _ReplayUnavailableError( + "Cannot retry runner call: upload source is closed or its consumed data is no longer replayable" + ) from cause + + async def _drop_cache(self) -> None: + cache, self.cache = self.cache, None + self.cache_size = 0 + if cache is not None: + with contextlib.suppress(OSError): + await _file_io(cache.close, threaded=cache._rolled) + + async def _retain(self, chunk: bytes) -> None: + size = self.cache_size + len(chunk) + retained = False + storage_error = None + try: + if size > self.config.retry_buffer_limit: + self._buffering = False + await self._drop_cache() + return + if self.cache is None: + self.cache = tempfile.SpooledTemporaryFile( + max_size=self.config.retry_buffer_memory_limit, mode="w+b" + ) + cache = self.cache + on_disk = size > self.config.retry_buffer_memory_limit + + def append() -> None: + nonlocal retained, storage_error + try: + if on_disk: + cache.rollover() # Spill before a large chunk grows the RAM buffer. + cache.seek(self.cache_size) + cache.write(chunk) + retained = True + except OSError as exc: + storage_error = exc + raise + + await _file_io(append, threaded=on_disk) + except OSError as exc: + storage_error = exc + finally: + # _file_io joins a cancelled write. Commit it exactly once even if + # cancellation prevents chunks() from receiving the result. + if retained: + self.cache_size = size + self._pending = None + elif storage_error is not None: + self._buffering = False + await self._drop_cache() + _LOG.warning("Upload replay storage failed; continuing without replay", exc_info=storage_error) + + async def _read_source(self) -> bytes: + try: + if isinstance(self.source, AsyncIterable): + if self._iterator is None: + self._iterator = aiter(self.source) + try: + chunk = await anext(self._iterator) + except StopAsyncIteration: + self.eof = True + return b"" + else: + chunk = await _file_io(lambda: self.source.read(_CHUNK_SIZE)) + if not chunk: + self.eof = True + if not isinstance(chunk, bytes): + raise TypeError("Upload streams must produce bytes") + return chunk + except BaseException: + self.failed = True + raise + + async def chunks(self) -> AsyncIterator[bytes]: + self.require_replayable() + if self._start_position is not None and self.started: + if self._pending is not None: + with contextlib.suppress(Exception): + await asyncio.shield(self._pending) + self._pending = None + try: + await _file_io(lambda: self.source.seek(self._start_position)) + except (OSError, ValueError) as exc: + self._start_position = None + self.failed = True + raise _ReplayUnavailableError("Cannot rewind upload source for retry") from exc + self.eof = self.failed = False + + # Only this attempt touches the cache. The pending read cannot append or + # discard it while we replay the prefix, so no lock or offsets are needed. + position = 0 + while position < self.cache_size: + cache = self.cache + + def read_cached() -> bytes: + cache.seek(position) + return cache.read(min(_CHUNK_SIZE, self.cache_size - position)) + + try: + chunk = await _file_io(read_cached, threaded=cache._rolled) + except OSError as exc: + self._buffering = False + await self._drop_cache() + raise _ReplayUnavailableError("Cannot read upload replay buffer") from exc + position += len(chunk) + yield chunk + + while self._pending is not None or not self.eof: + if self._pending is None: + self.started = True + self._pending = asyncio.create_task(self._read_source()) + # A cancelled attempt leaves the result available to its successor. + chunk = await asyncio.shield(self._pending) + if chunk and self._buffering: + await self._retain(chunk) + self._pending = None + if chunk: + yield chunk + + async def aclose(self) -> None: + if self.closed: + return + self.closed = True + await self.finish_attempt() + if self._pending is not None: + self._pending.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await self._pending + self._pending = None + await self._drop_cache() + + async def finish_attempt(self) -> None: + """Join the HTTP writer without cancelling the shared source read.""" + writer, self._writer = self._writer, None + await _cancel_writer(writer) + + +def _stream_state(payload: RunnerPayload | _StreamState) -> _StreamState | None: + if isinstance(payload, _StreamState): + return payload + if isinstance(payload, StreamPayload): + return _StreamState(payload) + if isinstance(payload, AsyncIterable) or callable(getattr(payload, "read", None)): + return _StreamState(StreamPayload(source=payload)) + if payload is not None and not isinstance(payload, (dict, bytes, bytearray, memoryview, str)): + raise TypeError(f"Unsupported runner payload: {type(payload).__name__}") + return None + + +class _UploadBody(Payload): + """An attempt borrows the source and cache; aiohttp must close neither.""" + + _autoclose = True + + def __init__(self, state: _StreamState) -> None: + super().__init__(state, content_type="application/octet-stream") + self.state = state + + def decode(self, encoding: str = "utf-8", errors: str = "strict") -> str: + raise TypeError("Streaming request bodies cannot be decoded") + + async def write(self, writer: Any) -> None: + self.state._writer = asyncio.current_task() + async for chunk in self.state.chunks(): + await writer.write(chunk) diff --git a/src/livepeer_gateway/selection.py b/src/livepeer_gateway/selection.py index 588c8ec..0118198 100644 --- a/src/livepeer_gateway/selection.py +++ b/src/livepeer_gateway/selection.py @@ -21,6 +21,7 @@ _live_runner_price_info_from_json, call_runner, ) +from .runner_payload import RunnerPayload, _ReplayUnavailableError, _StreamState, _stream_state _LOG = logging.getLogger(__name__) @@ -38,14 +39,19 @@ def __init__( self, candidates: Sequence[LiveRunnerInstance], *, - body: dict[str, Any] | None = None, + body: RunnerPayload = None, + content_type: str | None = None, method: str = "POST", signer_url: str | None = None, signer_headers: dict[str, str] | None = None, timeout: float = 5.0, ) -> None: self._candidates = list(candidates) - self._body = dict(body or {}) + self._body = dict(body) if isinstance(body, dict) else body + self._content_type = content_type + self._payload_state: _StreamState | None = None + self._closed = False + self._next_running = False self._method = method self._signer_url = signer_url self._signer_headers = signer_headers @@ -58,47 +64,84 @@ def candidates(self) -> tuple[LiveRunnerInstance, ...]: return tuple(self._candidates) async def next(self) -> LiveRunnerCallResult: - while self._next_index < len(self._candidates): - runner = self._candidates[self._next_index] - self._next_index += 1 + if self._closed: + raise LivepeerGatewayError("Runner selection cursor is closed") + if self._next_running: + raise LivepeerGatewayError("Runner selection cursor.next() is already running") + self._next_running = True + try: + if self._payload_state is None: + self._payload_state = _stream_state(self._body) + while self._next_index < len(self._candidates): + if self._payload_state is not None: + self._payload_state.require_replayable() + runner = self._candidates[self._next_index] + self._next_index += 1 + try: + kwargs: dict[str, Any] = { + "runner": runner, + "payload": self._payload_state if self._payload_state is not None else self._body, + "method": self._method, + "timeout": self._timeout, + } + if self._content_type is not None: + kwargs["content_type"] = self._content_type + if self._signer_url is not None: + kwargs["signer_url"] = self._signer_url + if self._signer_headers is not None: + kwargs["signer_headers"] = self._signer_headers + result = await call_runner(**kwargs) + except _ReplayUnavailableError: + raise + except Exception as e: + reason = str(e) + _LOG.debug( + "select_runner candidate failed: %s (%s)", + runner.url, + reason, + ) + self.rejections.append(RunnerRejection(url=runner.url, reason=reason)) + if self._payload_state is not None: + self._payload_state.require_replayable(e) + continue + + _LOG.debug("select_runner selected: %s", runner.url) + return result + + _LOG.debug( + "select_runner failed: all %d runners rejected", + len(self._candidates), + ) + raise NoRunnerAvailableError( + f"All runners failed ({len(self.rejections)} tried)", + rejections=list(self.rejections), + ) + except BaseException: try: - kwargs: dict[str, Any] = { - "runner": runner, - "payload": self._body, - "method": self._method, - "timeout": self._timeout, - } - if self._signer_url is not None: - kwargs["signer_url"] = self._signer_url - if self._signer_headers is not None: - kwargs["signer_headers"] = self._signer_headers - result = await call_runner(**kwargs) - except Exception as e: - reason = str(e) - _LOG.debug( - "select_runner candidate failed: %s (%s)", - runner.url, - reason, - ) - self.rejections.append(RunnerRejection(url=runner.url, reason=reason)) - continue + await self.aclose() + except Exception: + _LOG.warning("Runner cursor cleanup failed", exc_info=True) + raise + finally: + self._next_running = False - _LOG.debug("select_runner selected: %s", runner.url) - return result + async def aclose(self) -> None: + """Release the replay cache without closing the caller's upload source.""" + self._closed = True + if self._payload_state is not None: + await self._payload_state.aclose() - _LOG.debug( - "select_runner failed: all %d runners rejected", - len(self._candidates), - ) - raise NoRunnerAvailableError( - f"All runners failed ({len(self.rejections)} tried)", - rejections=list(self.rejections), - ) + async def __aenter__(self) -> RunnerSelectionCursor: + return self + + async def __aexit__(self, *exc: object) -> None: + await self.aclose() async def runner_selector( *, - body: dict[str, Any] | None = None, + body: RunnerPayload = None, + content_type: str | None = None, method: str = "POST", orchestrators: Sequence[str] | str | None = None, signer_url: str | None = None, @@ -134,6 +177,7 @@ async def runner_selector( return RunnerSelectionCursor( candidates, body=body, + content_type=content_type, method=method, signer_url=signer_url, signer_headers=signer_headers, diff --git a/tests/test_call_runner_uploads.py b/tests/test_call_runner_uploads.py new file mode 100644 index 0000000..b8e9364 --- /dev/null +++ b/tests/test_call_runner_uploads.py @@ -0,0 +1,639 @@ +from __future__ import annotations + +import asyncio +import io +import json +from contextlib import ExitStack, asynccontextmanager +from types import SimpleNamespace +from unittest import mock + +import pytest + +from livepeer_gateway import ( + LiveRunnerInstance, + RunnerSelectionCursor, + StreamPayload, + call_runner, +) +from livepeer_gateway import live_runner +from livepeer_gateway.errors import ( + LivepeerGatewayError, + LivepeerHTTPError, + NoRunnerAvailableError, +) + + +CHALLENGE = { + "payment_params": "opaque", + "manifest_id": "session", + "payment_url": "http://payment.example/session", +} + + +@pytest.fixture +def signer(monkeypatch): + monkeypatch.setattr(live_runner, "get_signer_info", mock.AsyncMock(return_value=SimpleNamespace(address="payer"))) + payment = mock.AsyncMock(return_value=(None, SimpleNamespace(payment="paid", seg_creds="segment"))) + monkeypatch.setattr(live_runner, "_get_runner_payment", payment) + return payment + + +@pytest.fixture +async def no_loop_errors(): + """A failed asyncio callback is logged, rather than failing the awaiting task.""" + loop = asyncio.get_running_loop() + previous = loop.get_exception_handler() + errors = [] + loop.set_exception_handler(lambda _, context: errors.append(context)) + try: + yield + # Future cancellation can queue callbacks after request cleanup returns. + await asyncio.sleep(0) + await asyncio.sleep(0) + finally: + loop.set_exception_handler(previous) + assert not errors, errors + + +@asynccontextmanager +async def serve(handler): + tasks = set() + errors = [] + + async def connected(reader, writer): + task = asyncio.current_task() + tasks.add(task) + try: + await handler(reader, writer) + except (ConnectionError, asyncio.IncompleteReadError): + pass + except Exception as exc: + errors.append(exc) + finally: + writer.close() + try: + await writer.wait_closed() + except ConnectionError: + pass + tasks.discard(task) + + server = await asyncio.start_server(connected, "127.0.0.1", 0) + async with server: + port = server.sockets[0].getsockname()[1] + try: + yield f"http://127.0.0.1:{port}" + finally: + if tasks: + await asyncio.wait_for(asyncio.gather(*tasks), 2) + if errors: + raise errors[0] + + +async def headers(reader): + raw = await reader.readuntil(b"\r\n\r\n") + lines = raw.decode().split("\r\n") + return dict(line.split(": ", 1) for line in lines[1:] if line) + + +async def body(reader, request_headers): + if request_headers.get("Transfer-Encoding") == "chunked": + chunks = [] + while True: + chunk = await read_chunk(reader) + if not chunk: + return b"".join(chunks) + chunks.append(chunk) + return await reader.readexactly(int(request_headers.get("Content-Length", 0))) + + +async def read_chunk(reader): + line = await reader.readline() + if not line: + raise asyncio.IncompleteReadError(b"", 1) + size = int(line, 16) + chunk = await reader.readexactly(size) + assert await reader.readexactly(2) == b"\r\n" + return chunk + + +async def respond(writer, status=200, data=None, *, fragment=False, location=None): + data = {"ok": True} if data is None else data + raw = json.dumps(data).encode() + redirect = f"Location: {location}\r\n" if location else "" + head = ( + f"HTTP/1.1 {status} Result\r\nContent-Type: application/json\r\n" + f"Content-Length: {len(raw)}\r\n{redirect}Connection: close\r\n\r\n".encode() + ) + writer.write(head if fragment else head + raw) + await writer.drain() + if fragment: + await asyncio.sleep(0.02) + writer.write(raw) + await writer.drain() + + +async def allow_upload(writer): + writer.write(b"HTTP/1.1 100 Continue\r\n\r\n") + await writer.drain() + + +class Source: + def __init__(self, *chunks): + self.chunks = iter(chunks) + self.iterations = 0 + self.reads = 0 + self.closed = False + + def __aiter__(self): + self.iterations += 1 + return self + + async def __anext__(self): + self.reads += 1 + try: + return next(self.chunks) + except StopIteration: + raise StopAsyncIteration from None + + async def aclose(self): + self.closed = True + + +def candidate(url): + return LiveRunnerInstance(url, "test", "runner", "single-shot", url, {}) + + +CUSTOM_TYPE = "application/vnd.test+json; charset=utf-8" + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("kwargs, expected, content_type", [ + ({}, b"", None), + ({"payload": None}, b"", None), + ({"payload": {}}, b"{}", "application/json"), + ({"payload": {"x": 1}}, b'{"x": 1}', "application/json"), + ({"payload": "héllo"}, "héllo".encode(), "application/octet-stream"), + ({"payload": b"\x00\xff"}, b"\x00\xff", "application/octet-stream"), + ({"payload": bytearray(b"abc")}, b"abc", "application/octet-stream"), + ({"payload": memoryview(b"abc")}, b"abc", "application/octet-stream"), + ({"payload": b""}, b"", "application/octet-stream"), + ({"payload": ""}, b"", "application/octet-stream"), + ({"payload": {"x": 1}, "content_type": CUSTOM_TYPE}, b'{"x": 1}', CUSTOM_TYPE), + ({"payload": b"abc", "content_type": CUSTOM_TYPE}, b"abc", CUSTOM_TYPE), + ({"payload": "abc", "content_type": CUSTOM_TYPE}, b"abc", CUSTOM_TYPE), + ({"payload": None, "content_type": CUSTOM_TYPE}, b"", CUSTOM_TYPE), +]) +async def test_request_bodies(kwargs, expected, content_type, stream): + async def handler(reader, writer): + request_headers = await headers(reader) + assert "Expect" not in request_headers + assert request_headers["Accept"] == "*/*" + assert request_headers.get("Content-Type") == content_type + assert await body(reader, request_headers) == expected + await respond(writer) + + async with serve(handler) as url: + result = await call_runner(url, stream=stream, **kwargs) + if stream: + async with result: + assert json.loads(b"".join([chunk async for chunk in result.aiter_bytes()])) == {"ok": True} + else: + assert result.data == {"ok": True} + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_early_payment_challenge_does_not_touch_source( + stream, signer, no_loop_errors +): + source = Source(b"one", b"two") + attempts = 0 + + async def handler(reader, writer): + nonlocal attempts + request_headers = await headers(reader) + attempts += 1 + assert request_headers["Expect"] == "100-continue" + if attempts == 1: + assert source.iterations == source.reads == 0 + await respond(writer, 402, CHALLENGE, fragment=True) + assert await reader.read() == b"" + return + assert request_headers["Livepeer-Payment"] == "paid" + assert source.iterations == source.reads == 0 + await allow_upload(writer) + assert await body(reader, request_headers) == b"onetwo" + await respond(writer) + + async with serve(handler) as url: + # Cancelling the upload writer as soon as final headers arrive can + # discard a fragmented 402 body on newer aiohttp. Bound the whole test + # so a regression fails rather than waiting forever for payment JSON. + result = await asyncio.wait_for( + call_runner( + url, + payload=StreamPayload(source=source, retry_buffer=False), + signer_url="http://signer.example", + payment_unit="fixed", + stream=stream, + ), + 2, + ) + if stream: + await result.aclose() + assert attempts == 2 + assert source.reads == 3 + assert not source.closed + signer.assert_awaited_once() + challenge = signer.call_args.args[0] + assert challenge.payment_params == CHALLENGE["payment_params"] + assert challenge.manifest_id == CHALLENGE["manifest_id"] + assert challenge.payment_url == CHALLENGE["payment_url"] + + +@pytest.mark.parametrize("status,kind", [(402, "nonseekable"), (503, "iterator"), (307, "file")]) +async def test_early_rejection_preserves_source_for_next_candidate(tmp_path, status, kind): + class NonSeekable(io.BytesIO): + def seekable(self): + return False + + with ExitStack() as sources: + if kind == "file": + path = tmp_path / "upload.bin" + path.write_bytes(b"skip-body") + source = sources.enter_context(path.open("rb")) + source.seek(5) + elif kind == "nonseekable": + source = sources.enter_context(NonSeekable(b"body")) + else: + source = Source(b"bo", b"dy") + + def assert_unread(): + assert not source.closed + if isinstance(source, Source): + assert source.iterations == source.reads == 0 + else: + assert source.tell() == (5 if kind == "file" else 0) + + async def reject(reader, writer): + await headers(reader) + assert_unread() + await respond(writer, status, location=second if status == 307 else None) + assert await reader.read() == b"" + + async def accept(reader, writer): + request_headers = await headers(reader) + assert_unread() + await allow_upload(writer) + assert await body(reader, request_headers) == b"body" + await respond(writer) + + async with serve(reject) as first, serve(accept) as second: + payload = source if kind == "file" else StreamPayload(source=source, retry_buffer=False) + async with RunnerSelectionCursor([candidate(first), candidate(second)], body=payload) as cursor: + result = await cursor.next() + assert result.data == {"ok": True} and result.runner_url == second + assert len(cursor.rejections) == 1 + assert not source.closed + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_expect_timeout_aborts_without_consumption(stream, no_loop_errors): + source = Source(b"never") + + async def handler(reader, writer): + await headers(reader) + assert await reader.read() == b"" + + async with serve(handler) as url: + with pytest.raises(LivepeerGatewayError): + await asyncio.wait_for( + call_runner(url, payload=source, stream=stream, timeout=0.05), 1 + ) + assert source.iterations == source.reads == 0 + assert not source.closed + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_expect_can_be_disabled(stream): + source = Source(b"direct") + + async def handler(reader, writer): + request_headers = await headers(reader) + assert "Expect" not in request_headers + assert await body(reader, request_headers) == b"direct" + await respond(writer) + + async with serve(handler) as url: + result = await call_runner( + url, payload=StreamPayload(source=source, expect100=False), stream=stream + ) + if stream: + await result.aclose() + assert not source.closed + + +async def test_buffered_upload_can_be_reused_by_cursor_until_close(): + source = Source(b"one", b"two") + received = [] + + async def handler(reader, writer): + request_headers = await headers(reader) + assert request_headers["Content-Type"] == "video/test" + await allow_upload(writer) + received.append(await body(reader, request_headers)) + await respond(writer) + + async with serve(handler) as url: + async with RunnerSelectionCursor( + [candidate(url), candidate(url)], body=source, content_type="video/test" + ) as cursor: + await cursor.next() + state = cursor._payload_state + cache = state.cache + assert cache is not None and not cache.closed + await cursor.next() + assert cache.closed and state.closed + with pytest.raises(LivepeerGatewayError, match="closed"): + await cursor.next() + assert received == [b"onetwo", b"onetwo"] + assert source.reads == 3 and not source.closed + + +async def test_seekable_file_replays_from_initial_position_without_buffering(): + source = io.BytesIO(b"prefix-data") + source.seek(7) + received = [] + + async def handler(reader, writer): + request_headers = await headers(reader) + await allow_upload(writer) + received.append(await body(reader, request_headers)) + await respond(writer, 503 if len(received) == 1 else 200) + + async with serve(handler) as url: + async with RunnerSelectionCursor( + [candidate(url), candidate(url)], + body=StreamPayload(source=source, retry_buffer=False), + ) as cursor: + await cursor.next() + assert cursor._payload_state.cache is None + assert received == [b"data", b"data"] + assert not source.closed + + +@pytest.mark.parametrize("options", [ + {"retry_buffer": False}, + {"retry_buffer_memory_limit": 3, "retry_buffer_limit": 6}, +]) +async def test_unreplayable_upload_finishes_but_stops_failover(options): + source = Source(b"prefix", b"over-limit", b"tail") + attempts = 0 + + async def handler(reader, writer): + nonlocal attempts + attempts += 1 + request_headers = await headers(reader) + await allow_upload(writer) + assert await body(reader, request_headers) == b"prefixover-limittail" + await respond(writer, 503) + + async with serve(handler) as url: + cursor = RunnerSelectionCursor( + [candidate(url), candidate(url)], body=StreamPayload(source=source, **options) + ) + with pytest.raises(LivepeerGatewayError, match="no longer replayable") as error: + await cursor.next() + assert isinstance(error.value.__cause__, LivepeerHTTPError) + assert cursor._payload_state.closed and cursor._payload_state.cache is None + assert attempts == 1 and not source.closed + + +@pytest.mark.parametrize("success", [False, True]) +async def test_exhausted_cursor_releases_cache(success): + source = Source(b"data") + + async def handler(reader, writer): + request_headers = await headers(reader) + await allow_upload(writer) + await body(reader, request_headers) + await respond(writer, 200 if success else 503) + + async with serve(handler) as url: + cursor = RunnerSelectionCursor([candidate(url)], body=StreamPayload(source=source, retry_buffer=not success)) + if success: + await cursor.next() + with pytest.raises(NoRunnerAvailableError): + await cursor.next() + assert cursor._payload_state.closed + assert cursor._payload_state.cache is None + assert not source.closed + + +async def test_disconnected_upload_failover_adopts_pending_source_read(): + read_started = asyncio.Event() + finish_read = asyncio.Event() + source_cancelled = False + source_reads = [] + rejections = 0 + + async def source(): + nonlocal source_cancelled + try: + source_reads.append(b"prefix") + yield b"prefix" + read_started.set() + await finish_read.wait() + source_reads.append(b"middle") + yield b"middle" + source_reads.append(b"tail") + yield b"tail" + except asyncio.CancelledError: + source_cancelled = True + raise + + async def reject(reader, writer): + nonlocal rejections + rejections += 1 + await headers(reader) + await allow_upload(writer) + assert await read_chunk(reader) == b"prefix" + await read_started.wait() + + async def accept(reader, writer): + request_headers = await headers(reader) + assert not source_cancelled + await allow_upload(writer) + assert await read_chunk(reader) == b"prefix" + finish_read.set() + assert await body(reader, request_headers) == b"middletail" + await respond(writer) + + async with serve(reject) as first, serve(accept) as second: + async with RunnerSelectionCursor( + [candidate(first), candidate(second)], body=source(), method="PUT" + ) as cursor: + assert (await asyncio.wait_for(cursor.next(), 2)).data == {"ok": True} + assert not source_cancelled + assert source_reads == [b"prefix", b"middle", b"tail"] + assert rejections == 1 + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_partial_upload_payment_retry_replays_prefix(stream, signer): + read_started = asyncio.Event() + finish_read = asyncio.Event() + attempts = 0 + + async def source(): + yield b"prefix" + read_started.set() + await finish_read.wait() + yield b"tail" + + async def handler(reader, writer): + nonlocal attempts + request_headers = await headers(reader) + attempts += 1 + await allow_upload(writer) + assert await read_chunk(reader) == b"prefix" + if attempts == 1: + await read_started.wait() + await respond(writer, 402, CHALLENGE, fragment=True) + else: + assert request_headers["Livepeer-Payment"] == "paid" + finish_read.set() + assert await body(reader, request_headers) == b"tail" + await respond(writer) + + async with serve(handler) as url: + result = await asyncio.wait_for( + call_runner( + url, + payload=source(), + signer_url="http://signer.example", + payment_unit="fixed", + stream=stream, + ), + 2, + ) + if stream: + await result.aclose() + signer.assert_awaited_once() + assert attempts == 2 + + +@pytest.mark.parametrize("mode", ["buffered", "stream", "cursor"]) +async def test_unbuffered_consumed_body_is_not_paid_or_retried(mode, signer): + source = Source(b"data") + + async def handler(reader, writer): + request_headers = await headers(reader) + await allow_upload(writer) + assert await body(reader, request_headers) == b"data" + await respond(writer, 402, CHALLENGE) + + async with serve(handler) as url: + with pytest.raises( + LivepeerGatewayError, match="no longer replayable" + ) as error: + payload = StreamPayload(source=source, retry_buffer=False) + if mode == "cursor": + await RunnerSelectionCursor([candidate(url)], body=payload, signer_url="http://signer.example").next() + else: + await call_runner(url, payload=payload, signer_url="http://signer.example", payment_unit="fixed", stream=mode == "stream") + assert isinstance(error.value.__cause__, LivepeerHTTPError) + signer.assert_not_called() + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_upload_cancellation_releases_cache(stream): + read_started = asyncio.Event() + + async def source(): + yield b"prefix" + read_started.set() + await asyncio.Future() + + async def handler(reader, writer): + await headers(reader) + await allow_upload(writer) + assert await read_chunk(reader) == b"prefix" + await reader.read() + + states = [] + make_state = live_runner._stream_state + + def capture_state(payload): + state = make_state(payload) + states.append(state) + return state + + with mock.patch.object(live_runner, "_stream_state", side_effect=capture_state): + async with serve(handler) as url: + task = asyncio.create_task( + call_runner(url, payload=source(), stream=stream) + ) + await read_started.wait() + cache = states[0].cache + assert cache is not None + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert states[0].closed and cache.closed + assert states[0]._writer is states[0]._pending is None + + +async def test_response_stream_owns_active_upload_until_close(): + read_started = asyncio.Event() + response_closed = asyncio.Event() + + async def source(): + yield b"prefix" + read_started.set() + await asyncio.Future() + + async def handler(reader, writer): + await headers(reader) + await allow_upload(writer) + assert await read_chunk(reader) == b"prefix" + await read_started.wait() + # Send response headers while the request source is still live. + writer.write(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n") + await writer.drain() + await reader.read() + response_closed.set() + + async with serve(handler) as url: + result = await call_runner(url, payload=source(), stream=True) + state = result._payload_state + cache = state.cache + assert not state.closed + await result.aclose() + assert state.closed and cache.closed + await asyncio.wait_for(response_closed.wait(), 1) + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_response_mode_controls_timeout_after_continue(stream): + async def source(): + await asyncio.sleep(0.1) + yield b"delayed" + + async def handler(reader, writer): + request_headers = await headers(reader) + # Native aiohttp also accepts 103 as permission; exercise that here + # while checking that permission ends only the continue-wait deadline. + writer.write(b"HTTP/1.1 103 Early Hints\r\n\r\n") + await writer.drain() + assert await body(reader, request_headers) == b"delayed" + await respond(writer) + + async with serve(handler) as url: + if stream: + async with await call_runner( + url, payload=source(), stream=True, timeout=0.05 + ) as result: + assert json.loads(b"".join([chunk async for chunk in result.aiter_bytes()])) == {"ok": True} + else: + with pytest.raises(LivepeerGatewayError): + await call_runner(url, payload=source(), timeout=0.05) diff --git a/tests/test_live_runner.py b/tests/test_live_runner.py index 3d542fd..e54ba4f 100644 --- a/tests/test_live_runner.py +++ b/tests/test_live_runner.py @@ -254,7 +254,7 @@ async def test_call_runner_can_attach_runner_instance(self) -> None: runner_id="runner-1", mode="persistent", orchestrator_url="https://service.example.com", - raw={"label": "echo"}, + raw={"app": "livepeer-sample/echo"}, ) def _request_body( diff --git a/tests/test_runner_payload.py b/tests/test_runner_payload.py new file mode 100644 index 0000000..6064a69 --- /dev/null +++ b/tests/test_runner_payload.py @@ -0,0 +1,258 @@ +from __future__ import annotations + +import asyncio +import io +import tempfile +import threading +from unittest import mock + +import pytest + +from livepeer_gateway import StreamPayload +from livepeer_gateway.errors import LivepeerGatewayError +from livepeer_gateway.runner_payload import _StreamState + + +@pytest.fixture +async def make_state(): + states = [] + + def make(source, **options): + state = _StreamState(StreamPayload(source=source, **options)) + states.append(state) + return state + + yield make + for state in states: + await state.aclose() + + +async def _chunks(*chunks): + for chunk in chunks: + yield chunk + + +async def collect(state): + return b"".join([chunk async for chunk in state.chunks()]) + + +@pytest.mark.parametrize("memory,total", [(0, 1), (1, 0), (2, 1), (-1, 10)]) +def test_invalid_buffer_limits(memory, total): + with pytest.raises(ValueError, match="limits"): + StreamPayload( + source=io.BytesIO(), + retry_buffer_memory_limit=memory, + retry_buffer_limit=total, + ) + + +def test_text_files_are_rejected(): + with pytest.raises(TypeError, match="binary"): + StreamPayload(source=io.StringIO("text")) + + +async def test_spills_and_replays_without_rereading_source(make_state): + chunks = [b"one", b"two", b"three"] + reads = [] + + async def source(): + for chunk in chunks: + reads.append(chunk) + yield chunk + + state = make_state(source(), retry_buffer_memory_limit=4, retry_buffer_limit=100) + assert await collect(state) == b"onetwothree" + cache = state.cache + assert cache._rolled + assert await collect(state) == b"onetwothree" + assert reads == chunks + await state.aclose() + assert cache.closed + + +@pytest.mark.parametrize("length,replayable", [(5, True), (6, False)]) +async def test_total_limit_is_checked_before_writing(length, replayable, make_state): + state = make_state(_chunks(b"x" * length), retry_buffer_memory_limit=2, retry_buffer_limit=5) + assert await collect(state) == b"x" * length + assert state.replayable is replayable + if replayable: + assert state.cache_size == 5 + assert await collect(state) == b"x" * length + else: + assert state.cache is None + with pytest.raises(LivepeerGatewayError, match="no longer replayable"): + await collect(state) + + +async def test_cancelled_attempt_replays_prefix_and_adopts_pending_read(make_state): + read_started = asyncio.Event() + finish_read = asyncio.Event() + reads = [] + + async def source(): + for chunk in (b"prefix", b"middle", b"end"): + reads.append(chunk) + if chunk == b"middle": + read_started.set() + await finish_read.wait() + yield chunk + + state = make_state(source()) + first = state.chunks() + assert await anext(first) == b"prefix" + upload = asyncio.create_task(anext(first)) + await read_started.wait() + upload.cancel() + with pytest.raises(asyncio.CancelledError): + await upload + assert not state._pending.cancelled() + # More than one interrupted attempt may replay the same prefix while + # the *one* source read remains pending. + second = state.chunks() + assert await anext(second) == b"prefix" + await second.aclose() + third = state.chunks() + assert await anext(third) == b"prefix" + finish_read.set() + assert b"".join([chunk async for chunk in third]) == b"middleend" + assert reads == [b"prefix", b"middle", b"end"] + + +async def test_pending_overflow_preserves_prefix_for_current_retry(make_state): + read_started = asyncio.Event() + finish_read = asyncio.Event() + prefix = b"x" * (64 * 1024 + 1) + + async def source(): + yield prefix + read_started.set() + await finish_read.wait() + yield b"new-data" + yield b"remainder" + + state = make_state(source(), retry_buffer_memory_limit=4, retry_buffer_limit=len(prefix)) + first = state.chunks() + assert await anext(first) == prefix + task = asyncio.create_task(anext(first)) + await read_started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + retry = state.chunks() + assert await anext(retry) == prefix[:-1] + finish_read.set() + # The pending chunk may complete before this retry finishes its cached prefix. + await asyncio.shield(state._pending) + assert b"".join([chunk async for chunk in retry]) == b"xnew-dataremainder" + assert not state.replayable + assert state.cache is None + + +@pytest.mark.parametrize("failure", ["create", "write", "rollover"]) +async def test_cache_failure_allows_current_upload_to_finish(failure, make_state): + state = make_state(_chunks(b"first", b"second"), retry_buffer_memory_limit=6, retry_buffer_limit=100) + factory = tempfile.SpooledTemporaryFile + + def make_cache(*args, **kwargs): + if failure == "create": + raise OSError("no space") + cache = factory(*args, **kwargs) + setattr(cache, failure, mock.Mock(side_effect=OSError("no space"))) + return cache + + with mock.patch( + "livepeer_gateway.runner_payload.tempfile.SpooledTemporaryFile", + side_effect=make_cache, + ): + assert await collect(state) == b"firstsecond" + assert not state.replayable + assert state.cache is None + + +async def test_failed_source_is_not_retried(make_state): + async def source(): + yield b"prefix" + raise RuntimeError("source failed") + + state = make_state(source()) + with pytest.raises(RuntimeError, match="source failed"): + await collect(state) + assert not state.replayable + + +async def test_closed_unread_file_is_not_retried(make_state): + file = io.BytesIO(b"data") + state = make_state(file) + file.close() + assert not state.replayable + with pytest.raises(LivepeerGatewayError, match="closed"): + await collect(state) + + +async def test_non_bytes_source_fails_clearly(make_state): + state = make_state(_chunks("text")) + with pytest.raises(TypeError, match="produce bytes"): + await collect(state) + + +async def test_replay_read_failure_disables_future_retries(make_state): + state = make_state(_chunks(b"data")) + assert await collect(state) == b"data" + state.cache.read = mock.Mock(side_effect=OSError("storage unavailable")) + with pytest.raises( + LivepeerGatewayError, match="Cannot read upload replay buffer" + ): + await collect(state) + assert not state.replayable and state.cache is None + + +async def test_failed_rewind_disables_future_retries(make_state): + source = io.BytesIO(b"data") + state = make_state(source, retry_buffer=False) + assert await collect(state) == b"data" + source.seek = mock.Mock(side_effect=OSError("seek failed")) + with pytest.raises(LivepeerGatewayError, match="Cannot rewind upload source"): + await collect(state) + assert not state.replayable + source.close() + + +@pytest.mark.parametrize("fail", [False, True]) +async def test_cancellation_joins_cache_write_before_retry(make_state, fail): + state = make_state(_chunks(b"prefix", b"middle", b"tail"), retry_buffer_memory_limit=2) + attempt = state.chunks() + assert await anext(attempt) == b"prefix" + cache = state.cache + entered, finish = asyncio.Event(), threading.Event() + loop = asyncio.get_running_loop() + write = cache.write + + def blocked_write(chunk): + loop.call_soon_threadsafe(entered.set) + assert finish.wait(2), "test did not release the cache write" + if fail: + raise OSError("disk full") + return write(chunk) + + with mock.patch.object(cache, "write", side_effect=blocked_write): + upload = asyncio.create_task(anext(attempt)) + try: + await asyncio.wait_for(entered.wait(), 1) + upload.cancel() + await asyncio.sleep(0) + assert not upload.done() + finally: + finish.set() + with pytest.raises(asyncio.CancelledError): + await upload + if fail: + assert not state.replayable and cache.closed + else: + assert await collect(state) == b"prefixmiddletail" + + +async def test_memory_cache_does_not_dispatch_file_io(make_state): + state = make_state(_chunks(b"one", b"two")) + with mock.patch("asyncio.to_thread", side_effect=AssertionError("unexpected disk I/O")): + assert await collect(state) == await collect(state) == b"onetwo" + await state.aclose() diff --git a/tests/test_selection.py b/tests/test_selection.py index e56d371..863bfd3 100644 --- a/tests/test_selection.py +++ b/tests/test_selection.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio from unittest import mock import pytest @@ -51,7 +52,7 @@ async def _call_runner( assert candidate.runner_id == "runner-c" assert candidate.mode == "persistent" assert candidate.orchestrator_url == "https://orch-b" - assert candidate.raw["label"] == "runner-c-label" + assert candidate.raw["metadata"] == "runner-c-metadata" assert candidate.price_info is not None assert candidate.price_info.price == 25 assert candidate.price_info.currency == "wei" @@ -499,7 +500,7 @@ def _discovery_entries() -> list[dict[str, object]]: "app": "app-c", "runner_id": "runner-c", "mode": "persistent", - "label": "runner-c-label", + "metadata": "runner-c-metadata", "price_info": { "price": 25, "currency": "wei", @@ -510,3 +511,68 @@ def _discovery_entries() -> list[dict[str, object]]: address="https://orch-b", ), ] + + +@pytest.mark.parametrize("body", [123, mock.Mock(read=mock.AsyncMock())]) +async def test_invalid_body_fails_before_attempting_candidates(body): + runner = LiveRunnerInstance("http://unused", "test", "runner", "single-shot", "http://unused", {}) + cursor = selection.RunnerSelectionCursor([runner, runner], body=body) + with mock.patch.object(selection, "call_runner") as call: + with pytest.raises(TypeError): + await cursor.next() + call.assert_not_called() + assert cursor.rejections == [] + + +async def test_cursor_cleanup_failure_preserves_original_error(): + cursor = selection.RunnerSelectionCursor([]) + with mock.patch.object(cursor, "aclose", side_effect=RuntimeError("cleanup failed")): + with pytest.raises(NoRunnerAvailableError): + await cursor.next() + + +async def test_overlapping_next_does_not_interrupt_or_advance_cursor(): + runner = LiveRunnerInstance("http://unused", "test", "runner", "single-shot", "http://unused", {}) + result = LiveRunnerCallResult({"ok": True}, runner_url=runner.url, runner=runner) + started, finish = asyncio.Event(), asyncio.Event() + + async def call(**kwargs): + started.set() + await finish.wait() + return result + + async with selection.RunnerSelectionCursor([runner, runner]) as cursor: + with mock.patch.object(selection, "call_runner", side_effect=call) as request: + active = asyncio.create_task(cursor.next()) + try: + await asyncio.wait_for(started.wait(), 1) + with pytest.raises(LivepeerGatewayError, match="already running"): + await asyncio.wait_for(cursor.next(), 1) + assert not active.done() + finally: + finish.set() + assert await active is result + assert await cursor.next() is result + assert request.await_count == 2 + + +async def test_cursor_cancellation_closes_upload_cache(): + async def source(): + yield b"data" + + started = asyncio.Event() + cursor = selection.RunnerSelectionCursor([LiveRunnerInstance("http://unused", "test", "runner", "single-shot", "http://unused", {})], body=source()) + + async def call(**kwargs): + assert await anext(kwargs["payload"].chunks()) == b"data" + started.set() + await asyncio.Future() + + with mock.patch.object(selection, "call_runner", side_effect=call): + attempt = asyncio.create_task(cursor.next()) + await asyncio.wait_for(started.wait(), 1) + cache = cursor._payload_state.cache + attempt.cancel() + with pytest.raises(asyncio.CancelledError): + await attempt + assert cache.closed and cursor._payload_state.closed