From ad4e926eb7b62c7dd8492454e26d222a4ba8001f Mon Sep 17 00:00:00 2001 From: Josh Allmann Date: Fri, 16 Oct 2026 12:07:48 -0700 Subject: [PATCH 1/2] Remove unused Label field --- CHANGELOG.md | 1 + examples/text/runners.toml | 1 - src/livepeer_gateway/live_runner.py | 6 ------ tests/test_live_runner.py | 2 +- tests/test_selection.py | 4 ++-- 5 files changed, 4 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a067f9f..39b0fa9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,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/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/src/livepeer_gateway/live_runner.py b/src/livepeer_gateway/live_runner.py index 8d80b3b..a9755c0 100644 --- a/src/livepeer_gateway/live_runner.py +++ b/src/livepeer_gateway/live_runner.py @@ -274,7 +274,6 @@ def __init__( runner_id: str = "", mode: str = "persistent", proxy: bool = False, - label: str = "", version: str = "", metadata: str = "", status: str = "ready", @@ -298,7 +297,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 +443,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 +605,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 +629,6 @@ async def register_runner( runner_id=runner_id, mode=mode, proxy=proxy, - label=label, version=version, metadata=metadata, status=status, 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_selection.py b/tests/test_selection.py index e56d371..ae48af8 100644 --- a/tests/test_selection.py +++ b/tests/test_selection.py @@ -51,7 +51,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 +499,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", From d430f6e896f8c9b161ffe08b1f75e45e1ef46139 Mon Sep 17 00:00:00 2001 From: Josh Allmann Date: Fri, 16 Oct 2026 18:57:09 -0700 Subject: [PATCH 2/2] Support streamable runner payloads Accept dicts, text, bytes, binary files, and async byte streams with an optional content type. Send omitted and None payloads as empty bodies. Add StreamPayload(source=...) for streaming data sources. To accommodate 402 (payment) and other error response codes without consuming the payload, the client sends an `Expect: 100-continue` and waits for the 100 before transmitting the body. This is something that most HTTP servers do by default already, including golang's. Bounded (non-streaming) inputs do not wait for the 100. In the case that there is still an error after payment, streaming input is buffered for replay during runner failover (optional, but enabled by default). Seekable files are rewound instead of buffered. This is API-compatible with existing callers, with a few tweaks: * `None` payloads now don't send a body, instead of `{}` * `stream=True` (response streaming) now honors `timeout` instead of using a fixed 10s Also remove an unused `Label` field from runner registration. --- CHANGELOG.md | 4 + README | 4 +- docs/runner-payloads.md | 193 ++++++++ pyproject.toml | 1 + src/livepeer_gateway/__init__.py | 3 + src/livepeer_gateway/_http_continue.py | 54 +++ src/livepeer_gateway/http.py | 139 ++++-- src/livepeer_gateway/live_runner.py | 288 ++++++----- src/livepeer_gateway/runner_payload.py | 272 +++++++++++ src/livepeer_gateway/selection.py | 116 +++-- tests/test_call_runner_uploads.py | 639 +++++++++++++++++++++++++ tests/test_runner_payload.py | 258 ++++++++++ tests/test_selection.py | 66 +++ 13 files changed, 1829 insertions(+), 208 deletions(-) create mode 100644 docs/runner-payloads.md create mode 100644 src/livepeer_gateway/_http_continue.py create mode 100644 src/livepeer_gateway/runner_payload.py create mode 100644 tests/test_call_runner_uploads.py create mode 100644 tests/test_runner_payload.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 39b0fa9..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 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/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 a9755c0..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 @@ -749,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 = ..., @@ -766,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 = ..., @@ -782,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, @@ -801,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_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 ae48af8..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 @@ -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