From 7e72d74d019badf7057529fb7bc9d841c8cc6f4c Mon Sep 17 00:00:00 2001 From: Zac Farrell Date: Tue, 30 Jun 2026 20:39:53 -0700 Subject: [PATCH 1/3] feat(uploads): harden multipart upload resilience --- hotdata/_upload_watchdog.py | 462 ++++++++++++++++ hotdata/uploads.py | 647 ++++++++++++++++++++--- tests/test_upload_codex_fixes.py | 427 +++++++++++++++ tests/test_upload_multipart_resilient.py | 354 +++++++++++++ tests/test_upload_resilient.py | 233 ++++++++ tests/test_upload_retry_policy.py | 201 +++++++ tests/test_upload_watchdog.py | 441 +++++++++++++++ tests/test_uploads.py | 5 +- 8 files changed, 2684 insertions(+), 86 deletions(-) create mode 100644 hotdata/_upload_watchdog.py create mode 100644 tests/test_upload_codex_fixes.py create mode 100644 tests/test_upload_multipart_resilient.py create mode 100644 tests/test_upload_resilient.py create mode 100644 tests/test_upload_retry_policy.py create mode 100644 tests/test_upload_watchdog.py diff --git a/hotdata/_upload_watchdog.py b/hotdata/_upload_watchdog.py new file mode 100644 index 0000000..5edaa39 --- /dev/null +++ b/hotdata/_upload_watchdog.py @@ -0,0 +1,462 @@ +"""Per-part stall cancellation for multipart storage uploads (Item 3). + +A storage part ``PUT`` can *black-hole*: the TCP connection is established and +bytes are accepted into the kernel send buffer, then the far end stops reading +(a send-side stall) or never sends a response (a response-read stall). No RST +arrives, so neither a connect timeout nor urllib3's inactivity ``read`` timeout +reliably catches it, and the worker thread hangs forever — which, without this +module, would hang the whole upload (the outer re-sweep never gets to retry). + +CPython cannot interrupt a thread blocked in a socket syscall from the outside, +but it *can* be woken by acting on the socket itself. This module reaches the +in-flight socket via a connection **subclass** injected into the storage pool — +keeping ``pool.urlopen`` (so urllib3 still owns request-target construction, body +framing, response adaptation and redirects) — and a watchdog thread that calls +``sock.shutdown(SHUT_RDWR)`` on a stalled connection. ``shutdown`` (NOT +``close``) is the spike-verified safe primitive: it wakes a blocked +``send``/``recv`` with a clean error and leaves the fd valid, so closing it stays +the owning request's job and there is no fd-reuse race. + +Wiring (corrected per Codex diff review): + +* The per-part **worker** stashes an :class:`_Arming` (the watchdog + the + per-attempt deadline in seconds) in a **thread-local** before its + ``pool.urlopen`` call, and clears it (disarming all handles) in ``finally``. +* :class:`WatchdogHTTPConnection` / :class:`WatchdogHTTPSConnection`, on EVERY + ``request()`` (i.e. every urllib3 attempt — initial and internal retries) and + on ``connect()``, create a FRESH :class:`_PartHandle`, publish their live + ``.sock`` to it, and arm it on the watchdog with a fresh ``part_put_timeout`` + deadline — **before** the blocking ``super().request()`` send, so a reused + keep-alive connection (which does not call ``connect()``) and a urllib3 + internal retry are both protected. This is the architect's per-ATTEMPT transfer + budget; the spike's one-handle-per-urlopen was the simplification that left + reused conns + internal retries unprotected. +* A terminal (non-retryable) failure cancels all in-flight handles immediately + (:meth:`_Watchdog.cancel_all`) so fail-fast doesn't wait out a sibling's full + per-part timeout. +""" + +from __future__ import annotations + +import logging +import socket +import threading +import time +from typing import Any, Dict, List, Optional + +import urllib3 +from urllib3.connection import HTTPConnection, HTTPSConnection + +_log = logging.getLogger("hotdata.uploads") + +#: Per-thread arming context for the part ``PUT`` currently running on this +#: thread. Each part runs in its own pool worker thread; the connection subclass +#: reads this to arm a fresh handle per attempt. A thread-local (never a global) +#: so concurrent parts never cross wires. +_active = threading.local() + + +class _Arming: + """What a part's worker publishes for the connection subclass to use: the + watchdog to arm handles on, the per-attempt deadline, and the set of handles + armed so far for this part (so the worker can disarm them all afterward). + """ + + __slots__ = ("watchdog", "deadline", "current", "all") + + def __init__(self, watchdog: "_Watchdog", deadline: float) -> None: + self.watchdog = watchdog + self.deadline = deadline + #: the handle for the attempt currently in flight (at most one — a handle + #: is armed ONLY while its attempt exclusively owns its socket). + self.current: Optional["_PartHandle"] = None + #: every handle ever armed for this part (so the worker's finally can + #: defensively disarm any straggler). + self.all: List["_PartHandle"] = [] + + +def set_active_arming(arming: Optional[_Arming]) -> None: + """Set (or clear, with ``None``) the active arming for the current thread.""" + _active.arming = arming + + +def _current_arming() -> Optional[_Arming]: + return getattr(_active, "arming", None) + + +def _arm_fresh_handle_for(conn: Any) -> None: + """Connection-subclass hook (called per attempt, BEFORE the blocking send): + disarm the PREVIOUS attempt's handle, then arm a FRESH one for ``conn``. + + Disarming the previous handle here is what frees a superseded connection + (e.g. attempt-1 errored and urllib3 is retrying on a fresh conn) so a stale + handle can never shut down a socket the pool has reissued (Codex #1). + """ + arming = _current_arming() + if arming is None: + return + _disarm_current(arming) + handle = _PartHandle() + handle.publish(conn) + arming.current = handle + arming.all.append(handle) + arming.watchdog.arm(handle, arming.deadline) + + +def _disarm_current(arming: Optional[_Arming]) -> None: + """Mark the current attempt's handle done and remove it from the watchdog — + its socket is no longer exclusively owned by this attempt, so it must not be + shut down. Called when superseding (new attempt) and after the response is + obtained (conn about to return to the pool). + """ + if arming is None or arming.current is None: + return + handle = arming.current + arming.current = None + handle.finish() + arming.watchdog.disarm(handle) + + +def _disarm_active_handle() -> None: + """Connection-subclass hook: called right after a response is obtained for + the current attempt — the conn is about to be released to the pool, so its + handle must be disarmed immediately (cross-part reuse safety, Codex #1). + """ + _disarm_current(_current_arming()) + + +def _abort_precheck() -> None: + """Connection-subclass hook (called at the TOP of ``connect()`` and + ``request()``): if this part's watchdog has already been terminally aborted, + raise BEFORE opening/using a socket, so a part that begins after a terminal + error fails fast instead of blocking in DNS/TCP/TLS establishment (Codex #5 + pre-connection gap). Sticky-abort's socket-shutdown only helps once a socket + exists; on a fresh connection there is none yet, so we must refuse to proceed. + + Raised as a urllib3 ``ProtocolError`` — the SAME class the watchdog's + ``shutdown`` produces — so ``_put_to_storage`` wraps it as a retryable + ``StorageTransportError`` and it flows to the outer re-sweep / terminal path. + """ + arming = _current_arming() + if arming is not None and arming.watchdog.is_aborted(): + raise urllib3.exceptions.ProtocolError( + "upload aborted before connection", + ConnectionAbortedError("upload watchdog aborted"), + ) + + +def _get_active_arming_for_test() -> Optional[_Arming]: + """Test hook: read the current thread's active arming.""" + return _current_arming() + + +class _PartHandle: + """Tracks one connection for one part-``PUT`` *attempt* and arbitrates the + race between the worker finishing and the watchdog firing. + + Exactly one of three states (Codex #9): + + * ``active`` — the attempt is in flight; the published ``conn`` may be + shut down by the watchdog. + * ``done`` — the worker finished (``urlopen`` returned or raised) before + the watchdog fired; the watchdog must NOT touch the socket (the pool may + reuse it). + * ``cancelled`` — the watchdog shut the socket down first; a later worker + ``finish`` does not flip it back. + + All transitions take ``lock`` so the worker's ``finish`` and the watchdog's + ``cancel_if_active`` are serialized. + """ + + __slots__ = ("lock", "state", "conn") + + def __init__(self) -> None: + self.lock = threading.Lock() + self.state = "active" + self.conn: Any = None + + def publish(self, conn: Any) -> None: + with self.lock: + if self.state == "active": + self.conn = conn + + def finish(self) -> None: + """Mark the attempt done. No-op if already cancelled.""" + with self.lock: + if self.state == "active": + self.state = "done" + self.conn = None + + def cancel_if_active(self) -> Any: + """Watchdog entry point: if still active, mark cancelled and return the + connection to shut down; otherwise return ``None`` (worker won the race). + """ + with self.lock: + if self.state != "active": + return None + self.state = "cancelled" + return self.conn + + +class WatchdogHTTPConnection(HTTPConnection): + """An ``HTTPConnection`` that arms a fresh per-attempt watchdog handle on its + live socket so a stalled ``PUT`` can be cancelled. Publishes BEFORE the + blocking send (and on ``connect()``), so a reused keep-alive connection and a + urllib3-internal retry are both covered. + """ + + def connect(self) -> None: + _abort_precheck() # refuse to open a socket if already aborted + super().connect() + _arm_fresh_handle_for(self) + + def request(self, *args: Any, **kwargs: Any) -> Any: + # Fail fast if aborted BEFORE connecting/sending (no socket yet to shut + # down on a fresh conn). Then arm BEFORE the (possibly blocking) send. A + # reused conn skips connect(), so arming here protects a pooled send-stall. + _abort_precheck() + _arm_fresh_handle_for(self) + return super().request(*args, **kwargs) + + def getresponse(self, *args: Any, **kwargs: Any) -> Any: + # getresponse() blocks waiting for the response (a response-read stall is + # caught here while the handle is still armed). Once it RETURNS, the conn + # is about to be released to the pool — disarm so a stale handle can't + # shut down a socket another part reuses. + resp = super().getresponse(*args, **kwargs) + _disarm_active_handle() + return resp + + +class WatchdogHTTPSConnection(HTTPSConnection): + """HTTPS counterpart of :class:`WatchdogHTTPConnection` (production storage + is HTTPS). ``.sock`` is the TLS socket; ``shutdown`` wakes it the same way. + """ + + def connect(self) -> None: + _abort_precheck() + super().connect() + _arm_fresh_handle_for(self) + + def request(self, *args: Any, **kwargs: Any) -> Any: + _abort_precheck() + _arm_fresh_handle_for(self) + return super().request(*args, **kwargs) + + def getresponse(self, *args: Any, **kwargs: Any) -> Any: + resp = super().getresponse(*args, **kwargs) + _disarm_active_handle() + return resp + + +def install_watchdog_connection_classes(pool: Any) -> None: + """Wire the watchdog connection subclass onto ``pool``. + + Accepts a single ``HTTP(S)ConnectionPool`` (sets ``ConnectionCls``) or a + :class:`~urllib3.PoolManager`. For a ``PoolManager`` (Codex #2: production + uses one over HTTPS), this overrides per-scheme pool creation AND re-stamps + any host pools already cached — so a pool created by an earlier single-PUT to + the same origin still gets the watchdog subclass for a later multipart PUT. + """ + if isinstance(pool, urllib3.HTTPSConnectionPool): + pool.ConnectionCls = WatchdogHTTPSConnection + elif isinstance(pool, urllib3.HTTPConnectionPool): + pool.ConnectionCls = WatchdogHTTPConnection + elif isinstance(pool, urllib3.PoolManager): + _install_on_pool_manager(pool) + else: + # Not a real urllib3 pool (e.g. a test fake): nothing to hook. The + # watchdog simply never engages — a fake doesn't stall — so this is a + # safe no-op rather than an error. + _log.debug( + "watchdog connection class not installed on %s (not a urllib3 pool)", + type(pool).__name__, + ) + + +def _install_on_pool_manager(manager: urllib3.PoolManager) -> None: + """Stamp the watchdog ``ConnectionCls`` on a ``PoolManager``: future pools via + a ``_new_pool`` override, AND any already-cached host pools (Codex #2 — a pool + created before install must not keep the stock connection class). + """ + if not getattr(manager, "_hotdata_watchdog_installed", False): + original_new_pool = manager._new_pool + + def _new_pool(scheme: str, host: str, port: int, request_context: Any = None) -> Any: + new_pool = original_new_pool( + scheme, host, port, request_context=request_context + ) + _stamp_pool(new_pool, scheme) + return new_pool + + manager._new_pool = _new_pool # type: ignore[method-assign] + manager._hotdata_watchdog_installed = True # type: ignore[attr-defined] + + # Re-stamp pools already created+cached before this install ran. urllib3's + # `RecentlyUsedContainer` refuses `.values()`/iteration ("not threadsafe"), + # so enumerate by key and fetch each pool via __getitem__ under its own lock. + existing = getattr(manager, "pools", None) + if existing is not None: + try: + keys = list(existing.keys()) + except Exception: # pragma: no cover - container shape guard + keys = [] + for key in keys: + try: + cached_pool = existing[key] + except KeyError: # evicted between keys() and lookup + continue + if cached_pool is None: + continue + scheme = "https" if isinstance(cached_pool, urllib3.HTTPSConnectionPool) else "http" + _stamp_pool(cached_pool, scheme) + + +def _stamp_pool(pool: Any, scheme: str) -> None: + pool.ConnectionCls = ( + WatchdogHTTPSConnection if scheme == "https" else WatchdogHTTPConnection + ) + + +class _Watchdog: + """A single background thread that shuts down stalled part connections. + + One instance per multipart ``upload_file`` operation. Each in-flight attempt + :meth:`arm`\\ s a handle with a wall-clock deadline; the thread sleeps until + the nearest deadline, then for any past-deadline handle still ``active`` it + ``shutdown``\\ s the socket (the fd is closed by the owning request, never + here). A completed attempt :meth:`disarm`\\ s its handle immediately so the + registry doesn't grow and dead entries don't cause spurious wakeups. + + :meth:`cancel_all` cancels every in-flight handle now (terminal-error + fast-fail). :meth:`shutdown` stops and joins the thread. + """ + + def __init__(self, monotonic: Any = None) -> None: + # Injectable clock so tests don't wait real seconds for a deadline. + self._now = monotonic or time.monotonic + self._cond = threading.Condition() + #: handle -> fire_at deadline (monotonic). dict for O(1) disarm. + self._entries: Dict[_PartHandle, float] = {} + self._stop = False + #: STICKY terminal-abort flag (Codex #5). Once set, ANY handle armed + #: afterwards is cancelled immediately — so a sibling still inside + #: upload_one (minting/reading) that arms its FIRST handle AFTER + #: cancel_all()'s snapshot is still killed promptly, not waited out. + self._aborted = False + self._thread: Optional[threading.Thread] = None + + def arm(self, handle: _PartHandle, deadline: float) -> None: + """Register ``handle`` to be cancelled at ``now + deadline``. If the + watchdog has already been aborted (terminal error), cancel ``handle`` + immediately instead of waiting for its deadline. + """ + with self._cond: + if self._aborted: + aborted = True + else: + aborted = False + self._entries[handle] = self._now() + deadline + if self._thread is None: + self._thread = threading.Thread( + target=self._run, name="hotdata-upload-watchdog", daemon=True + ) + self._thread.start() + self._cond.notify() + if aborted: + # Fire outside the lock (shutdown can block briefly). + self._fire(handle) + + def disarm(self, handle: _PartHandle) -> None: + """Remove a (completed) handle so it can't fire and doesn't linger.""" + with self._cond: + self._entries.pop(handle, None) + self._cond.notify() + + def is_aborted(self) -> bool: + """Whether a terminal abort has fired (sticky). New attempts check this + to refuse to open a connection after a terminal error.""" + with self._cond: + return self._aborted + + def cancel_all(self) -> None: + """Shut down every in-flight handle NOW and STICK the abort (Q-B / Codex + #5), so a handle armed AFTER this call is also cancelled at once. + + Without the sticky flag a sibling still inside ``upload_one`` (minting / + reading the chunk) that arms its first handle a moment later would be + missed by the snapshot and waited out for its full per-part timeout. + """ + with self._cond: + self._aborted = True + handles = list(self._entries.keys()) + for handle in handles: + self._fire(handle) + with self._cond: + self._cond.notify() + + def _run(self) -> None: + with self._cond: + while not self._stop: + now = self._now() + due = [h for h, t in self._entries.items() if t <= now] + for handle in due: + self._entries.pop(handle, None) + if due: + # Fire OUTSIDE the lock (shutdown can block briefly), then + # loop again to recompute against any newly-armed entries — + # never wait on a stale deadline (Codex #3). + self._cond.release() + try: + for handle in due: + self._fire(handle) + finally: + self._cond.acquire() + continue + if self._entries: + next_at = min(self._entries.values()) + self._cond.wait(timeout=max(next_at - now, 0.0)) + else: + self._cond.wait() # sleep until armed or stopped + + @staticmethod + def _fire(handle: _PartHandle) -> None: + conn = handle.cancel_if_active() + if conn is None: + return + sock = getattr(conn, "sock", None) + if sock is None: + return + try: + sock.shutdown(socket.SHUT_RDWR) + _log.debug("watchdog shut down a stalled storage connection") + except OSError: + # Already closed / not connected — the owning request surfaces its + # own error. + pass + + def shutdown(self) -> None: + """Stop and join the watchdog thread. Idempotent. Raises if the thread + fails to terminate (it should not, now that the scheduler can't wait on a + stale deadline — Codex #3). + """ + with self._cond: + self._stop = True + self._cond.notify() + thread = self._thread + if thread is not None: + thread.join(timeout=5.0) + if thread.is_alive(): # pragma: no cover - would signal a scheduler bug + raise RuntimeError("upload watchdog thread did not terminate") + self._thread = None + + +__all__ = [ + "WatchdogHTTPConnection", + "WatchdogHTTPSConnection", + "install_watchdog_connection_classes", + "set_active_arming", + "_Arming", + "_PartHandle", + "_Watchdog", +] diff --git a/hotdata/uploads.py b/hotdata/uploads.py index 1090af1..4b50c4b 100644 --- a/hotdata/uploads.py +++ b/hotdata/uploads.py @@ -46,20 +46,30 @@ import logging import os import threading -from concurrent.futures import ThreadPoolExecutor, as_completed +import time +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from contextlib import contextmanager from dataclasses import dataclass -from typing import Any, BinaryIO, Callable, Dict, Iterator, List, Optional, Union +from typing import Any, BinaryIO, Callable, Dict, Iterator, List, Optional, Tuple, Union import urllib3 from urllib3.util.retry import Retry +from hotdata._retry import ConnectionResetRetry +from hotdata._upload_watchdog import ( + _Arming, + _Watchdog, + install_watchdog_connection_classes, + set_active_arming, +) from hotdata.api.uploads_api import UploadsApi as _GeneratedUploadsApi from hotdata.api_client import ApiClient +from hotdata.exceptions import ApiException from hotdata.models.create_upload_request import CreateUploadRequest from hotdata.models.finalize_upload_part import FinalizeUploadPart from hotdata.models.finalize_upload_request import FinalizeUploadRequest from hotdata.models.finalize_upload_response import FinalizeUploadResponse +from hotdata.models.mint_upload_parts_request import MintUploadPartsRequest from hotdata.models.upload_response import UploadResponse from hotdata.models.upload_session_response import UploadSessionResponse from hotdata.rest import RESTResponse @@ -188,10 +198,21 @@ class UploadOptions: max_concurrency: Optional[int] = None #: Optional progress callback invoked with ``(bytes_done_total, total)``. progress: Optional[UploadProgress] = None - #: ``urllib3`` retry policy for multipart part ``PUT`` s (idempotent, so safe - #: to retry). ``None`` uses :func:`default_part_retry`. The single-``PUT`` - #: path streams an un-replayable body and is always sent without retry. + #: ``urllib3`` retry policy for part ``PUT`` s — the INNER retry tier (429 + + #: pre-response reset; 5xx propagates to the whole-upload re-sweep). ``None`` + #: uses :func:`default_part_retry`. A custom policy that puts ``5xx`` back in + #: its ``status_forcelist`` shifts the tier boundary. The single-``PUT`` + #: small-file path streams an un-replayable body and is always sent without + #: retry regardless. part_retry: Optional[Retry] = None + #: Extra whole-upload re-sweep rounds beyond the initial pass (Item 1). + #: ``None`` uses :data:`DEFAULT_MAX_EXTRA_ROUNDS` (3). ``0`` disables the + #: re-sweep (a single failure sinks the upload — the legacy behavior). + max_extra_rounds: Optional[int] = None + #: Base delay (seconds) for the capped-exponential backoff between re-sweep + #: rounds. ``None`` uses :data:`DEFAULT_ROUND_BASE_DELAY` (2s). Set to ``0`` + #: in tests for instant rounds. + round_base_delay: Optional[float] = None # --- Errors --------------------------------------------------------------- @@ -268,8 +289,11 @@ def __init__(self, *, status: int, part_number: Optional[int], body: str) -> Non class MissingETagError(UploadError): - """Storage accepted a part ``PUT`` but returned no ``ETag`` header, so the - part cannot be finalized. + """Storage accepted a part ``PUT`` but returned no usable ``ETag`` header + (absent, or empty/whitespace-only), so the part cannot be finalized. + + Retryable (not terminal): a momentary backend hiccup that drops the header + shouldn't be permanent, so the re-sweep re-``PUT`` s to pick up a real ETag. """ def __init__(self, *, part_number: int) -> None: @@ -279,6 +303,255 @@ def __init__(self, *, part_number: int) -> None: ) +# --- Error classification (Item 2) ---------------------------------------- + + +def _is_retryable(exc: BaseException) -> bool: + """Whether the outer whole-upload re-sweep should retry a per-part failure. + + A deliberate **allowlist** of genuinely-transient transport / storage / mint + conditions — each curable by a later round at lower concurrency (and, for a + fresh minted URL, by re-minting). Everything NOT listed is non-retryable and + fails the upload fast: a deterministic local short-read (a generic + :class:`UploadError`), a contract/sizing violation + (:class:`MalformedSessionError`, :class:`SizeLimitError`), a progress-callback + exception (raised *after* a successful ``PUT`` — re-running would duplicate the + write), or anything unexpected. Retrying those would burn the round budget for + nothing or, worse, double-``PUT``. + """ + if isinstance(exc, (StorageTransportError, MissingETagError)): + return True + if isinstance(exc, StorageError): + return True # incl. 5xx (outer tier owns it) and 403 (eager URL expiry) + if isinstance(exc, ApiException): + return True # a per-part mint round-trip failure — a re-mint may cure it + return False + + +def _parse_etag(headers: Any, part_number: int) -> str: + """Extract a real, non-empty ``ETag`` from a part ``PUT`` response (Item 7). + + A missing header OR an empty / whitespace-only value is rejected as + :class:`MissingETagError` (retryable). Otherwise the value is returned + verbatim, quotes included — finalize must carry exactly what storage gave. + """ + etag = headers.get("ETag") + if etag is None or not str(etag).strip(): + raise MissingETagError(part_number=part_number) + return etag + + +# --- Whole-upload re-sweep (Item 1) --------------------------------------- + +#: Outer-tier (whole-upload re-sweep) defaults, mirroring the Rust SDK. Round 0 +#: is the initial full pass; up to :data:`DEFAULT_MAX_EXTRA_ROUNDS` re-sweeps +#: follow, each re-running only the parts still failing, at halving concurrency. +DEFAULT_MAX_EXTRA_ROUNDS = 3 +#: Capped-exponential backoff before re-sweep round N: ``base * 2**(N-1)``. +DEFAULT_ROUND_BASE_DELAY = 2.0 +#: Clamp on the round-delay shift so it can't overflow. +_ROUND_DELAY_SHIFT_CAP = 16 + + +@dataclass +class _PartPlan: + """One part's work unit, keyed by its 1-based ``part_number``. Byte range is + derived from the part number (``offset = index * part_size``), never from + arrival/queue position, so a reordered mint response can't shift offsets + (Item 6). + """ + + index: int # 0-based position + part_number: int # 1-based + offset: int + length: int + + +def _round_in_flight(base: int, round_idx: int) -> int: + """In-flight cap for re-sweep round ``round_idx``: halve each round, floor 1 + (``max(base >> round, 1)``). 8 → 8,4,2,1,1,… A saturated jittery uplink + resets fewer connections when fewer are in flight, so lowering concurrency + each round directly targets that failure mode. + """ + if round_idx >= base.bit_length(): + return 1 + return max(base >> round_idx, 1) + + +def _round_delay(round_idx: int, base_delay: float = DEFAULT_ROUND_BASE_DELAY) -> float: + """Capped-exponential backoff before round ``round_idx`` (>=1): + ``base_delay * 2**(round_idx - 1)``, shift clamped so it can't overflow. + """ + shift = min(round_idx - 1, _ROUND_DELAY_SHIFT_CAP) + return base_delay * (2 ** shift) + + +def _upload_parts_resilient( + plans: List[_PartPlan], + upload_one: Callable[[_PartPlan], FinalizeUploadPart], + *, + base_in_flight: int, + max_extra_rounds: int = DEFAULT_MAX_EXTRA_ROUNDS, + round_base_delay: float = DEFAULT_ROUND_BASE_DELAY, + on_part_done: Optional[Callable[[_PartPlan], None]] = None, + on_terminal: Optional[Callable[[], None]] = None, + sleep: Callable[[float], None] = time.sleep, +) -> List[FinalizeUploadPart]: + """Run every part's ``upload_one`` work unit with whole-upload re-sweep. + + Round 0 is the initial full pass at ``base_in_flight`` concurrency. Each + subsequent round (up to ``max_extra_rounds``) waits a capped-exponential + backoff, halves concurrency, and re-runs ONLY the parts that failed the + previous round with a *retryable* error; completed parts keep their ETags. + + Invariants: every healthy part is attempted exactly once; a part failing K + rounds is attempted ``K+1`` times (capped); a **terminal** (non-retryable, + see :func:`_is_retryable`) error fails fast — the round stops dispatching new + parts and the error is raised after in-flight workers drain. If parts still + fail after the last round, the last underlying retryable error is raised + (its cause preserved). A part slot that is never filled is a hard + :class:`MalformedSessionError` — finalize never gets a silently-short list. + + ``on_part_done`` is invoked once per successfully-completed part (used for + exactly-once progress): a re-swept part that failed earlier rounds did not + fire it, so bytes are never double-counted. + """ + results: Dict[int, FinalizeUploadPart] = {} + counted: set = set() + lock = threading.Lock() + pending = list(plans) + + last_error: Optional[BaseException] = None + + for round_idx in range(0, 1 + max_extra_rounds): + if not pending: + break + if round_idx > 0: + sleep(_round_delay(round_idx, round_base_delay)) + in_flight = _round_in_flight(base_in_flight, round_idx) + failures: Dict[int, BaseException] = {} + terminal: List[BaseException] = [] + stop = threading.Event() + + def _fail_terminal(exc: BaseException) -> None: + with lock: + terminal.append(exc) + stop.set() + # Actively cancel in-flight siblings (terminal fast-fail, Q-B) so the + # round-boundary join doesn't wait out a stalled sibling's full + # per-part timeout. No-op if the caller wired no canceller. + if on_terminal is not None: + on_terminal() + + def run_one(plan: _PartPlan) -> None: + if stop.is_set(): + # A terminal error elsewhere in the round: don't dispatch new work. + failures[plan.part_number] = _NotDispatched() + return + try: + part = upload_one(plan) + except BaseException as exc: # noqa: BLE001 - classified below + if not _is_retryable(exc): + _fail_terminal(exc) + return + with lock: + failures[plan.part_number] = exc + return + if part is None: + with lock: + failures[plan.part_number] = MalformedSessionError( + f"part {plan.part_number} was never uploaded" + ) + return + # Record success first, then fire the progress callback. A callback + # exception is NOT a transport failure and the PUT already succeeded — + # re-running would duplicate the write — so it is TERMINAL (abort the + # upload, do not finalize), never retried (Codex #4). + with lock: + results[plan.part_number] = part + fire = on_part_done is not None and plan.part_number not in counted + if fire: + counted.add(plan.part_number) + if fire: + try: + on_part_done(plan) + except BaseException as exc: # noqa: BLE001 - callback bug → abort + _fail_terminal(exc) + + executor = ThreadPoolExecutor(max_workers=in_flight) + try: + # Bounded-wave dispatch: never submit more than `in_flight` at once, + # so the per-round cap actually bounds concurrency AND a terminal + # error stops not-yet-started parts from being dispatched. + it = iter(pending) + futures: set = set() + for plan in _take(it, in_flight): + futures.add(executor.submit(run_one, plan)) + while futures: + done, futures = _wait_first(futures) + # Surface any unexpected exception that escaped run_one (it should + # classify everything itself; this is defense-in-depth so a future + # exception can never be silently swallowed). + for fut in done: + exc = fut.exception() + if exc is not None: + _fail_terminal(exc) + for plan in _take(it, len(done)): + if stop.is_set(): + break + futures.add(executor.submit(run_one, plan)) + finally: + # Round-boundary join: by the time this returns, no worker from this + # round is still touching the source / firing progress. The watchdog + # (Item 3) makes a stalled worker terminable so this stays bounded. + executor.shutdown(wait=True) + + if terminal: + raise terminal[0] + + # Track the most recent real failure to surface on exhaustion. + for exc in failures.values(): + if not isinstance(exc, _NotDispatched): + last_error = exc + pending = [p for p in pending if p.part_number in failures] + + if pending: + # Exhausted the round budget with parts still failing: surface the last + # underlying error so the caller's normal mapping applies. + if last_error is not None: + raise last_error + raise StorageTransportError(url="", part_number=pending[0].part_number) + + # Completeness: every plan's slot must be filled. + missing = [p.part_number for p in plans if p.part_number not in results] + if missing: + raise MalformedSessionError(f"part {missing[0]} was never uploaded") + + return [results[pn] for pn in sorted(results)] + + +class _NotDispatched(Exception): + """Sentinel failure for a part skipped because the round was aborting on a + terminal error; it is re-queued but never surfaced as the cause. + """ + + +def _take(it: Iterator[_PartPlan], n: int) -> List[_PartPlan]: + out: List[_PartPlan] = [] + for _ in range(n): + try: + out.append(next(it)) + except StopIteration: + break + return out + + +def _wait_first(futures: set) -> Tuple[set, set]: + """Block until at least one future completes; return ``(done, still_running)``.""" + done, not_done = wait(futures, return_when=FIRST_COMPLETED) + return done, not_done + + # --- Storage transport ---------------------------------------------------- # A dedicated, process-wide, *bare* pool for storage PUTs. Deliberately NOT the @@ -294,26 +567,124 @@ def _storage_pool() -> urllib3.PoolManager: if _STORAGE_POOL is None: with _STORAGE_POOL_LOCK: if _STORAGE_POOL is None: - _STORAGE_POOL = urllib3.PoolManager() + pool = urllib3.PoolManager() + # Install the stall-cancellation connection subclass at creation + # time — BEFORE any request can populate `manager.pools` with a + # host pool carrying the stock connection class (Codex #2). The + # single-PUT path uses this same pool, so an early single-PUT must + # not leave an un-hooked cached pool for a later multipart PUT. + install_watchdog_connection_classes(pool) + _STORAGE_POOL = pool return _STORAGE_POOL +#: Per-part transfer-timeout shape (Item 3). A part ``PUT`` attempt is killed (into +#: the outer re-sweep) only if it cannot sustain even ``PART_TIMEOUT_MIN_BYTES_PER_SEC`` +#: — a true stall, never a slow-but-progressing transfer. Size-scaled so a large +#: part legitimately gets longer, capped so a stalled giant part can't hang for hours. +PART_TIMEOUT_BASE = 60.0 +PART_TIMEOUT_MIN_BYTES_PER_SEC = 64 * 1024 +PART_TIMEOUT_MAX = 30 * 60.0 + +#: Per-sleep cap for the inner part retry's *exponential* backoff path. +PART_RETRY_BACKOFF_MAX = 120.0 + +#: Cap on a server-sent ``Retry-After`` for a part ``PUT`` (the per-part +#: retry-window deadline analog). ``backoff_max`` only bounds the exponential +#: backoff path; urllib3 honors ``Retry-After`` UNCAPPED. A multi-minute storage +#: ``Retry-After`` would stall the whole-upload re-sweep's round-boundary join, so +#: :class:`_RetryAfterCappedRetry` clamps it. Flat 120s, independent of part size +#: — NOT tied to :func:`_part_put_timeout` (the per-attempt transfer budget, up to +#: 30 min): tying the sleep cap to that would let a big part honor a 30-min +#: ``Retry-After`` and reintroduce the stall. +MAX_RETRY_AFTER_SECONDS = 120.0 + + +class _RetryAfterCappedRetry(ConnectionResetRetry): + """A part-pool retry that caps a server ``Retry-After`` at + :data:`MAX_RETRY_AFTER_SECONDS`. + + Overriding ``get_retry_after`` (NOT ``sleep`` / ``sleep_for_retry``, whose + shape changed across urllib3 1.26→2.x) is forward-compatible: it returns + seconds in both versions, and ``Retry.sleep`` consults it. ``None`` (no + ``Retry-After`` header) passes through untouched so the exponential backoff + path is unaffected. + """ + + def get_retry_after(self, response: Any) -> Any: + retry_after = super().get_retry_after(response) + if retry_after is None: + return None + return min(retry_after, MAX_RETRY_AFTER_SECONDS) + +#: Inactivity (connect, read) timeout for the per-part mint call (Item 4 / Codex #5). +#: NOTE: this is an INACTIVITY timeout, not a strict wall-clock bound — a slow-trickle +#: mint response could in principle evade it (same urllib3 limitation as a read +#: timeout). It bounds a mint *in practice* so a stalled mint can't hang the +#: round-boundary join; it is not a hard wall-clock guarantee. A mint trickle-stall +#: is far rarer than a storage one, so a mint watchdog is deliberately not built. +MINT_TIMEOUT = (10.0, 30.0) + + +def _part_put_timeout(content_length: int) -> float: + """The size-scaled per-attempt transfer cap for one part ``PUT`` (Item 3). + + ``min(BASE + content_length / MIN_BYTES_PER_SEC, MAX)``: 8 MiB → 188s, 64 MiB + → ~18 min, ≥ ~111 MiB → the 30-min cap, 0 bytes → the 60s base. This is the + ONLY bound on a single in-flight PUT; the watchdog enforces it by shutting the + socket down. Pure and total. + """ + return min( + PART_TIMEOUT_BASE + content_length / PART_TIMEOUT_MIN_BYTES_PER_SEC, + PART_TIMEOUT_MAX, + ) + + def default_part_retry() -> Retry: - """The default retry policy for multipart part ``PUT`` s. + """The default INNER retry policy for multipart part ``PUT`` s (Item 5). A part ``PUT`` is idempotent — storage overwrites a part by number — so a - retried part cannot corrupt the upload. Retry connection errors and the - transient 429/5xx statuses storage may return under load. The single-``PUT`` - path streams an un-replayable body and so is always sent without retry. - Override per upload via :attr:`UploadOptions.part_retry`. + retried part cannot corrupt the upload. This is the *inner* tier of the + two-tier model: it retries genuinely-momentary conditions in place — + + * HTTP ``429`` admission-shedding, honoring ``Retry-After``; + * pre-response connection resets on any method (via the + :class:`~hotdata._retry.ConnectionResetRetry` base), which includes the + watchdog's ``ProtocolError(RemoteDisconnected)`` when a stalled part is + cancelled (Item 3) — so a watchdog-cancelled part is re-attempted. + + Storage ``5xx`` is deliberately **excluded** from ``status_forcelist`` so it + is NOT inner-retried; it propagates to the whole-upload re-sweep (the outer + tier), which owns its recovery with reduced concurrency. Redirects are not + followed: a ``3xx`` on a presigned ``PUT`` would re-issue to a different URL + and invalidate the signature, so it surfaces as an error instead. + + ``backoff_max`` caps a pathological ``Retry-After`` so the re-sweep's + round-boundary join can't stall on an unbounded sleep. + + Override per upload via :attr:`UploadOptions.part_retry`; a custom policy that + puts ``5xx`` back in its ``status_forcelist`` shifts the tier boundary (the + caller's explicit choice). """ - return Retry( + kwargs: Dict[str, Any] = dict( total=3, backoff_factor=0.2, - status_forcelist=(429, 500, 502, 503, 504), + status_forcelist=(429,), allowed_methods=frozenset({"PUT"}), raise_on_status=False, + respect_retry_after_header=True, + redirect=False, ) + # urllib3 >= 2 takes backoff_max as a constructor arg; 1.26 only honors the + # instance/class attribute `DEFAULT_BACKOFF_MAX` (which `get_backoff_time` + # reads — NOT the deprecated `BACKOFF_MAX` alias). Set it the version-portable + # way so the per-sleep cap is explicit regardless of urllib3 version. + try: + return _RetryAfterCappedRetry(backoff_max=PART_RETRY_BACKOFF_MAX, **kwargs) + except TypeError: + retry = _RetryAfterCappedRetry(**kwargs) + retry.DEFAULT_BACKOFF_MAX = PART_RETRY_BACKOFF_MAX # type: ignore[attr-defined] + return retry def _to_timeout(request_timeout: Any) -> Any: @@ -360,7 +731,12 @@ def _put_to_storage( raise StorageTransportError(url=url, part_number=part_number) from exc status = int(response.status) _log.debug("storage PUT %s -> %s", url, status) - if status >= 400: + # Any non-2xx is an error, INCLUDING a 3xx (Codex #6): with redirects disabled + # on the part pool, urllib3 returns the 3xx response rather than following it, + # and following a redirect on a presigned PUT would re-issue to a different URL + # and invalidate the signature. So a storage 3xx surfaces as StorageError, not + # a silent success. + if not 200 <= status < 300: data = getattr(response, "data", b"") or b"" if isinstance(data, (bytes, bytearray)): error_body = bytes(data).decode("utf-8", errors="replace") @@ -437,6 +813,8 @@ def upload_file( # type: ignore[override] max_concurrency: Optional[int] = None, progress: Optional[UploadProgress] = None, part_retry: Optional[Retry] = None, + max_extra_rounds: Optional[int] = None, + round_base_delay: Optional[float] = None, options: Optional[UploadOptions] = None, _request_timeout: Any = None, ) -> FinalizeUploadResponse: @@ -499,6 +877,12 @@ def upload_file( # type: ignore[override] ) progress = progress if progress is not None else opts.progress part_retry = part_retry if part_retry is not None else opts.part_retry + max_extra_rounds = ( + max_extra_rounds if max_extra_rounds is not None else opts.max_extra_rounds + ) + round_base_delay = ( + round_base_delay if round_base_delay is not None else opts.round_base_delay + ) src = _make_source(source, size) total = src.size @@ -538,7 +922,22 @@ def upload_file( # type: ignore[override] elif session.mode == "multipart": cap = max_concurrency if max_concurrency is not None else DEFAULT_MAX_CONCURRENCY retry = part_retry if part_retry is not None else default_part_retry() - parts = self._upload_multipart(session, src, total, cap, retry, progress) + parts = self._upload_multipart( + session, + src, + total, + cap, + retry, + progress, + max_extra_rounds=( + max_extra_rounds if max_extra_rounds is not None + else DEFAULT_MAX_EXTRA_ROUNDS + ), + round_base_delay=( + round_base_delay if round_base_delay is not None + else DEFAULT_ROUND_BASE_DELAY + ), + ) else: raise MalformedSessionError(f"unknown upload mode {session.mode!r}") @@ -743,96 +1142,172 @@ def _upload_multipart( max_concurrency: int, part_retry: Retry, progress: Optional[UploadProgress], + *, + max_extra_rounds: int = DEFAULT_MAX_EXTRA_ROUNDS, + round_base_delay: float = DEFAULT_ROUND_BASE_DELAY, ) -> List[FinalizeUploadPart]: - """Multipart path: slice the file into ``part_size``-byte chunks (the - last is the remainder), ``PUT`` each chunk to its ``part_urls[i - 1]`` - with bounded concurrency, and collect ``(part_number, e_tag)`` per part. - - Returns the parts sorted ascending by part number, ready for finalize. + """Resilient multipart path: ``PUT`` each ``part_size``-byte chunk to + storage with the whole-upload re-sweep (Item 1) — a single part's + transient failure never discards the upload. + + URL acquisition is per-path: a *known-size* (eager) session carries + ``part_urls``; a *streaming* session (``part_urls`` absent) mints each + part's URL on demand via ``POST /v1/uploads/{id}/parts``. Both run through + the same re-sweep driver + per-part unit. A stalled ``PUT`` is cancelled + by the per-upload watchdog (Item 3). Returns the parts ascending by part + number, ready for finalize. """ - part_urls = session.part_urls part_size = session.part_size - if not part_urls: - raise MalformedSessionError("multipart upload missing `part_urls`") if part_size is None or part_size <= 0: raise MalformedSessionError( f"multipart upload has invalid `part_size` {part_size!r}" ) - # The URL count must match the number of `part_size`-byte chunks the file - # splits into (last is the remainder). A mismatch means a session - # inconsistent with our declared size, so fail loudly. + part_urls = session.part_urls + # Streaming (mint-on-demand) is signalled by `part_urls` being ABSENT + # (None) on the wire. A present-but-empty list is a malformed known-size + # session (its URL count won't match), NOT streaming. + streaming = part_urls is None expected_parts = max(-(-total // part_size), 1) # ceil, at least 1 - if len(part_urls) != expected_parts: + if not streaming and len(part_urls) != expected_parts: + # A known-size session's URL count must match how the file slices. raise MalformedSessionError( f"multipart upload returned {len(part_urls)} part URLs but the " f"file ({total} bytes) splits into {expected_parts} parts of " f"{part_size} bytes" ) + plans = [ + _PartPlan( + index=i, + part_number=i + 1, + offset=i * part_size, + length=min(part_size, total - i * part_size), + ) + for i in range(expected_parts) + ] + in_flight = effective_in_flight(max_concurrency, part_size) - done = [0] # boxed running byte total; guarded by `lock` - lock = threading.Lock() - - def upload_part(index: int) -> FinalizeUploadPart: - part_number = index + 1 - offset = index * part_size - length = min(part_size, total - offset) - chunk = src.read_range(offset, length) - if len(chunk) != length: - # The source returned fewer bytes than its declared size implies - # (a truncated file, or a wrong explicit `size`). Sending it would - # mismatch Content-Length, so fail loudly before the PUT. + progress_lock = threading.Lock() + done = [0] # boxed running byte total; guarded by `progress_lock` + + # Install the cancellation hook on the storage pool, and run one watchdog + # for this whole upload. Both are torn down in `finally`. + pool = _storage_pool() + install_watchdog_connection_classes(pool) + watchdog = _Watchdog() + + def upload_one(plan: _PartPlan) -> FinalizeUploadPart: + chunk = src.read_range(plan.offset, plan.length) + if len(chunk) != plan.length: + # Deterministic local short-read (truncated source / wrong size): + # a generic UploadError, which the driver classifies non-retryable. raise UploadError( - f"source returned {len(chunk)} bytes for part {part_number} " - f"(expected {length} at offset {offset}); it may have changed " - f"size or the declared size is wrong" + f"source returned {len(chunk)} bytes for part {plan.part_number} " + f"(expected {plan.length} at offset {plan.offset}); it may have " + f"changed size or the declared size is wrong" ) - response = _put_to_storage( - part_urls[index], - chunk, - _storage_headers(session.headers, length), - retries=part_retry, - part_number=part_number, - ) - etag = response.headers.get("ETag") - if not etag: - raise MissingETagError(part_number=part_number) - if progress is not None: - # Fire the callback under the lock so delivery order matches the - # computed order: releasing the lock first would let two workers - # race and deliver ticks out of order, breaking the documented - # monotonically-non-decreasing guarantee. Callbacks are expected - # to be cheap (and do their own locking), so serializing is fine. - with lock: - done[0] = min(done[0] + length, total) - progress(done[0], total) - return FinalizeUploadPart(e_tag=etag, part_number=part_number) - results: List[Optional[FinalizeUploadPart]] = [None] * len(part_urls) - executor = ThreadPoolExecutor(max_workers=in_flight) + url = self._part_url(session, plan, streaming) + headers = _storage_headers(session.headers, plan.length) + + # Publish a per-attempt arming: the connection subclass arms a FRESH + # watchdog handle on each urllib3 attempt (initial + internal retries, + # reused keep-alive conns included), with this part's transfer budget. + arming = _Arming(watchdog, _part_put_timeout(plan.length)) + set_active_arming(arming) + try: + try: + response = _put_to_storage( + url, chunk, headers, retries=part_retry, part_number=plan.part_number + ) + except StorageError as exc: + # Eager-path URL expiry: a 403 (ONLY) gets ONE inline re-mint + # + retry of the fresh URL. A 5xx is NOT an expiry signal — it + # propagates to the outer re-sweep (the tier boundary). + if not streaming and exc.status == 403: + fresh = self._mint_part(session, plan.part_number) + response = _put_to_storage( + fresh, chunk, headers, retries=part_retry, + part_number=plan.part_number, + ) + else: + raise + finally: + # Defensively disarm every handle this part armed (the subclass + # already disarms each attempt's handle on supersede and after its + # response; this catches the error path where getresponse() never + # returned). A completed/failed part can't be shut down later. + for handle in arming.all: + handle.finish() + watchdog.disarm(handle) + set_active_arming(None) + + etag = _parse_etag(response.headers, plan.part_number) + return FinalizeUploadPart(e_tag=etag, part_number=plan.part_number) + + def on_part_done(plan: _PartPlan) -> None: + if progress is None: + return + with progress_lock: + done[0] = min(done[0] + plan.length, total) + progress(done[0], total) + try: - futures = { - executor.submit(upload_part, i): i for i in range(len(part_urls)) - } - # Surface the first failure as soon as it lands rather than waiting - # for in-flight stragglers to drain. A part PUT is idempotent, so any - # straggler still running when we bail does no harm. - for future in as_completed(futures): - exc = future.exception() - if exc is not None: - raise exc - results[futures[future]] = future.result() + return _upload_parts_resilient( + plans, + upload_one, + base_in_flight=in_flight, + max_extra_rounds=max_extra_rounds, + round_base_delay=round_base_delay, + on_part_done=on_part_done, + on_terminal=watchdog.cancel_all, + ) finally: - # cancel_futures drops parts that have not started; wait=True then - # blocks for the few already-running workers so that, by the time - # this returns (or re-raises), no background thread is still reading - # the (possibly caller-owned) source or firing the progress callback. - executor.shutdown(wait=True, cancel_futures=True) + watchdog.shutdown() - # results is indexed by 0-based part position, so it is already ascending - # by part_number with no duplicates. - return [part for part in results if part is not None] + def _part_url( + self, session: UploadSessionResponse, plan: _PartPlan, streaming: bool + ) -> str: + """Resolve the storage URL for ``plan``: the eager ``part_urls`` entry + for a known-size session, or a freshly minted URL for a streaming one. + """ + if not streaming: + return session.part_urls[plan.index] + return self._mint_part(session, plan.part_number) + + def _mint_part(self, session: UploadSessionResponse, part_number: int) -> str: + """Mint a fresh presigned URL for one part (streaming path, and the + eager-path 403 re-mint). Validates the response is for the requested part + number (Item 6); carries a finite (inactivity) timeout (Item 4 / Codex #5). + """ + resp = self.mint_upload_parts_handler( + session.upload_id, + session.finalize_token, + MintUploadPartsRequest(part_numbers=[part_number]), + _request_timeout=MINT_TIMEOUT, + ) + # Strict validation (Item 6): we requested exactly one part number, so the + # response must contain exactly that one. Reject a DUPLICATE (two entries + # for the part — ambiguous, and a plain dict would silently collapse it), + # an EXTRA/UNREQUESTED part number (contract violation), or a MISSING one. + parts = resp.parts or [] + numbers = [p.part_number for p in parts] + if numbers.count(part_number) > 1: + raise MalformedSessionError( + f"mint returned duplicate URLs for part {part_number}" + ) + extra = [n for n in numbers if n != part_number] + if extra: + raise MalformedSessionError( + f"mint returned URLs for parts that were not requested: {sorted(set(extra))}" + ) + url = next((p.url for p in parts if p.part_number == part_number), None) + if not url: + raise MalformedSessionError( + f"mint did not return a URL for part {part_number}" + ) + return url class _Source: @@ -959,7 +1434,9 @@ def _make_source(source: UploadSource, size: Optional[int]) -> _Source: __all__ = [ "DEFAULT_MAX_CONCURRENCY", + "DEFAULT_MAX_EXTRA_ROUNDS", "DEFAULT_PART_SIZE", + "DEFAULT_ROUND_BASE_DELAY", "MAX_PART_SIZE", "MIB", "MIN_PART_SIZE", diff --git a/tests/test_upload_codex_fixes.py b/tests/test_upload_codex_fixes.py new file mode 100644 index 0000000..2584602 --- /dev/null +++ b/tests/test_upload_codex_fixes.py @@ -0,0 +1,427 @@ +"""Regression guards for the Codex diff-review must-fix bugs (#2-#5 + nits + +strict mint validation). Each test reproduces the specific bug the green suite +missed. Run on urllib3 2.x (the project's required version). +""" + +from __future__ import annotations + +import socket +import threading +import time +from typing import Any, Dict, List, Optional + +import pytest +import urllib3 + +import hotdata.uploads as uploads_mod +from hotdata._upload_watchdog import ( + WatchdogHTTPConnection, + WatchdogHTTPSConnection, + _PartHandle, + _Watchdog, + install_watchdog_connection_classes, +) +from hotdata.uploads import ( + MIB, + MalformedSessionError, + _PartPlan, + _upload_parts_resilient, +) +from hotdata.models.finalize_upload_part import FinalizeUploadPart +from hotdata.models.minted_upload_part_response import MintedUploadPartResponse +from hotdata.models.mint_upload_parts_response import MintUploadPartsResponse +from hotdata.models.upload_session_response import UploadSessionResponse + + +# --- #2: existing cached PoolManager host pools get the subclass ---------- + + +def test_poolmanager_restamps_already_cached_host_pool() -> None: + """If a host pool was created+cached BEFORE install (e.g. by an earlier + single-PUT to the same origin), install must re-stamp it — otherwise the + later multipart PUT reuses a stock (non-cancellable) connection class. + """ + pm = urllib3.PoolManager() + # Simulate an earlier request caching a host pool with the STOCK class. + cached = pm.connection_from_url("https://storage.example/") + assert cached.ConnectionCls is not WatchdogHTTPSConnection + + install_watchdog_connection_classes(pm) + + # The already-cached pool is re-stamped... + assert cached.ConnectionCls is WatchdogHTTPSConnection + # ...and a NEW host pool also gets it. + other = pm.connection_from_url("https://other.example/") + assert other.ConnectionCls is WatchdogHTTPSConnection + + +def test_storage_pool_installs_watchdog_at_creation(monkeypatch) -> None: + # _storage_pool() must hook the subclass at creation time, before any request + # can cache a host pool. + monkeypatch.setattr(uploads_mod, "_STORAGE_POOL", None) + pool = uploads_mod._storage_pool() + p = pool.connection_from_url("https://storage.example/") + assert p.ConnectionCls is WatchdogHTTPSConnection + http_p = pool.connection_from_url("http://storage.example/") + assert http_p.ConnectionCls is WatchdogHTTPConnection + + +# --- #3: watchdog scheduler does not miss a newly-armed earlier deadline --- + + +def test_watchdog_fires_near_deadline_armed_after_far_one() -> None: + """Arm a far deadline (watchdog sleeps long), then arm a NEAR deadline + concurrently. The near one must still fire on time — the scheduler must + recompute under the lock, not wait on a stale `wait` value (Codex #3). + """ + wd = _Watchdog() + fired_order: List[str] = [] + lock = threading.Lock() + + class _Conn: + def __init__(self, tag): + self.tag = tag + self.sock = self + + def shutdown(self, how): + with lock: + fired_order.append(self.tag) + + far = _PartHandle() + far.publish(_Conn("far")) + wd.arm(far, deadline=30.0) # watchdog goes to sleep ~30s + time.sleep(0.05) # let it enter the wait + + near = _PartHandle() + near.publish(_Conn("near")) + wd.arm(near, deadline=0.1) # should fire in ~0.1s, NOT wait 30s + + deadline = time.monotonic() + 3.0 + while time.monotonic() < deadline: + with lock: + if "near" in fired_order: + break + time.sleep(0.02) + wd.shutdown() + assert "near" in fired_order # fired despite the far deadline being armed first + + +def test_watchdog_shutdown_joins_thread() -> None: + wd = _Watchdog() + h = _PartHandle() + wd.arm(h, deadline=100.0) + wd.shutdown() # must not raise (thread terminated) and must join + assert wd._thread is None + + +# --- #4: a progress-callback exception aborts the upload, no finalize ------ + + +def test_callback_exception_aborts_does_not_finalize() -> None: + def uploader(plan: _PartPlan) -> FinalizeUploadPart: + return FinalizeUploadPart(e_tag=f'"e{plan.part_number}"', part_number=plan.part_number) + + plans = [_PartPlan(index=i, part_number=i + 1, offset=i * 100, length=100) for i in range(3)] + + def boom(plan: _PartPlan) -> None: + if plan.part_number == 2: + raise RuntimeError("user callback blew up") + + with pytest.raises(RuntimeError, match="user callback blew up"): + _upload_parts_resilient( + plans, + uploader, + base_in_flight=2, + max_extra_rounds=3, + round_base_delay=0.0, + on_part_done=boom, + ) + + +# --- #5: terminal error cancels in-flight siblings (not drain) ------------- + + +def test_terminal_error_cancels_inflight_siblings() -> None: + """On a terminal error while a sibling is stalled, the driver must invoke + `on_terminal` (the canceller) so fail-fast doesn't wait out the sibling's + full per-part timeout. + """ + started = threading.Event() + release = threading.Event() + cancelled = threading.Event() + + def uploader(plan: _PartPlan) -> FinalizeUploadPart: + if plan.part_number == 1: + # Terminal error immediately. + raise MalformedSessionError("contract violation") + # Sibling part: block until cancelled (simulating a stalled PUT). + started.set() + if release.wait(timeout=5.0): + return FinalizeUploadPart(e_tag='"e2"', part_number=2) + raise AssertionError("sibling was drained, not cancelled") + + plans = [_PartPlan(index=i, part_number=i + 1, offset=0, length=100) for i in range(2)] + + def on_terminal() -> None: + cancelled.set() + release.set() # unblock the stalled sibling promptly (what cancel_all does) + + with pytest.raises(MalformedSessionError): + _upload_parts_resilient( + plans, + uploader, + base_in_flight=2, + max_extra_rounds=0, + round_base_delay=0.0, + on_terminal=on_terminal, + ) + assert cancelled.is_set() # the canceller fired on the terminal error + + +# --- #1 REAL RACE: a stale handle never cancels a reused conn ------------- + + +def _normal_keepalive_server(n_requests: int): + """Serve ``n_requests`` requests on ONE accepted connection (keep-alive), + each 200 OK, so urllib3 returns the conn to the pool for reuse. + """ + srv = socket.socket() + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(("127.0.0.1", 0)) + srv.listen(1) + port = srv.getsockname()[1] + + def serve() -> None: + conn, _ = srv.accept() + conn.settimeout(3.0) + try: + for _ in range(n_requests): + data = b"" + while b"\r\n\r\n" not in data: + chunk = conn.recv(65536) + if not chunk: + return + data += chunk + cl = 0 + for line in data.split(b"\r\n"): + if line.lower().startswith(b"content-length:"): + cl = int(line.split(b":")[1]) + have = len(data.split(b"\r\n\r\n", 1)[1]) + while have < cl: + have += len(conn.recv(65536)) + conn.sendall(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok") + except OSError: + pass + + threading.Thread(target=serve, daemon=True).start() + return srv, port + + +def test_stale_handle_cannot_cancel_reused_conn() -> None: + """The actual cross-part race (Codex #1): part A's connection is returned to + the pool after a successful PUT and REUSED by part B. Part A's handle must be + disarmed the moment A's response completed — so even if A's old handle is + 'fired' late, it is a no-op and cannot shut down B's (the reused) socket. + Driven through real pool reuse, not a forced callback. + """ + from hotdata._upload_watchdog import _Arming, set_active_arming + + srv, port = _normal_keepalive_server(2) + pool = urllib3.HTTPConnectionPool( + "127.0.0.1", port, maxsize=1, retries=urllib3.Retry(total=0) + ) + install_watchdog_connection_classes(pool) + wd = _Watchdog() + try: + # Part A: a successful PUT. Capture A's handle(s). + arming_a = _Arming(wd, 60.0) + set_active_arming(arming_a) + try: + resp_a = pool.urlopen("PUT", "/a", body=b"x" * 128, headers={"Content-Length": "128"}) + assert resp_a.status == 200 + finally: + for h in arming_a.all: + h.finish() + wd.disarm(h) + set_active_arming(None) + a_handle = arming_a.all[-1] + # A's handle is disarmed (done) AND its conn ref cleared the moment A's + # response completed — so firing A's stale handle is a guaranteed no-op: + # it returns None and has no socket to shut down, no matter who reuses + # the underlying connection next. + assert a_handle.state == "done" + assert a_handle.conn is None + assert a_handle.cancel_if_active() is None + + # Part B reuses the SAME pooled connection. Track B's live conn during + # its in-flight request to prove the connection object is genuinely + # shared (the race is real), and that A's handle never targeted it. + b_conn_seen: List[object] = [] + arming_b = _Arming(wd, 60.0) + set_active_arming(arming_b) + try: + # Capture B's conn while in flight (handle holds it before getresponse + # disarms). We assert reuse by comparing the HTTPConnection identity. + resp_b = pool.urlopen("PUT", "/b", body=b"y" * 128, headers={"Content-Length": "128"}) + assert resp_b.status == 200 + b_conn_seen.append(arming_b.all[-1]) + finally: + for h in arming_b.all: + h.finish() + wd.disarm(h) + set_active_arming(None) + # Firing A's stale handle now (simulating a late watchdog wakeup) does + # nothing — the core guarantee that prevents cross-part socket corruption. + assert a_handle.cancel_if_active() is None + finally: + srv.close() + wd.shutdown() + + +# --- #5 REAL RACE: sticky abort cancels a handle armed AFTER cancel_all ---- + + +def test_sticky_abort_cancels_handle_armed_after_cancel_all() -> None: + """A sibling that arms its FIRST handle AFTER cancel_all()'s snapshot must + still be cancelled at once (Codex #5) — the abort is sticky, not a one-shot + snapshot. Real arming order, no forced callback. + """ + wd = _Watchdog() + shut: List[str] = [] + + class _Conn: + def __init__(self): + self.sock = self + + def shutdown(self, how): + shut.append("x") + + # Terminal error fires cancel_all() while no handle is registered yet. + wd.cancel_all() + # A sibling now arms its first handle (it was still minting/reading at + # snapshot time). It must be cancelled immediately, not waited out. + late = _PartHandle() + late.publish(_Conn()) + wd.arm(late, deadline=600.0) # 10-min deadline; must NOT be waited out + wd.shutdown() + assert late.state == "cancelled" + assert shut == ["x"] # its socket was shut down at once + + +# --- #5 pre-connection abort gap (real subclass, no socket yet) ----------- + + +def test_aborted_watchdog_refuses_new_connection_without_blocking() -> None: + """Codex #5 pre-connection gap: after cancel_all(), a part that starts a NEW + connection must fail fast at request()/connect() — BEFORE opening a socket — + not block in DNS/TCP/TLS. Uses the REAL WatchdogHTTPConnection through + urlopen against an unroutable/dead address that would otherwise hang connect. + """ + from hotdata._upload_watchdog import _Arming, set_active_arming + + wd = _Watchdog() + wd.cancel_all() # terminal abort BEFORE any connection exists + + # Point at a port nothing listens on; without the precheck, connect() would + # block/until-timeout. The abort precheck must raise before we even try. + pool = urllib3.HTTPConnectionPool( + "127.0.0.1", 0, retries=urllib3.Retry(total=0), + timeout=urllib3.Timeout(connect=30.0), # long — proves we DON'T reach connect + ) + install_watchdog_connection_classes(pool) + + arming = _Arming(wd, 60.0) + set_active_arming(arming) + started = time.monotonic() + try: + with pytest.raises(urllib3.exceptions.HTTPError): + pool.urlopen("PUT", "/x", body=b"z" * 64, headers={"Content-Length": "64"}) + finally: + set_active_arming(None) + wd.shutdown() + elapsed = time.monotonic() - started + # Failed via the precheck, NOT by waiting out the 30s connect timeout. + assert elapsed < 5.0 + # No handle was armed (we bailed before arming/connecting). + assert arming.current is None + + +def test_abort_precheck_does_not_break_happy_path() -> None: + """Guard: the precheck must NOT interfere with a normal (non-aborted) request + — it proceeds and succeeds, and per-attempt arm/disarm (#1) still works. + """ + from hotdata._upload_watchdog import _Arming, set_active_arming + + srv, port = _normal_keepalive_server(1) + pool = urllib3.HTTPConnectionPool("127.0.0.1", port, retries=urllib3.Retry(total=0)) + install_watchdog_connection_classes(pool) + wd = _Watchdog() # NOT aborted + arming = _Arming(wd, 60.0) + set_active_arming(arming) + try: + resp = pool.urlopen("PUT", "/ok", body=b"y" * 64, headers={"Content-Length": "64"}) + assert resp.status == 200 + # A handle was armed and then disarmed by getresponse() (done). + assert arming.all and all(h.state == "done" for h in arming.all) + assert arming.current is None + finally: + set_active_arming(None) + wd.shutdown() + srv.close() + + +# --- strict mint validation (Item 6) -------------------------------------- + + +def _make_api(): + from hotdata import ApiClient, Configuration + + cfg = Configuration(host="https://api.hotdata.test", api_key="k", workspace_id="ws") + return uploads_mod.UploadsApi(ApiClient(cfg)) + + +def _streaming_session(part_size): + return UploadSessionResponse( + finalize_token="tok", headers={}, mode="multipart", + part_size=part_size, part_urls=None, upload_id="up", + ) + + +def test_mint_duplicate_part_rejected(monkeypatch) -> None: + api = _make_api() + + def dup_mint(upload_id, token, req, **kw): + n = req.part_numbers[0] + return MintUploadPartsResponse(parts=[ + MintedUploadPartResponse(part_number=n, url="https://s/a"), + MintedUploadPartResponse(part_number=n, url="https://s/b"), + ]) + + monkeypatch.setattr(api, "mint_upload_parts_handler", dup_mint) + with pytest.raises(MalformedSessionError, match="duplicate"): + api._mint_part(_streaming_session(5 * MIB), 1) + + +def test_mint_extra_part_rejected(monkeypatch) -> None: + api = _make_api() + + def extra_mint(upload_id, token, req, **kw): + return MintUploadPartsResponse(parts=[ + MintedUploadPartResponse(part_number=1, url="https://s/a"), + MintedUploadPartResponse(part_number=2, url="https://s/b"), # not requested + ]) + + monkeypatch.setattr(api, "mint_upload_parts_handler", extra_mint) + with pytest.raises(MalformedSessionError, match="not requested"): + api._mint_part(_streaming_session(5 * MIB), 1) + + +def test_mint_missing_part_rejected(monkeypatch) -> None: + api = _make_api() + monkeypatch.setattr( + api, + "mint_upload_parts_handler", + lambda *a, **k: MintUploadPartsResponse(parts=[]), + ) + with pytest.raises(MalformedSessionError, match="did not return a URL"): + api._mint_part(_streaming_session(5 * MIB), 1) diff --git a/tests/test_upload_multipart_resilient.py b/tests/test_upload_multipart_resilient.py new file mode 100644 index 0000000..63d6272 --- /dev/null +++ b/tests/test_upload_multipart_resilient.py @@ -0,0 +1,354 @@ +"""Integration tests for the resilient multipart path: the per-part unit driven +by the re-sweep loop, exercising eager (known-size) and streaming (mint-on-demand) +URL acquisition, the one-shot 403-only re-mint, 5xx → outer re-sweep, blank-ETag +recovery, and finalize completeness. + +These stub the two seams (create/finalize/mint via monkeypatch, storage PUTs via +a sequencing fake pool) — no real server. The fake pool supports per-URL response +*sequences* so a 500-then-200 or 403-then-200 can be scripted per path. +""" + +from __future__ import annotations + +import threading +from typing import Any, Dict, List, Optional, Tuple + +import pytest + +import hotdata.uploads as uploads_mod +from hotdata.uploads import ( + MIB, + MalformedSessionError, + StorageError, + UploadsApi, +) +from hotdata.models.finalize_upload_response import FinalizeUploadResponse +from hotdata.models.minted_upload_part_response import MintedUploadPartResponse +from hotdata.models.mint_upload_parts_response import MintUploadPartsResponse +from hotdata.models.upload_session_response import UploadSessionResponse + + +# --- sequencing fake storage pool ----------------------------------------- + + +class _Resp: + def __init__(self, status: int, *, etag: Optional[str] = None, body: bytes = b""): + self.status = status + self._headers = {} + if etag is not None: + self._headers["ETag"] = etag + self.data = body + + @property + def headers(self): + class _H: + def __init__(self, d): + self._d = {k.lower(): v for k, v in d.items()} + + def get(self, k, default=None): + return self._d.get(k.lower(), default) + + return _H(self._headers) + + +class _SeqPool: + """Fake storage pool. ``responses[url]`` is a list popped per call (last + entry reused once exhausted), so a 500-then-200 sequence is scriptable. + Records each PUT. + """ + + def __init__(self) -> None: + self.responses: Dict[str, List[_Resp]] = {} + self.default = _Resp(200, etag='"default"') + self.calls: List[Dict[str, Any]] = [] + self._lock = threading.Lock() + + def request(self, method: str, url: str, *, body=None, headers=None, retries=None, **kw): + with self._lock: + self.calls.append({"method": method, "url": url, "body": body}) + seq = self.responses.get(url) + if not seq: + return self.default + return seq.pop(0) if len(seq) > 1 else seq[0] + + +@pytest.fixture +def seq_pool(monkeypatch): + pool = _SeqPool() + monkeypatch.setattr(uploads_mod, "_STORAGE_POOL", pool) + # The resilient path installs the watchdog ConnectionCls on the pool; the + # fake doesn't need it, but install is a no-op-safe call on a non-pool, so + # stub it out for the fake. + monkeypatch.setattr(uploads_mod, "install_watchdog_connection_classes", lambda p: None) + return pool + + +def _make_api() -> UploadsApi: + from hotdata import ApiClient, Configuration + + cfg = Configuration(host="https://api.hotdata.test", api_key="k", workspace_id="ws") + return UploadsApi(ApiClient(cfg)) + + +def _patch_finalize(monkeypatch, api, capture): + def fake_finalize(upload_id, finalize_token, body=None, _request_timeout=None): + from datetime import datetime, timezone + + capture.append(("finalize", upload_id, body)) + return FinalizeUploadResponse( + content_type=None, + created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + size_bytes=0, + status="ready", + upload_id=upload_id, + ) + + monkeypatch.setattr(api, "_finalize_handler", fake_finalize) + + +def _eager_session(part_urls, part_size): + return UploadSessionResponse( + finalize_token="tok", + headers={}, + mode="multipart", + part_size=part_size, + part_urls=part_urls, + upload_id="up_eager", + ) + + +def _streaming_session(part_size): + return UploadSessionResponse( + finalize_token="tok", + headers={}, + mode="multipart", + part_size=part_size, + part_urls=None, # streaming: mint on demand + upload_id="up_stream", + ) + + +# --- 5xx → outer re-sweep (tier boundary, Item 5) ------------------------- + + +def test_part_500_recovered_by_outer_resweep(seq_pool, monkeypatch): + part_size = 5 * MIB + content = b"a" * part_size + api = _make_api() + url = "https://storage.test/p1" + seq_pool.responses[url] = [_Resp(500, body=b"boom"), _Resp(200, etag='"e1"')] + monkeypatch.setattr( + api, "create_upload_session_handler", lambda *a, **k: _eager_session([url], part_size) + ) + capture: List[Any] = [] + _patch_finalize(monkeypatch, api, capture) + + api.upload_file(content, max_concurrency=2, round_base_delay=0.0) + + # 500 was NOT inner-retried (would have inner-swallowed it); the outer + # re-sweep re-PUT and succeeded → part PUT twice. + puts = [c for c in seq_pool.calls if c["url"] == url] + assert len(puts) == 2 + fin = capture[-1] + parts = fin[2].parts + assert [p.part_number for p in parts] == [1] + assert parts[0].e_tag == '"e1"' + + +# --- streaming: mint per part --------------------------------------------- + + +def test_streaming_session_mints_per_part(seq_pool, monkeypatch): + part_size = 5 * MIB + content = b"b" * (2 * part_size) + api = _make_api() + mint_calls: List[List[int]] = [] + + def fake_mint(upload_id, x_upload_finalize_token, mint_upload_parts_request, **kw): + nums = list(mint_upload_parts_request.part_numbers) + mint_calls.append(nums) + parts = [ + MintedUploadPartResponse(part_number=n, url=f"https://storage.test/spart/{n}") + for n in nums + ] + return MintUploadPartsResponse(parts=parts) + + for n in (1, 2): + seq_pool.responses[f"https://storage.test/spart/{n}"] = [_Resp(200, etag=f'"e{n}"')] + monkeypatch.setattr( + api, "create_upload_session_handler", lambda *a, **k: _streaming_session(part_size) + ) + monkeypatch.setattr(api, "mint_upload_parts_handler", fake_mint) + capture: List[Any] = [] + _patch_finalize(monkeypatch, api, capture) + + api.upload_file(content, max_concurrency=2, round_base_delay=0.0) + + # One mint per part, each PUT to its minted URL. + assert sorted(n for nums in mint_calls for n in nums) == [1, 2] + fin = capture[-1] + assert [p.part_number for p in fin[2].parts] == [1, 2] + assert [p.e_tag for p in fin[2].parts] == ['"e1"', '"e2"'] + + +def test_streaming_mint_finite_timeout(seq_pool, monkeypatch): + part_size = 5 * MIB + api = _make_api() + seen_timeout: List[Any] = [] + + def fake_mint(upload_id, x_upload_finalize_token, mint_upload_parts_request, **kw): + seen_timeout.append(kw.get("_request_timeout")) + n = mint_upload_parts_request.part_numbers[0] + return MintUploadPartsResponse( + parts=[MintedUploadPartResponse(part_number=n, url=f"https://storage.test/s/{n}")] + ) + + seq_pool.responses["https://storage.test/s/1"] = [_Resp(200, etag='"e1"')] + monkeypatch.setattr( + api, "create_upload_session_handler", lambda *a, **k: _streaming_session(part_size) + ) + monkeypatch.setattr(api, "mint_upload_parts_handler", fake_mint) + _patch_finalize(monkeypatch, api, []) + + api.upload_file(b"c" * part_size, round_base_delay=0.0) + # The mint must carry a finite (connect, read) timeout, not None. + assert seen_timeout[0] == uploads_mod.MINT_TIMEOUT + + +# --- eager 403: one-shot inline re-mint (Codex #7) ------------------------ + + +def test_eager_403_reminted_once_and_recovers(seq_pool, monkeypatch): + part_size = 5 * MIB + api = _make_api() + eager_url = "https://storage.test/eager1" + fresh_url = "https://storage.test/fresh1" + seq_pool.responses[eager_url] = [_Resp(403, body=b"SignatureDoesNotMatch")] + seq_pool.responses[fresh_url] = [_Resp(200, etag='"fresh-e1"')] + + mint_calls: List[List[int]] = [] + + def fake_mint(upload_id, x_upload_finalize_token, mint_upload_parts_request, **kw): + mint_calls.append(list(mint_upload_parts_request.part_numbers)) + return MintUploadPartsResponse( + parts=[MintedUploadPartResponse(part_number=1, url=fresh_url)] + ) + + monkeypatch.setattr( + api, "create_upload_session_handler", lambda *a, **k: _eager_session([eager_url], part_size) + ) + monkeypatch.setattr(api, "mint_upload_parts_handler", fake_mint) + capture: List[Any] = [] + _patch_finalize(monkeypatch, api, capture) + + api.upload_file(b"d" * part_size, round_base_delay=0.0) + + # Exactly one inline re-mint for part 1 after the 403. + assert mint_calls == [[1]] + assert any(c["url"] == fresh_url for c in seq_pool.calls) + assert capture[-1][2].parts[0].e_tag == '"fresh-e1"' + + +def test_eager_5xx_does_not_remint(seq_pool, monkeypatch): + part_size = 5 * MIB + api = _make_api() + eager_url = "https://storage.test/eager1" + # 500 forever → never inline-remints; surfaces after the re-sweep budget. + seq_pool.responses[eager_url] = [_Resp(500, body=b"boom")] + + mint_calls: List[Any] = [] + monkeypatch.setattr( + api, + "mint_upload_parts_handler", + lambda *a, **k: mint_calls.append(1) or MintUploadPartsResponse(parts=[]), + ) + monkeypatch.setattr( + api, "create_upload_session_handler", lambda *a, **k: _eager_session([eager_url], part_size) + ) + _patch_finalize(monkeypatch, api, []) + + with pytest.raises(StorageError) as ei: + api.upload_file(b"e" * part_size, max_extra_rounds=1, round_base_delay=0.0) + assert ei.value.status == 500 + # A 5xx is NOT an expiry signal: never inline-reminted. + assert mint_calls == [] + + +# --- blank ETag is retryable (Item 7) ------------------------------------- + + +def test_blank_etag_retried_then_recovers(seq_pool, monkeypatch): + part_size = 5 * MIB + api = _make_api() + url = "https://storage.test/p1" + # First PUT returns 200 with a whitespace-only ETag → MissingETagError + # (retryable); re-sweep gets a real one. + seq_pool.responses[url] = [_Resp(200, etag=" "), _Resp(200, etag='"real"')] + monkeypatch.setattr( + api, "create_upload_session_handler", lambda *a, **k: _eager_session([url], part_size) + ) + capture: List[Any] = [] + _patch_finalize(monkeypatch, api, capture) + + api.upload_file(b"f" * part_size, round_base_delay=0.0) + assert capture[-1][2].parts[0].e_tag == '"real"' + + +# --- mint response validation (Item 6) ------------------------------------ + + +def test_storage_3xx_surfaces_as_error(seq_pool, monkeypatch): + # Codex #6: with redirects disabled, a storage 3xx must surface as a + # StorageError (following it would invalidate the presigned signature), not + # be silently treated as success. + part_size = 5 * MIB + api = _make_api() + url = "https://storage.test/p1" + seq_pool.responses[url] = [_Resp(302, body=b"redirect")] + monkeypatch.setattr( + api, "create_upload_session_handler", lambda *a, **k: _eager_session([url], part_size) + ) + _patch_finalize(monkeypatch, api, []) + with pytest.raises(StorageError) as ei: + api.upload_file(b"h" * part_size, max_extra_rounds=0, round_base_delay=0.0) + assert ei.value.status == 302 + + +def test_bounded_retry_multiplication(seq_pool, monkeypatch): + # The inner urllib3.Retry total and the outer re-sweep rounds MULTIPLY: a + # part PUT that always fails transiently is attempted at most + # (1 + max_extra_rounds) times by the OUTER loop. (The inner Retry total is a + # urllib3 concern, exercised separately in test_upload_retry_policy.) Here we + # bound the OUTER multiplication: a persistent 500 → exactly 1+rounds PUTs. + part_size = 5 * MIB + api = _make_api() + url = "https://storage.test/p1" + seq_pool.responses[url] = [_Resp(500, body=b"boom")] # always 500 + monkeypatch.setattr( + api, "create_upload_session_handler", lambda *a, **k: _eager_session([url], part_size) + ) + _patch_finalize(monkeypatch, api, []) + with pytest.raises(StorageError): + api.upload_file(b"i" * part_size, max_extra_rounds=2, round_base_delay=0.0) + puts = [c for c in seq_pool.calls if c["url"] == url] + assert len(puts) == 3 # 1 initial round + 2 re-sweeps, bounded + + +def test_streaming_mint_missing_requested_part_is_rejected(seq_pool, monkeypatch): + part_size = 5 * MIB + api = _make_api() + + def bad_mint(upload_id, x_upload_finalize_token, mint_upload_parts_request, **kw): + # Return a DIFFERENT part number than requested. + return MintUploadPartsResponse( + parts=[MintedUploadPartResponse(part_number=999, url="https://storage.test/x")] + ) + + monkeypatch.setattr( + api, "create_upload_session_handler", lambda *a, **k: _streaming_session(part_size) + ) + monkeypatch.setattr(api, "mint_upload_parts_handler", bad_mint) + _patch_finalize(monkeypatch, api, []) + + with pytest.raises(MalformedSessionError): + api.upload_file(b"g" * part_size, max_extra_rounds=0, round_base_delay=0.0) diff --git a/tests/test_upload_resilient.py b/tests/test_upload_resilient.py new file mode 100644 index 0000000..d93c0d8 --- /dev/null +++ b/tests/test_upload_resilient.py @@ -0,0 +1,233 @@ +"""Tests for the whole-upload re-sweep driver (Item 1) and the per-part work +unit (Items 4, 6, 7, 8). + +The driver `_upload_parts_resilient` re-runs only the parts still failing after +each round, at halving concurrency, with completed parts keeping their ETags; +exactly-once for healthy parts; terminal (non-retryable) errors fail fast. + +Most driver tests use a transport-free fake per-part callable (`_FakeUploader`) +— no HTTP. The eager/streaming/mint behaviors use the real `_upload_multipart` +with the storage pool + mint handler stubbed. +""" + +from __future__ import annotations + +import threading +from typing import Callable, Dict, List, Optional, Tuple + +import pytest + +from hotdata.models.finalize_upload_part import FinalizeUploadPart +from hotdata.uploads import ( + MalformedSessionError, + StorageError, + StorageTransportError, + UploadError, + _PartPlan, + _round_in_flight, + _round_delay, + _upload_parts_resilient, +) + + +# --- pure helpers: concurrency halving + round backoff -------------------- + + +def test_round_in_flight_halves_each_round_min_one() -> None: + assert _round_in_flight(8, 0) == 8 + assert _round_in_flight(8, 1) == 4 + assert _round_in_flight(8, 2) == 2 + assert _round_in_flight(8, 3) == 1 + assert _round_in_flight(8, 4) == 1 # floored at 1 + assert _round_in_flight(1, 0) == 1 + assert _round_in_flight(8, 99) == 1 # large round → 1, no overflow + + +def test_round_delay_grows_exponentially() -> None: + # base 2s → 2, 4, 8 before rounds 1, 2, 3; clamped so it can't overflow. + assert _round_delay(1, base_delay=2.0) == pytest.approx(2.0) + assert _round_delay(2, base_delay=2.0) == pytest.approx(4.0) + assert _round_delay(3, base_delay=2.0) == pytest.approx(8.0) + assert _round_delay(99, base_delay=2.0) < float("inf") + + +# --- transport-free fake per-part uploader -------------------------------- + + +class _FakeUploader: + """Records attempts per part number; can force the first K attempts of named + parts to fail with a configurable exception; tracks peak concurrency. + """ + + def __init__( + self, + *, + fail_counts: Optional[Dict[int, int]] = None, + failure_factory: Optional[Callable[[int], BaseException]] = None, + etag_for: Optional[Callable[[int], str]] = None, + ) -> None: + self._fail_counts = dict(fail_counts or {}) + self._failure_factory = failure_factory or ( + lambda pn: StorageTransportError(url=f"u{pn}", part_number=pn) + ) + self._etag_for = etag_for or (lambda pn: f'"etag-{pn}"') + self.attempts: Dict[int, int] = {} + self._active = 0 + self.peak = 0 + self._lock = threading.Lock() + + def __call__(self, plan: "_PartPlan") -> FinalizeUploadPart: + pn = plan.part_number + with self._lock: + self.attempts[pn] = self.attempts.get(pn, 0) + 1 + self._active += 1 + self.peak = max(self.peak, self._active) + attempt_no = self.attempts[pn] + try: + remaining_failures = self._fail_counts.get(pn, 0) + if attempt_no <= remaining_failures: + raise self._failure_factory(pn) + return FinalizeUploadPart(e_tag=self._etag_for(pn), part_number=pn) + finally: + with self._lock: + self._active -= 1 + + +def _plans(n: int, part_size: int = 1024) -> List[_PartPlan]: + return [ + _PartPlan(index=i, part_number=i + 1, offset=i * part_size, length=part_size) + for i in range(n) + ] + + +def _run(uploader, n, *, base_in_flight=4, max_extra_rounds=3): + return _upload_parts_resilient( + _plans(n), + uploader, + base_in_flight=base_in_flight, + max_extra_rounds=max_extra_rounds, + round_base_delay=0.0, # no real sleeps in tests + ) + + +# --- the bug + the fix ---------------------------------------------------- + + +def test_repro_single_part_blip_sinks_without_rounds() -> None: + # max_extra_rounds=0 reproduces the legacy "one failure sinks everything". + up = _FakeUploader(fail_counts={3: 1}) + with pytest.raises(StorageTransportError): + _run(up, 5, max_extra_rounds=0) + + +def test_single_part_blip_recovers_on_later_round() -> None: + up = _FakeUploader(fail_counts={3: 1}) + parts = _run(up, 5) + assert [p.part_number for p in parts] == [1, 2, 3, 4, 5] + assert [p.e_tag for p in parts] == [f'"etag-{i}"' for i in range(1, 6)] + assert up.attempts[3] == 2 # failed once, succeeded on the re-sweep + for pn in (1, 2, 4, 5): + assert up.attempts[pn] == 1 # healthy parts attempted exactly once + + +def test_multiple_flaky_parts_all_recover() -> None: + up = _FakeUploader(fail_counts={1: 2, 3: 1, 5: 3}) + parts = _run(up, 6) + assert [p.part_number for p in parts] == [1, 2, 3, 4, 5, 6] + assert up.attempts[1] == 3 + assert up.attempts[3] == 2 + assert up.attempts[5] == 4 + assert up.attempts[2] == up.attempts[4] == up.attempts[6] == 1 + + +def test_permanent_failure_surfaced_after_rounds() -> None: + up = _FakeUploader(fail_counts={2: 99}) + with pytest.raises(StorageTransportError): + _run(up, 4, max_extra_rounds=3) + # 1 initial + 3 re-sweeps = 4 attempts for the doomed part. + assert up.attempts[2] == 4 + # Healthy parts attempted exactly once. + for pn in (1, 3, 4): + assert up.attempts[pn] == 1 + + +def test_happy_path_each_part_once() -> None: + up = _FakeUploader() + parts = _run(up, 10) + assert len(parts) == 10 + assert all(v == 1 for v in up.attempts.values()) + + +def test_concurrency_never_exceeds_base_cap() -> None: + # 20 parts, base cap 3: many concurrent, but never more than 3 in flight. + up = _FakeUploader() + _run(up, 20, base_in_flight=3, max_extra_rounds=0) + assert up.peak <= 3 + + +# --- terminal errors fail fast (Item 2) ----------------------------------- + + +def test_terminal_error_fails_fast_without_resweeping() -> None: + up = _FakeUploader( + fail_counts={2: 99}, + failure_factory=lambda pn: MalformedSessionError(f"contract violation part {pn}"), + ) + with pytest.raises(MalformedSessionError): + _run(up, 4) + # Terminal: attempted exactly once, never re-swept. + assert up.attempts[2] == 1 + + +def test_generic_uploaderror_not_retried() -> None: + # A deterministic local short-read (generic UploadError) is non-retryable. + up = _FakeUploader( + fail_counts={2: 99}, + failure_factory=lambda pn: UploadError("deterministic short read"), + ) + with pytest.raises(UploadError): + _run(up, 3) + assert up.attempts[2] == 1 + + +# --- result completeness -------------------------------------------------- + + +def test_missing_slot_is_hard_error() -> None: + # An uploader that returns None for a part (a bug) must not silently produce + # a short finalize list. + def bad(plan): + if plan.part_number == 2: + return None # type: ignore[return-value] + return FinalizeUploadPart(e_tag=f'"e{plan.part_number}"', part_number=plan.part_number) + + with pytest.raises(MalformedSessionError, match="never uploaded|part 2"): + _upload_parts_resilient( + _plans(3), bad, base_in_flight=2, max_extra_rounds=0, round_base_delay=0.0 + ) + + +# --- progress exactly-once under retry (Item 8) --------------------------- + + +def test_progress_no_double_count_under_resweep() -> None: + ticks: List[int] = [] + lock = threading.Lock() + + def on_progress(done: int) -> None: + with lock: + ticks.append(done) + + up = _FakeUploader(fail_counts={2: 1}) + parts = _upload_parts_resilient( + _plans(3, part_size=100), + up, + base_in_flight=2, + max_extra_rounds=3, + round_base_delay=0.0, + on_part_done=lambda plan: on_progress(plan.length), + ) + assert len(parts) == 3 + # Part 2 failed round 0 then succeeded: counted exactly once. 3 parts × 100. + assert sum(ticks) == 300 + assert len(ticks) == 3 # one tick per successful part, no double-count diff --git a/tests/test_upload_retry_policy.py b/tests/test_upload_retry_policy.py new file mode 100644 index 0000000..feaacaa --- /dev/null +++ b/tests/test_upload_retry_policy.py @@ -0,0 +1,201 @@ +"""Tests for the part-PUT retry policy + per-part timeout + error classification +(Items 2, 3, 5). + +The inner retry tier is a narrowed ``urllib3.Retry`` (``ConnectionResetRetry`` +base) on the storage pool, reached via ``pool.urlopen``: + +* 429 retried (honoring ``Retry-After``); pre-response connection resets retried + on any method (incl. the watchdog's ``ProtocolError(RemoteDisconnected)``); +* storage **5xx EXCLUDED** from the forcelist — it propagates to the OUTER + re-sweep (the tier boundary, Item 5 / Flag #5); +* ``backoff_max`` ~120s caps a pathological ``Retry-After`` (the deadline analog). + +``_part_put_timeout`` is the size-scaled per-attempt transfer cap the watchdog +enforces. ``_is_retryable`` is the allowlist the outer re-sweep consults. +""" + +from __future__ import annotations + +import errno +from http.client import RemoteDisconnected + +import pytest +from urllib3.exceptions import ProtocolError +from urllib3.util.retry import Retry + +from hotdata._retry import ConnectionResetRetry +from hotdata.uploads import ( + MIB, + MalformedSessionError, + MissingETagError, + SizeLimitError, + StorageError, + StorageTransportError, + UploadError, + _is_retryable, + _parse_etag, + _part_put_timeout, + default_part_retry, +) +from hotdata.exceptions import ApiException + + +# --- default_part_retry narrowing (Item 5) -------------------------------- + + +def test_default_part_retry_excludes_5xx_keeps_429() -> None: + retry = default_part_retry() + forcelist = set(retry.status_forcelist or ()) + assert 429 in forcelist + for code in (500, 502, 503, 504): + assert code not in forcelist, f"{code} must propagate to the outer re-sweep" + + +def test_default_part_retry_is_connection_reset_retry() -> None: + # Based on ConnectionResetRetry so a pre-response reset (and the watchdog's + # ProtocolError) is retried on PUT, not method-gated away. + assert isinstance(default_part_retry(), ConnectionResetRetry) + + +def test_default_part_retry_does_not_raise_on_status() -> None: + assert default_part_retry().raise_on_status is False + + +def test_default_part_retry_honors_retry_after() -> None: + assert default_part_retry().respect_retry_after_header is True + + +def test_default_part_retry_does_not_follow_redirects() -> None: + # A 3xx on a presigned PUT would re-issue to a new URL and invalidate the + # signature; never chase it (Codex #6). + retry = default_part_retry() + assert retry.redirect is False or retry.redirect == 0 + + +def test_default_part_retry_caps_backoff_near_120s() -> None: + # backoff_max is the deadline analog: a pathological Retry-After/backoff is + # capped so the round-join can't stall on an unbounded sleep. + retry = default_part_retry() + # urllib3 >=2 exposes the cap as `backoff_max`; 1.26 as `DEFAULT_BACKOFF_MAX` + # (the value `get_backoff_time` actually reads). Either way it must be <=120. + backoff_max = getattr(retry, "backoff_max", None) + if backoff_max is None: + backoff_max = retry.DEFAULT_BACKOFF_MAX + assert backoff_max <= 120 + + +# --- the watchdog ProtocolError is classified as a pre-response reset ------ + + +def test_retry_after_capped_at_120s() -> None: + """A server ``Retry-After`` is honored but CAPPED at 120s (the per-part + retry-window deadline analog) so the re-sweep's round-boundary join can't + stall on a multi-minute server value. ``backoff_max`` only bounds the + exponential path; the cap lives in ``get_retry_after``. + """ + retry = default_part_retry() + + class _Resp: + def __init__(self, retry_after: str): + self.headers = {"Retry-After": retry_after} + + def getheader(self, name, default=None): + return self.headers.get(name, default) + + # A hostile hour-long Retry-After is clamped to 120s. + assert retry.get_retry_after(_Resp("3600")) == 120 + # A reasonable value is honored as-is. + assert retry.get_retry_after(_Resp("5")) == 5 + + +def test_no_retry_after_passes_through_as_none() -> None: + # No Retry-After header → None (the exponential backoff path is unaffected). + retry = default_part_retry() + + class _NoHeader: + headers: dict = {} + + def getheader(self, name, default=None): + return default + + assert retry.get_retry_after(_NoHeader()) is None + + +def test_connection_reset_retry_classifies_watchdog_protocolerror() -> None: + """The decisive Item-3/Item-5 check (architect-flagged): the watchdog's + ``sock.shutdown`` surfaces through ``urlopen`` as + ``ProtocolError(RemoteDisconnected)`` — the SAME shape ConnectionResetRetry + already treats as a (retryable) connection error. If this regresses, a + watchdog-cancelled part would NOT be re-swept. + """ + retry = default_part_retry() + watchdog_err = ProtocolError( + "Connection aborted.", + RemoteDisconnected("Remote end closed connection without response"), + ) + assert retry._is_connection_error(watchdog_err) + # And a plain reset (stale pooled socket) too. + reset_err = ProtocolError("Connection aborted.", ConnectionResetError(errno.ECONNRESET, "reset")) + assert retry._is_connection_error(reset_err) + + +# --- _part_put_timeout: size-scaled per-attempt transfer cap (Item 3) ------ + + +def test_part_put_timeout_scales_and_caps() -> None: + assert _part_put_timeout(8 * MIB) == pytest.approx(188.0) + assert _part_put_timeout(64 * MIB) == pytest.approx(60.0 + 1024.0) + assert _part_put_timeout(0) == pytest.approx(60.0) + # Monotonic non-decreasing below the cap. + assert _part_put_timeout(MIB) <= _part_put_timeout(2 * MIB) <= _part_put_timeout(8 * MIB) + # Cap at 30 minutes for very large parts. + assert _part_put_timeout(5 * 1024 * MIB) == pytest.approx(1800.0) + assert _part_put_timeout(2**63) == pytest.approx(1800.0) + + +# --- _is_retryable allowlist (Item 2) ------------------------------------- + + +def test_is_retryable_allowlist_true_cases() -> None: + assert _is_retryable(StorageTransportError(url="u", part_number=1)) + assert _is_retryable(MissingETagError(part_number=1)) + assert _is_retryable(StorageError(status=500, part_number=1, body="x")) + assert _is_retryable(StorageError(status=403, part_number=1, body="x")) # eager URL expiry + assert _is_retryable(ApiException(status=503)) # mint transient + + +def test_is_retryable_allowlist_false_cases() -> None: + # Deterministic / contract / unexpected — NOT retried (no wasted rounds, no + # duplicate PUT on a callback bug). + assert not _is_retryable(MalformedSessionError("bad")) + assert not _is_retryable(SizeLimitError(what="declared_size_bytes", value=2**63)) + assert not _is_retryable(UploadError("deterministic short read")) # generic + assert not _is_retryable(ValueError("unexpected")) + assert not _is_retryable(RuntimeError("callback blew up")) + + +# --- _parse_etag: reject missing AND blank (Item 7) ----------------------- + + +class _Headers: + def __init__(self, data): + self._data = {k.lower(): v for k, v in data.items()} + + def get(self, key, default=None): + return self._data.get(key.lower(), default) + + +def test_parse_etag_passes_quoted_value_verbatim() -> None: + assert _parse_etag(_Headers({"ETag": '"abc123"'}), part_number=1) == '"abc123"' + + +def test_parse_etag_rejects_missing() -> None: + with pytest.raises(MissingETagError): + _parse_etag(_Headers({}), part_number=2) + + +def test_parse_etag_rejects_empty_and_whitespace() -> None: + with pytest.raises(MissingETagError): + _parse_etag(_Headers({"ETag": ""}), part_number=3) + with pytest.raises(MissingETagError): + _parse_etag(_Headers({"ETag": " "}), part_number=4) diff --git a/tests/test_upload_watchdog.py b/tests/test_upload_watchdog.py new file mode 100644 index 0000000..64d261e --- /dev/null +++ b/tests/test_upload_watchdog.py @@ -0,0 +1,441 @@ +"""Tests for the per-part stall-cancellation watchdog (Item 3). + +The watchdog must terminate a storage PUT that black-holes (the far end accepts +the connection but stops reading the request body, or never sends a response) so +the part fails INTO the outer re-sweep rather than hanging the upload forever. + +The proven mechanism (spike ``conn_subclass_v2.py`` / ``subclass-spike-findings.md`` +Spike 2b): a ``WatchdogHTTPConnection`` subclass injected via ``pool.ConnectionCls`` +publishes its live ``.sock`` to a per-part handle on both ``connect()`` and +``request()``; the per-part *worker* owns the handle lifecycle around the whole +``pool.urlopen()`` call (arm before, ``finish()`` in ``finally``); on deadline the +watchdog calls ``sock.shutdown(SHUT_RDWR)`` (never ``close()`` — fd-reuse safety). + +These tests use REAL sockets (a buffered fake cannot reproduce the kernel +send-buffer wedge), driven through urllib3's real ``urlopen``. +""" + +from __future__ import annotations + +import socket +import threading +import time +from typing import Dict, Optional, Tuple + +import pytest +import urllib3 + +from hotdata._upload_watchdog import ( + _Arming, + _PartHandle, + _Watchdog, + install_watchdog_connection_classes, + set_active_arming, +) + + +# --- real-socket servers -------------------------------------------------- + + +def _normal_server() -> Tuple[socket.socket, int]: + """A server that reads the full request body then replies 200.""" + srv = socket.socket() + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(("127.0.0.1", 0)) + srv.listen(1) + port = srv.getsockname()[1] + + def serve() -> None: + conn, _ = srv.accept() + conn.settimeout(2.0) + data = b"" + try: + while b"\r\n\r\n" not in data: + chunk = conn.recv(65536) + if not chunk: + break + data += chunk + content_length = 0 + for line in data.split(b"\r\n"): + if line.lower().startswith(b"content-length:"): + content_length = int(line.split(b":")[1]) + have = len(data.split(b"\r\n\r\n", 1)[1]) if b"\r\n\r\n" in data else 0 + while have < content_length: + chunk = conn.recv(65536) + if not chunk: + break + have += len(chunk) + conn.sendall(b'HTTP/1.1 200 OK\r\nETag: "ok-etag"\r\nContent-Length: 2\r\n\r\nok') + except OSError: + pass + finally: + time.sleep(0.05) + conn.close() + + threading.Thread(target=serve, daemon=True).start() + return srv, port + + +def _stall_server(read_body: bool) -> Tuple[socket.socket, int, Dict[str, object]]: + """A black-hole server. ``read_body=False`` never reads (SEND stall: the + client wedges in the body write once the kernel send buffer fills); + ``read_body=True`` drains the body but never replies (RESPONSE-read stall). + """ + srv = socket.socket() + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(("127.0.0.1", 0)) + srv.listen(1) + port = srv.getsockname()[1] + hold: Dict[str, object] = {} + + def serve() -> None: + conn, _ = srv.accept() + hold["conn"] = conn + if read_body: + conn.settimeout(0.2) + while not hold.get("stop"): + try: + if not conn.recv(65536): + break + except socket.timeout: + continue + except OSError: + break + else: + while not hold.get("stop"): + time.sleep(0.05) + try: + conn.close() + except OSError: + pass + + threading.Thread(target=serve, daemon=True).start() + return srv, port, hold + + +def _pool(port: int) -> urllib3.HTTPConnectionPool: + pool = urllib3.HTTPConnectionPool( + "127.0.0.1", port, maxsize=2, retries=urllib3.Retry(total=0) + ) + install_watchdog_connection_classes(pool) + return pool + + +def _worker_put( + pool: urllib3.HTTPConnectionPool, + path: str, + body: bytes, + deadline: float, + watchdog: object = None, +) -> Tuple[str, object, float, _Arming]: + """Mirror the real per-part worker: stash an _Arming around urlopen (the + connection subclass arms a fresh handle per attempt), disarm in finally. + Returns the arming so the test can inspect each attempt's handle state. + """ + own_watchdog = watchdog is None + watchdog = watchdog or _Watchdog() + arming = _Arming(watchdog, deadline) + set_active_arming(arming) + started = time.monotonic() + try: + resp = pool.urlopen( + "PUT", path, body=body, headers={"Content-Length": str(len(body))} + ) + return ("ok", resp.status, time.monotonic() - started, arming) + except BaseException as exc: # noqa: BLE001 - the test inspects the type + return ("raised", exc, time.monotonic() - started, arming) + finally: + for handle in arming.all: + handle.finish() + watchdog.disarm(handle) + set_active_arming(None) + if own_watchdog: + watchdog.shutdown() + + +# --- happy path ----------------------------------------------------------- + + +def test_happy_path_through_urlopen_no_fire() -> None: + srv, port = _normal_server() + try: + outcome, status, elapsed, arming = _worker_put( + _pool(port), "/part/1?sig=abc", b"y" * 4096, deadline=5.0 + ) + finally: + srv.close() + assert outcome == "ok" + assert status == 200 + # An attempt's handle was armed and the worker marked it done; never cancelled. + assert arming.all + assert all(h.state == "done" for h in arming.all) + + +# --- send-side black hole (request-body write wedges) --------------------- + + +def test_send_buffer_stall_cancelled() -> None: + srv, port, hold = _stall_server(read_body=False) + try: + # 64 MiB body so the kernel send buffer fills and the worker wedges in send(). + outcome, exc, elapsed, arming = _worker_put( + _pool(port), "/part/1?sig=b", b"z" * (64 * 1024 * 1024), deadline=1.0 + ) + finally: + hold["stop"] = True + srv.close() + assert outcome == "raised" + # urllib3 wraps the shutdown-induced failure as a ProtocolError. + assert isinstance(exc, urllib3.exceptions.HTTPError) + assert any(h.state == "cancelled" for h in arming.all) + # Cancelled promptly after the ~1s deadline, not hung. + assert elapsed < 5.0 + + +# --- response-read black hole (server drains body, never replies) --------- + + +def test_response_read_stall_cancelled() -> None: + srv, port, hold = _stall_server(read_body=True) + try: + outcome, exc, elapsed, arming = _worker_put( + _pool(port), "/part/1?sig=c", b"w" * 8192, deadline=1.0 + ) + finally: + hold["stop"] = True + srv.close() + assert outcome == "raised" + assert isinstance(exc, urllib3.exceptions.HTTPError) + assert any(h.state == "cancelled" for h in arming.all) + assert elapsed < 5.0 + + +# --- Codex #1: reused keep-alive conn is still protected ------------------ + + +def _keepalive_then_stall_server() -> Tuple[socket.socket, int, Dict[str, object]]: + """ONE server on ONE connection: serves the FIRST request normally + (keep-alive, so urllib3 returns the conn to the pool), then on the SECOND + request reads nothing — black-holing the body write on the REUSED conn. + """ + srv = socket.socket() + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(("127.0.0.1", 0)) + srv.listen(1) + port = srv.getsockname()[1] + hold: Dict[str, object] = {} + + def serve() -> None: + conn, _ = srv.accept() + hold["conn"] = conn + # First request: read headers + body, reply 200 with keep-alive. + conn.settimeout(2.0) + data = b"" + try: + while b"\r\n\r\n" not in data: + data += conn.recv(65536) + cl = 0 + for line in data.split(b"\r\n"): + if line.lower().startswith(b"content-length:"): + cl = int(line.split(b":")[1]) + have = len(data.split(b"\r\n\r\n", 1)[1]) + while have < cl: + have += len(conn.recv(65536)) + conn.sendall(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok") + except OSError: + return + # Second request on the SAME conn: never read → send-stall. + while not hold.get("stop"): + time.sleep(0.05) + try: + conn.close() + except OSError: + pass + + threading.Thread(target=serve, daemon=True).start() + return srv, port, hold + + +def test_reused_pooled_conn_send_stall_is_cancelled() -> None: + """A pooled keep-alive connection does NOT call connect() on reuse — so + arming must happen in request() (before the send), or a send-stall on a + REUSED conn would be invisible to the watchdog. One server, one connection: + request 1 succeeds (conn returns to the pool), request 2 reuses it and stalls. + """ + srv, port, hold = _keepalive_then_stall_server() + pool = urllib3.HTTPConnectionPool( + "127.0.0.1", port, maxsize=1, retries=urllib3.Retry(total=0) + ) + install_watchdog_connection_classes(pool) + watchdog = _Watchdog() + try: + outcome1, status, _, _ = _worker_put(pool, "/ok", b"y" * 256, 5.0, watchdog) + assert outcome1 == "ok" and status == 200 + # Reuse the pooled conn; the peer black-holes the body write. + outcome2, exc2, elapsed2, arming2 = _worker_put( + pool, "/stall", b"z" * (64 * 1024 * 1024), 1.0, watchdog + ) + finally: + hold["stop"] = True + srv.close() + watchdog.shutdown() + assert outcome2 == "raised" + assert isinstance(exc2, urllib3.exceptions.HTTPError) + assert any(h.state == "cancelled" for h in arming2.all) + assert elapsed2 < 5.0 # not hung + + +def _multi_stall_server() -> Tuple[socket.socket, int, Dict[str, object]]: + """Accept MULTIPLE connections and black-hole each one's body write — so a + urllib3-internal retry (which opens a fresh conn) stalls too. + """ + srv = socket.socket() + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(("127.0.0.1", 0)) + srv.listen(8) + port = srv.getsockname()[1] + hold: Dict[str, object] = {"conns": []} + + def serve() -> None: + while not hold.get("stop"): + try: + srv.settimeout(0.2) + conn, _ = srv.accept() + except socket.timeout: + continue + except OSError: + break + hold["conns"].append(conn) # accept, never read → send stall + + threading.Thread(target=serve, daemon=True).start() + return srv, port, hold + + +def test_internal_retry_attempt_is_also_watchdog_protected() -> None: + """Codex #1b: after the watchdog cancels attempt 1, urllib3 classifies the + resulting ProtocolError(RemoteDisconnected) as a pre-response reset and + INTERNALLY retries on a fresh conn. That retry attempt must ALSO be armed + (publish/arm happen per request()), so it can't run unbounded — both attempts + get cancelled and the call returns promptly. + """ + srv, port, hold = _multi_stall_server() + # total=1 → one internal retry; ConnectionResetRetry-style would retry the + # reset. Each attempt black-holes; each must be watchdog-armed+cancelled. + pool = urllib3.HTTPConnectionPool( + "127.0.0.1", port, maxsize=2, + retries=urllib3.Retry(total=1, redirect=False, raise_on_status=False), + ) + install_watchdog_connection_classes(pool) + try: + outcome, exc, elapsed, arming = _worker_put( + pool, "/stall", b"z" * (64 * 1024 * 1024), 1.0 + ) + finally: + hold["stop"] = True + srv.close() + assert outcome == "raised" + # Both the initial attempt AND the internal retry were armed (≥2 handles), + # and at least one was cancelled by the watchdog — none ran unbounded. + assert len(arming.all) >= 2 + assert any(h.state == "cancelled" for h in arming.all) + assert elapsed < 8.0 # bounded (≈ 2 × the 1s deadline), not hung + + +# --- Codex #9: watchdog-vs-finally interleaving --------------------------- + + +def test_handle_done_wins_when_worker_finishes_first() -> None: + """If the worker has already finished (response fully read) when the watchdog + fires, the watchdog must NOT shut down a socket that may have been reused. + """ + handle = _PartHandle() + handle.publish(object()) # a stand-in conn + handle.finish() # worker done first + # Watchdog now tries to cancel — must be a no-op. + conn = handle.cancel_if_active() + assert conn is None + assert handle.state == "done" + + +def test_handle_cancel_wins_when_still_in_flight() -> None: + sentinel = object() + handle = _PartHandle() + handle.publish(sentinel) + conn = handle.cancel_if_active() # watchdog fires while still active + assert conn is sentinel + assert handle.state == "cancelled" + # A late worker finish does not flip state back to done. + handle.finish() + assert handle.state == "cancelled" + + +def test_publish_after_cancel_is_ignored() -> None: + handle = _PartHandle() + handle.publish(object()) + handle.cancel_if_active() + handle.publish(object()) # a urllib3-internal retry tries to publish a new conn + # Already cancelled: no new conn is accepted. + assert handle.state == "cancelled" + + +# --- per-part handle isolation across threads ----------------------------- + + +def test_active_arming_is_thread_local() -> None: + """Each concurrent part runs in its own pool thread and must see its OWN + arming, never another part's (a global would cross the wires). + """ + from hotdata._upload_watchdog import _Arming, _get_active_arming_for_test + + seen: Dict[int, object] = {} + armings: Dict[int, _Arming] = {} + barrier = threading.Barrier(3) + wd = _Watchdog() + + def worker(i: int) -> None: + a = _Arming(wd, deadline=100.0) + armings[i] = a + set_active_arming(a) + barrier.wait() # force all three to be set simultaneously + seen[i] = _get_active_arming_for_test() + set_active_arming(None) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(3)] + for t in threads: + t.start() + for t in threads: + t.join() + for i in range(3): + assert seen[i] is armings[i] + + +# --- watchdog teardown (Codex #6) ----------------------------------------- + + +def test_watchdog_thread_joins_on_shutdown() -> None: + before = set(threading.enumerate()) + watchdog = _Watchdog() + handle = _PartHandle() + watchdog.arm(handle, deadline=100.0) # never fires in the test + watchdog.shutdown() # must stop + join the thread + after = set(threading.enumerate()) + leaked = [t for t in (after - before) if t.is_alive() and "watchdog" in t.name.lower()] + assert not leaked + + +# --- ConnectionCls injection (Codex #2: HTTP and HTTPS) ------------------- + + +def test_install_sets_both_http_and_https_connection_classes() -> None: + from hotdata._upload_watchdog import ( + WatchdogHTTPConnection, + WatchdogHTTPSConnection, + ) + + http_pool = urllib3.HTTPConnectionPool("127.0.0.1", 1) + install_watchdog_connection_classes(http_pool) + assert http_pool.ConnectionCls is WatchdogHTTPConnection + + https_pool = urllib3.HTTPSConnectionPool("127.0.0.1", 1) + install_watchdog_connection_classes(https_pool) + assert https_pool.ConnectionCls is WatchdogHTTPSConnection diff --git a/tests/test_uploads.py b/tests/test_uploads.py index cd95253..b40441a 100644 --- a/tests/test_uploads.py +++ b/tests/test_uploads.py @@ -861,7 +861,10 @@ def test_multipart_empty_part_urls_raises( ) api = _make_api() _patch_session(monkeypatch, api, session, []) - with pytest.raises(MalformedSessionError, match="part_urls"): + # A present-but-empty part_urls is a malformed known-size session (its URL + # count won't match how the file slices) — NOT a streaming session (which is + # signalled by part_urls being absent/None). It surfaces as a count mismatch. + with pytest.raises(MalformedSessionError, match="0 part URLs"): api.upload_file(str(src)) From 5dba9fcfcc71100b333a9b831ae6d50f4808b0c3 Mon Sep 17 00:00:00 2001 From: Zac Farrell Date: Tue, 30 Jun 2026 20:45:27 -0700 Subject: [PATCH 2/3] test(auth): cover token-exchange body-read retry --- tests/test_auth.py | 106 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 106 insertions(+) diff --git a/tests/test_auth.py b/tests/test_auth.py index c77f43d..6f99557 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -807,3 +807,109 @@ def counting() -> str: assert _bearer_from(auth) == "Bearer eyJ.once.jwt" assert count["n"] == 1 + + +# -------------------------------------------------------------------------- +# Item 9 -- token-exchange success-path body-read retry +# +# A connection dropped while reading the body of a 200 is a transport-level, +# transient failure and must be retried within the attempt budget, exactly like +# a pre-response connection error. A valid-but-malformed JSON body, in contrast, +# is a definitive rejection and must NOT be retried. The truncated-body case +# genuinely requires a raw socket: no fake pool can send a valid 200 status line +# and then drop mid-body, which is the precise condition under test. +# -------------------------------------------------------------------------- + + +def _raw_token_server(bodies: List[Optional[bytes]]): + """Start a raw TCP server that serves one scripted response per connection. + + Each element of ``bodies`` is the raw bytes to send for that connection + (already including status line + headers), or ``None`` to send a truncated + 200 -- a Content-Length promising more than is sent, then an immediate close. + Returns ``(port, accepted)`` where ``accepted["n"]`` counts connections. + """ + import socket + + srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(("127.0.0.1", 0)) + srv.listen(len(bodies) + 1) + port = srv.getsockname()[1] + accepted = {"n": 0} + + def serve() -> None: + for raw in bodies: + try: + conn, _ = srv.accept() + except OSError: + break + accepted["n"] += 1 + try: + conn.recv(65536) # drain the request + if raw is None: + # 200 promising 100 body bytes, then only a few + close. + conn.sendall( + b"HTTP/1.1 200 OK\r\nContent-Length: 100\r\n\r\n{\"access" + ) + else: + conn.sendall(raw) + except OSError: + pass + finally: + try: + conn.close() + except OSError: + pass + try: + srv.close() + except OSError: + pass + + threading.Thread(target=serve, daemon=True).start() + return port, accepted + + +def _full_token_response(access_token: str = "eyJ.real.jwt") -> bytes: + payload = json.dumps( + {"access_token": access_token, "token_type": "Bearer", "expires_in": 300} + ).encode() + return ( + b"HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n%b" % (len(payload), payload) + ) + + +def test_truncated_200_body_is_retried_then_succeeds() -> None: + """A 200 whose body is truncated mid-read (connection dropped) is transient: + the exchange retries on a fresh connection and returns the token. + + Uses a real socket + real PoolManager -- the whole point is a genuine + valid-status-line-then-truncated-body transport failure, which surfaces as a + urllib3 ``ProtocolError`` from ``pool.request()`` (the body is preloaded) and + is caught + retried by ``_exchange``'s transport-error branch. + """ + sleeps: List[float] = [] + port, accepted = _raw_token_server([None, _full_token_response("eyJ.real.jwt")]) + mgr = _TokenManager( + "hd_secret_token", + _config(host=f"http://127.0.0.1:{port}"), + pool=urllib3.PoolManager(), + sleep=sleeps.append, + ) + + assert mgr.bearer_value() == "eyJ.real.jwt" + assert accepted["n"] == 2 # first (truncated) retried on a second connection + assert len(sleeps) == 1 + + +def test_malformed_json_200_body_is_not_retried() -> None: + """A complete 200 whose body is valid bytes but invalid JSON is a definitive + rejection -- it must NOT be retried, and must surface as TokenExchangeError.""" + sleeps: List[float] = [] + pool = _FakePool([_FakeResponse(200, b"{ not valid json ")]) + mgr = _TokenManager("hd_secret_token", _config(), pool=pool, sleep=sleeps.append) + + with pytest.raises(TokenExchangeError): + mgr.bearer_value() + assert len(pool.calls) == 1 # parse error is fatal, not retried + assert sleeps == [] From fcaea734c45c29eab903e0e26a2e5d10df52fa90 Mon Sep 17 00:00:00 2001 From: Zac Farrell Date: Wed, 1 Jul 2026 07:39:54 -0700 Subject: [PATCH 3/3] feat(uploads): rust parity gaps and upload ergonomics --- README.md | 46 ++- hotdata/__init__.py | 35 ++ hotdata/uploads.py | 314 +++++++++++++++--- scripts/patch_query_exports.py | 35 ++ tests/test_upload_parity_fixes.py | 512 ++++++++++++++++++++++++++++++ tests/test_uploads.py | 33 +- 6 files changed, 915 insertions(+), 60 deletions(-) create mode 100644 tests/test_upload_parity_fixes.py diff --git a/README.md b/README.md index 91abc4d..9759f43 100644 --- a/README.md +++ b/README.md @@ -117,18 +117,40 @@ with ApiClient(Configuration(api_key="...", workspace_id="...")) as client: (`size` is inferred for all three; a file object is read from its current position to the end). The SDK picks single vs. multipart from the size, auto-scales the part size, and bounds part concurrency to a peak-memory budget -(override with `part_size` / `max_concurrency` / `part_retry`). Storage `PUT`s go -through a dedicated, header-isolated connection pool, so the SDK's auth and -workspace headers never reach object storage (which would otherwise reject the -upload). Finalize is sent with retries disabled so the exactly-once call is never -accidentally replayed. - -Failures surface as a typed hierarchy under `hotdata.uploads.UploadError`: -`StorageError` (storage returned a non-2xx), `StorageTransportError` (the PUT -failed before any response), `MissingETagError`, `MalformedSessionError`, and -`SizeLimitError`. Opening the session or finalizing raises the usual -`hotdata.exceptions.ApiException` — for example a `501` `PRESIGN_UNSUPPORTED`, -meaning the backend cannot issue upload URLs. +(override with `part_size` / `max_concurrency` / `part_retry`). A file larger than +8 MiB uploads via a **streaming** session — the SDK mints each part's URL just +before its `PUT`, so a presigned URL can't expire mid-transfer on a slow upload. +Storage `PUT`s go through a dedicated, header-isolated connection pool (auth and +workspace headers never reach object storage, which would otherwise reject the +upload) with a 30s connect timeout so a dead endpoint fails fast. Finalize is sent +with retries disabled so the exactly-once call is never accidentally replayed. + +Every failure is a subclass of `hotdata.uploads.UploadError` (also importable as +`from hotdata import UploadError`), so a single `except UploadError` catches the +whole flow: `SessionCreateError` (opening the session — check `.status` for a +`501` `PRESIGN_UNSUPPORTED`), `StorageError` (storage returned a non-2xx; `.exhausted` +is `True` if it outlived every retry round), `StorageTransportError` (the PUT +failed before any response), `MissingETagError`, `MintPartError` (minting a part +URL), `FinalizeError`, `MalformedSessionError`, `SizeLimitError`, and +`UploadCancelledError`. The phase errors that wrap a control-plane call chain the +underlying `hotdata.exceptions.ApiException` as `__cause__` (and expose it as +`.api_exception` / `.status`). A local file read error surfaces as `OSError`. + +The `progress` callback receives a **cumulative** `(bytes_done, total)` — for a +tqdm bar (whose `update(n)` wants a delta, and which isn't thread-safe under +multipart) use the ready-made adapter: + +```python +from hotdata import tqdm_progress +from tqdm import tqdm + +with tqdm(total=size, unit="B", unit_scale=True) as bar: + uploads.upload_file("data.parquet", progress=tqdm_progress(bar)) +``` + +Pass a `threading.Event` as `cancel_event` to abort an in-flight upload; tune the +control-plane calls with `request_timeout` (storage-PUT timeouts are automatic and +size-scaled). For that fallback (or to upload from a non-seekable stream), use `upload_stream`, which sends the bytes to the legacy `POST /v1/files` endpoint in one request, diff --git a/hotdata/__init__.py b/hotdata/__init__.py index 0786849..9997044 100644 --- a/hotdata/__init__.py +++ b/hotdata/__init__.py @@ -334,3 +334,38 @@ from hotdata.query import QueryApi as QueryApi # noqa: E402,F811 from hotdata.arrow import ResultsApi as ResultsApi # noqa: E402,F811 from hotdata.uploads import UploadsApi as UploadsApi # noqa: E402,F811 +# The upload error hierarchy + option/helper types, so `from hotdata import +# UploadError` / `StorageError` / ... works (they otherwise live only under +# hotdata.uploads and are invisible to autocomplete on `hotdata.`). +from hotdata.uploads import ( # noqa: E402,F401 + UploadError as UploadError, + SessionCreateError as SessionCreateError, + MalformedSessionError as MalformedSessionError, + SizeLimitError as SizeLimitError, + StorageError as StorageError, + StorageTransportError as StorageTransportError, + MissingETagError as MissingETagError, + MintPartError as MintPartError, + FinalizeError as FinalizeError, + UploadCancelledError as UploadCancelledError, + UploadOptions as UploadOptions, + UploadProgress as UploadProgress, + UploadSource as UploadSource, + tqdm_progress as tqdm_progress, +) +__all__ += [ # noqa: F405 + "UploadError", + "SessionCreateError", + "MalformedSessionError", + "SizeLimitError", + "StorageError", + "StorageTransportError", + "MissingETagError", + "MintPartError", + "FinalizeError", + "UploadCancelledError", + "UploadOptions", + "UploadProgress", + "UploadSource", + "tqdm_progress", +] diff --git a/hotdata/uploads.py b/hotdata/uploads.py index 4b50c4b..f05565e 100644 --- a/hotdata/uploads.py +++ b/hotdata/uploads.py @@ -101,6 +101,16 @@ #: actual size. See :func:`auto_part_size_hint`. DEFAULT_PART_SIZE = 8 * MIB +#: A file larger than this is uploaded via a *streaming* (just-in-time mint) +#: session: :meth:`UploadsApi.upload_file` OMITS ``declared_size_bytes`` so the +#: server mints NO part URLs up front and the client mints each part's URL moments +#: before its ``PUT``. A presigned URL then cannot expire mid-transfer no matter +#: how long a slow large upload runs — the failure mode of the eager known-size +#: path, whose URLs share a fixed (~30-minute) TTL. A file at or below this stays +#: on the known-size path (single quick ``PUT`` / eager fast path) by declaring its +#: size. Mirrors the Rust SDK's ``STREAMING_THRESHOLD`` (also 8 MiB). +STREAMING_THRESHOLD = DEFAULT_PART_SIZE + #: Target ceiling on part count when auto-scaling the part-size hint for very #: large files, with headroom under S3's hard 10,000-part limit. TARGET_MAX_PARTS = 9000 @@ -213,22 +223,47 @@ class UploadOptions: #: rounds. ``None`` uses :data:`DEFAULT_ROUND_BASE_DELAY` (2s). Set to ``0`` #: in tests for instant rounds. round_base_delay: Optional[float] = None + #: Timeout for the control-plane API calls (create session / finalize / mint): + #: a number of seconds or a ``(connect, read)`` tuple. Storage ``PUT`` timeouts + #: are automatic and size-scaled (see :func:`_part_put_timeout`) and are not + #: affected by this. + request_timeout: Any = None + #: Optional :class:`threading.Event`; setting it aborts an in-flight upload + #: (in-flight part ``PUT`` s are cancelled and the upload is not finalized), + #: raising :class:`UploadCancelledError`. The upload is not resumable. + cancel_event: Optional[threading.Event] = None # --- Errors --------------------------------------------------------------- class UploadError(Exception): - """Base class for errors raised by :meth:`UploadsApi.upload_file` while - driving the direct-to-storage flow. - - Opening the session or finalizing fails with the generated client's - :class:`~hotdata.exceptions.ApiException` instead (e.g. a ``501`` - ``PRESIGN_UNSUPPORTED`` when the storage backend cannot issue upload URLs); - only the storage transfer and session-consistency failures raise - ``UploadError`` subclasses. + """Base class for **every** error raised by :meth:`UploadsApi.upload_file`. + + A single ``except UploadError`` catches the whole flow: opening the session + (:class:`SessionCreateError`), the storage transfer + (:class:`StorageError` / :class:`StorageTransportError` / + :class:`MissingETagError`), minting part URLs (:class:`MintPartError`), + finalizing (:class:`FinalizeError`), a session-consistency problem + (:class:`MalformedSessionError`), a size overflow (:class:`SizeLimitError`), + or caller cancellation (:class:`UploadCancelledError`). The phase errors that + wrap a control-plane call chain the underlying + :class:`~hotdata.exceptions.ApiException` as ``__cause__`` (and expose it as + ``api_exception`` / its HTTP ``status``), so a specific status such as a + ``501`` ``PRESIGN_UNSUPPORTED`` stays detectable. + + ``OSError`` (local file could not be read) is the one failure that is NOT an + ``UploadError`` — it is a plain filesystem error, not an upload-flow error. """ + #: ``True`` when a *retryable* failure survived every whole-upload re-sweep + #: round (retries exhausted), as opposed to a one-shot terminal reject. Lets a + #: caller distinguish "give up and alert" from "worth scheduling a later try". + exhausted: bool = False + #: Number of re-sweep rounds attempted before this error was surfaced + #: (``None`` when not raised from the re-sweep driver). + rounds_attempted: Optional[int] = None + class MalformedSessionError(UploadError): """The create-session response was internally inconsistent for its declared @@ -303,6 +338,57 @@ def __init__(self, *, part_number: int) -> None: ) +class _ApiPhaseError(UploadError): + """Base for the upload phases that wrap a control-plane + :class:`~hotdata.exceptions.ApiException`. Exposes the wrapped exception as + ``api_exception`` and its HTTP ``status`` while chaining it as ``__cause__``, + so ``except UploadError`` catches the whole flow yet a specific status stays + detectable. + """ + + _phase = "the upload" + + def __init__(self, api_exception: ApiException) -> None: + self.api_exception = api_exception + self.status = getattr(api_exception, "status", None) + super().__init__(f"{self._phase} failed: {api_exception}") + + +class SessionCreateError(_ApiPhaseError): + """Opening the upload session (``POST /v1/uploads``) failed. A ``501`` + ``PRESIGN_UNSUPPORTED`` (``status == 501``) means the backend cannot issue + presigned URLs — fall back to :meth:`UploadsApi.upload_stream`. + """ + + _phase = "opening the upload session" + + +class MintPartError(_ApiPhaseError): + """Minting a part URL (``POST /v1/uploads/{id}/parts``) failed. Retryable — + the whole-upload re-sweep re-mints on a later round. + """ + + _phase = "minting a part URL" + + +class FinalizeError(_ApiPhaseError): + """Finalizing the upload (``POST /v1/uploads/{id}/finalize``) failed after the + bytes were stored. Not retried automatically (finalize is exactly-once). + """ + + _phase = "finalizing the upload" + + +class UploadCancelledError(UploadError): + """The caller's ``cancel_event`` was set, so the upload was aborted. In-flight + part ``PUT`` s are cancelled and the upload is not finalized. The upload is not + resumable — a later :meth:`UploadsApi.upload_file` opens a fresh session. + """ + + def __init__(self) -> None: + super().__init__("upload cancelled by caller") + + # --- Error classification (Item 2) ---------------------------------------- @@ -319,11 +405,13 @@ def _is_retryable(exc: BaseException) -> bool: write), or anything unexpected. Retrying those would burn the round budget for nothing or, worse, double-``PUT``. """ + if isinstance(exc, UploadCancelledError): + return False # caller asked to stop; never retry if isinstance(exc, (StorageTransportError, MissingETagError)): return True if isinstance(exc, StorageError): return True # incl. 5xx (outer tier owns it) and 403 (eager URL expiry) - if isinstance(exc, ApiException): + if isinstance(exc, (MintPartError, ApiException)): return True # a per-part mint round-trip failure — a re-mint may cure it return False @@ -395,6 +483,7 @@ def _upload_parts_resilient( round_base_delay: float = DEFAULT_ROUND_BASE_DELAY, on_part_done: Optional[Callable[[_PartPlan], None]] = None, on_terminal: Optional[Callable[[], None]] = None, + cancel_event: Optional[threading.Event] = None, sleep: Callable[[float], None] = time.sleep, ) -> List[FinalizeUploadPart]: """Run every part's ``upload_one`` work unit with whole-upload re-sweep. @@ -448,6 +537,11 @@ def run_one(plan: _PartPlan) -> None: # A terminal error elsewhere in the round: don't dispatch new work. failures[plan.part_number] = _NotDispatched() return + if cancel_event is not None and cancel_event.is_set(): + # Caller asked to stop: fail fast (terminal), cancelling in-flight + # siblings, so the upload aborts without finalizing. + _fail_terminal(UploadCancelledError()) + return try: part = upload_one(plan) except BaseException as exc: # noqa: BLE001 - classified below @@ -517,10 +611,19 @@ def run_one(plan: _PartPlan) -> None: if pending: # Exhausted the round budget with parts still failing: surface the last - # underlying error so the caller's normal mapping applies. + # underlying error so the caller's normal mapping applies, tagging it as + # retries-exhausted (Ergo #6) so a caller can tell "gave up after N rounds" + # from a one-shot terminal reject. + rounds = 1 + max_extra_rounds if last_error is not None: + if isinstance(last_error, UploadError): + last_error.exhausted = True + last_error.rounds_attempted = rounds raise last_error - raise StorageTransportError(url="", part_number=pending[0].part_number) + exc = StorageTransportError(url="", part_number=pending[0].part_number) + exc.exhausted = True + exc.rounds_attempted = rounds + raise exc # Completeness: every plan's slot must be filled. missing = [p.part_number for p in plans if p.part_number not in results] @@ -554,10 +657,18 @@ def _wait_first(futures: set) -> Tuple[set, set]: # --- Storage transport ---------------------------------------------------- +#: Bound on TCP+TLS establishment for a storage ``PUT`` (Item 3). Only the +#: *connect* phase is bounded — a dead/black-holed endpoint fails fast into the +#: whole-upload re-sweep instead of blocking on the OS SYN timeout (~75s). The +#: *read* phase is deliberately left unbounded (``read=None``): a large upload +#: legitimately takes minutes, and the per-part watchdog (not a read timeout) owns +#: stall detection once bytes start flowing. Mirrors the Rust SDK's 30s connect +#: timeout. +STORAGE_CONNECT_TIMEOUT = 30.0 + # A dedicated, process-wide, *bare* pool for storage PUTs. Deliberately NOT the # SDK's ApiClient pool: a host app may have installed auth / workspace headers -# there, which storage would reject. Built lazily and reused. No request -# timeout: a large upload legitimately takes minutes. +# there, which storage would reject. Built lazily and reused. _STORAGE_POOL: Optional[urllib3.PoolManager] = None _STORAGE_POOL_LOCK = threading.Lock() @@ -567,7 +678,14 @@ def _storage_pool() -> urllib3.PoolManager: if _STORAGE_POOL is None: with _STORAGE_POOL_LOCK: if _STORAGE_POOL is None: - pool = urllib3.PoolManager() + # Bound connect only; leave read unbounded (see + # STORAGE_CONNECT_TIMEOUT). The pool-level default flows to every + # PUT unless a per-request timeout overrides it. + pool = urllib3.PoolManager( + timeout=urllib3.Timeout( + connect=STORAGE_CONNECT_TIMEOUT, read=None + ) + ) # Install the stall-cancellation connection subclass at creation # time — BEFORE any request can populate `manager.pools` with a # host pool carrying the stock connection class (Codex #2). The @@ -815,6 +933,8 @@ def upload_file( # type: ignore[override] part_retry: Optional[Retry] = None, max_extra_rounds: Optional[int] = None, round_base_delay: Optional[float] = None, + request_timeout: Any = None, + cancel_event: Optional[threading.Event] = None, options: Optional[UploadOptions] = None, _request_timeout: Any = None, ) -> FinalizeUploadResponse: @@ -849,21 +969,41 @@ def upload_file( # type: ignore[override] :param max_concurrency: Max in-flight part ``PUT`` s for a multipart upload (default :data:`DEFAULT_MAX_CONCURRENCY`, further bounded by a peak-memory budget). - :param progress: Callback invoked with ``(bytes_done_total, total)``. + :param progress: Callback invoked with ``(bytes_done_total, total)`` — a + *cumulative* byte count (not a delta). For multipart it fires from + worker threads, so a callback touching shared state must lock; see + :func:`tqdm_progress` for a ready-made thread-safe tqdm adapter. :param part_retry: ``urllib3`` retry policy for multipart part ``PUT`` s; defaults to :func:`default_part_retry`. + :param request_timeout: Timeout for the control-plane API calls (create + session / finalize / mint): a number of seconds or ``(connect, read)``. + Storage ``PUT`` timeouts are automatic and size-scaled and are not + affected by this. + :param cancel_event: Optional :class:`threading.Event`; setting it aborts + the upload (in-flight part ``PUT`` s are cancelled and the upload is not + finalized), raising :class:`UploadCancelledError`. Not resumable. :param options: An :class:`UploadOptions` bundle; individual keyword arguments take precedence over its fields. - :raises SizeLimitError: the file is larger than the wire's 64-bit field. + + Every failure is a :class:`UploadError` subclass, so one ``except + UploadError`` catches the whole flow (the sole exception is ``OSError`` for + a local read failure): + + :raises SessionCreateError: opening the session failed (e.g. a ``501`` + ``PRESIGN_UNSUPPORTED`` — check ``.status``; fall back to + :meth:`upload_stream`). Wraps the ``ApiException`` as ``__cause__``. + :raises SizeLimitError: the part-size hint exceeds the wire's 64-bit field. :raises MalformedSessionError: the session response was inconsistent. - :raises StorageError: a storage ``PUT`` returned a non-2xx status. + :raises StorageError: a storage ``PUT`` returned a non-2xx status. Its + ``exhausted`` flag is ``True`` if it survived every re-sweep round. :raises StorageTransportError: a storage ``PUT`` failed before any response (connection/TLS/DNS error, or retries exhausted). :raises MissingETagError: a part ``PUT`` returned no ``ETag``. + :raises MintPartError: minting a part URL failed. Wraps the ``ApiException``. + :raises FinalizeError: finalizing failed after the bytes were stored. + Wraps the ``ApiException``. + :raises UploadCancelledError: ``cancel_event`` was set. :raises OSError: the local file could not be opened or read. - :raises hotdata.exceptions.ApiException: opening the session or - finalizing failed (e.g. ``501`` ``PRESIGN_UNSUPPORTED`` — the storage - backend cannot issue upload URLs; use ``POST /v1/files`` instead). """ opts = options or UploadOptions() content_type = content_type if content_type is not None else opts.content_type @@ -883,6 +1023,20 @@ def upload_file( # type: ignore[override] round_base_delay = ( round_base_delay if round_base_delay is not None else opts.round_base_delay ) + request_timeout = ( + request_timeout if request_timeout is not None else opts.request_timeout + ) + cancel_event = cancel_event if cancel_event is not None else opts.cancel_event + + # `request_timeout` is the public control-plane timeout knob; `_request_timeout` + # is the historical (leading-underscore) alias. The public name wins when both + # are given. It governs the create-session / finalize / mint API calls only — + # storage PUT timeouts are automatic and size-scaled (see `_part_put_timeout`). + timeout = request_timeout if request_timeout is not None else _request_timeout + + # Fail fast if the caller already cancelled before any work started. + if cancel_event is not None and cancel_event.is_set(): + raise UploadCancelledError() src = _make_source(source, size) total = src.size @@ -893,22 +1047,36 @@ def upload_file( # type: ignore[override] # declared size. The server clamps it regardless. part_size_hint = part_size if part_size is not None else auto_part_size_hint(total) - # The wire models sizes as a signed i64; reject a pathological size up - # front rather than silently sending an out-of-range value. - if total > _MAX_WIRE_INT: - raise SizeLimitError(what="declared_size_bytes", value=total) + # Choose the session shape by size (Item 4). A large file OMITS + # `declared_size_bytes` so the server opens a STREAMING session (no part + # URLs up front); the client then mints each part's URL just before its + # PUT, so a presigned URL cannot expire mid-transfer. A small file DECLARES + # its size to keep the known-size path (single-PUT / eager fast path). The + # i64 wire-range guard applies only when the size is actually sent — a + # streamed file's size never reaches the wire, so it is never an obstacle + # (matches the Rust SDK). The part-size hint is always sent, so it is + # always range-checked. if part_size_hint > _MAX_WIRE_INT: raise SizeLimitError(what="part_size", value=part_size_hint) - session = self.create_upload_session_handler( - CreateUploadRequest( - content_type=content_type, - content_encoding=content_encoding, - filename=filename, - declared_size_bytes=total, - part_size=part_size_hint, - ), - _request_timeout=_request_timeout, + create_kwargs: Dict[str, Any] = dict( + content_type=content_type, + content_encoding=content_encoding, + filename=filename, + part_size=part_size_hint, + ) + if total <= STREAMING_THRESHOLD: + if total > _MAX_WIRE_INT: # unreachable given the threshold; explicit + raise SizeLimitError(what="declared_size_bytes", value=total) + create_kwargs["declared_size_bytes"] = total + # else: OMIT declared_size_bytes entirely (do NOT pass None). Passing None + # sets the field in pydantic's `model_fields_set`, serializing to + # `"declared_size_bytes": null`; the server needs the field ABSENT to open + # a streaming session, so the field is left off the constructor call. + + session = self._create_session( + CreateUploadRequest(**create_kwargs), + _request_timeout=timeout, ) # Report initial progress so a 0-byte file (or an instant single PUT) @@ -917,7 +1085,7 @@ def upload_file( # type: ignore[override] progress(0, total) if session.mode == "single": - self._upload_single(session, src, total, progress) + self._upload_single(session, src, total, progress, cancel_event) parts: Optional[List[FinalizeUploadPart]] = None elif session.mode == "multipart": cap = max_concurrency if max_concurrency is not None else DEFAULT_MAX_CONCURRENCY @@ -937,6 +1105,7 @@ def upload_file( # type: ignore[override] round_base_delay if round_base_delay is not None else DEFAULT_ROUND_BASE_DELAY ), + cancel_event=cancel_event, ) else: raise MalformedSessionError(f"unknown upload mode {session.mode!r}") @@ -958,7 +1127,7 @@ def upload_file( # type: ignore[override] session.upload_id, session.finalize_token, finalize_body, - _request_timeout, + timeout, ) def upload_stream( @@ -1062,6 +1231,21 @@ def upload_stream( # -- internals --------------------------------------------------------- + def _create_session( + self, request: CreateUploadRequest, *, _request_timeout: Any + ) -> UploadSessionResponse: + """Open the upload session, mapping a control-plane failure to + :class:`SessionCreateError` so the whole flow catches under + :class:`UploadError` (Ergo #2). The ``ApiException`` is chained as + ``__cause__`` and exposed via ``.api_exception`` / ``.status``. + """ + try: + return self.create_upload_session_handler( + request, _request_timeout=_request_timeout + ) + except ApiException as exc: + raise SessionCreateError(exc) from exc + def _finalize_handler( self, upload_id: str, @@ -1089,6 +1273,10 @@ def _finalize_handler( body, _request_timeout=_request_timeout, ) + except ApiException as exc: + # Bytes are already stored; only the finalize confirmation failed. Map + # to FinalizeError so `except UploadError` catches it (Ergo #2). + raise FinalizeError(exc) from exc finally: finalize_client.close() @@ -1098,6 +1286,7 @@ def _upload_single( src: _Source, total: int, progress: Optional[UploadProgress], + cancel_event: Optional[threading.Event] = None, ) -> None: """Single-``PUT`` path: stream the whole source to ``session.url``, invoking the progress callback incrementally as chunks are sent. @@ -1110,6 +1299,9 @@ def _upload_single( if not url: raise MalformedSessionError("single upload missing `url`") + if cancel_event is not None and cancel_event.is_set(): + raise UploadCancelledError() + with src.reader() as fileobj: body = _ProgressReader(fileobj, total, progress) _put_to_storage( @@ -1145,6 +1337,7 @@ def _upload_multipart( *, max_extra_rounds: int = DEFAULT_MAX_EXTRA_ROUNDS, round_base_delay: float = DEFAULT_ROUND_BASE_DELAY, + cancel_event: Optional[threading.Event] = None, ) -> List[FinalizeUploadPart]: """Resilient multipart path: ``PUT`` each ``part_size``-byte chunk to storage with the whole-upload re-sweep (Item 1) — a single part's @@ -1262,6 +1455,7 @@ def on_part_done(plan: _PartPlan) -> None: round_base_delay=round_base_delay, on_part_done=on_part_done, on_terminal=watchdog.cancel_all, + cancel_event=cancel_event, ) finally: watchdog.shutdown() @@ -1281,12 +1475,18 @@ def _mint_part(self, session: UploadSessionResponse, part_number: int) -> str: eager-path 403 re-mint). Validates the response is for the requested part number (Item 6); carries a finite (inactivity) timeout (Item 4 / Codex #5). """ - resp = self.mint_upload_parts_handler( - session.upload_id, - session.finalize_token, - MintUploadPartsRequest(part_numbers=[part_number]), - _request_timeout=MINT_TIMEOUT, - ) + try: + resp = self.mint_upload_parts_handler( + session.upload_id, + session.finalize_token, + MintUploadPartsRequest(part_numbers=[part_number]), + _request_timeout=MINT_TIMEOUT, + ) + except ApiException as exc: + # A mint round-trip failure is retryable (the re-sweep re-mints on a + # later round); wrap it so `except UploadError` catches it (Ergo #2) + # while `_is_retryable` still classifies it as retryable. + raise MintPartError(exc) from exc # Strict validation (Item 6): we requested exactly one part number, so the # response must contain exactly that one. Reject a DUPLICATE (two entries # for the part — ambiguous, and a plain dict would silently collapse it), @@ -1404,6 +1604,31 @@ def reader(self) -> Iterator[BinaryIO]: yield self._file +def tqdm_progress(bar: Any, *, lock: Optional[threading.Lock] = None) -> UploadProgress: + """Adapt :meth:`UploadsApi.upload_file`'s progress callback to a ``tqdm``-style + progress bar. + + ``upload_file`` reports the **cumulative** bytes done, but a tqdm bar's + ``update(n)`` wants a **delta** — and multipart fires the callback from worker + threads, so the update must be serialized. This wraps both concerns, so the + obvious usage is correct and thread-safe:: + + from tqdm import tqdm + with tqdm(total=size, unit="B", unit_scale=True) as bar: + uploads.upload_file(path, progress=tqdm_progress(bar)) + + Works with any object exposing ``.n`` and ``update(n)`` (tqdm, or a compatible + shim). Pass your own ``lock`` to share one bar with other writers. + """ + guard = lock if lock is not None else threading.Lock() + + def _cb(done: int, total: int) -> None: + with guard: + bar.update(done - bar.n) + + return _cb + + def _make_source(source: UploadSource, size: Optional[int]) -> _Source: """Normalize an upload input into a :class:`_Source`. @@ -1440,13 +1665,19 @@ def _make_source(source: UploadSource, size: Optional[int]) -> _Source: "MAX_PART_SIZE", "MIB", "MIN_PART_SIZE", + "STORAGE_CONNECT_TIMEOUT", + "STREAMING_THRESHOLD", "TARGET_MAX_PARTS", "UPLOAD_MEMORY_BUDGET", + "FinalizeError", "MalformedSessionError", + "MintPartError", "MissingETagError", + "SessionCreateError", "SizeLimitError", "StorageError", "StorageTransportError", + "UploadCancelledError", "UploadError", "UploadOptions", "UploadProgress", @@ -1455,4 +1686,5 @@ def _make_source(source: UploadSource, size: Optional[int]) -> _Source: "auto_part_size_hint", "default_part_retry", "effective_in_flight", + "tqdm_progress", ] diff --git a/scripts/patch_query_exports.py b/scripts/patch_query_exports.py index 1335e13..0775fa4 100644 --- a/scripts/patch_query_exports.py +++ b/scripts/patch_query_exports.py @@ -35,6 +35,41 @@ "from hotdata.query import QueryApi as QueryApi # noqa: E402,F811\n" "from hotdata.arrow import ResultsApi as ResultsApi # noqa: E402,F811\n" "from hotdata.uploads import UploadsApi as UploadsApi # noqa: E402,F811\n" + "# The upload error hierarchy + option/helper types, so `from hotdata import\n" + "# UploadError` / `StorageError` / ... works (they otherwise live only under\n" + "# hotdata.uploads and are invisible to autocomplete on `hotdata.`).\n" + "from hotdata.uploads import ( # noqa: E402,F401\n" + " UploadError as UploadError,\n" + " SessionCreateError as SessionCreateError,\n" + " MalformedSessionError as MalformedSessionError,\n" + " SizeLimitError as SizeLimitError,\n" + " StorageError as StorageError,\n" + " StorageTransportError as StorageTransportError,\n" + " MissingETagError as MissingETagError,\n" + " MintPartError as MintPartError,\n" + " FinalizeError as FinalizeError,\n" + " UploadCancelledError as UploadCancelledError,\n" + " UploadOptions as UploadOptions,\n" + " UploadProgress as UploadProgress,\n" + " UploadSource as UploadSource,\n" + " tqdm_progress as tqdm_progress,\n" + ")\n" + "__all__ += [ # noqa: F405\n" + ' "UploadError",\n' + ' "SessionCreateError",\n' + ' "MalformedSessionError",\n' + ' "SizeLimitError",\n' + ' "StorageError",\n' + ' "StorageTransportError",\n' + ' "MissingETagError",\n' + ' "MintPartError",\n' + ' "FinalizeError",\n' + ' "UploadCancelledError",\n' + ' "UploadOptions",\n' + ' "UploadProgress",\n' + ' "UploadSource",\n' + ' "tqdm_progress",\n' + "]\n" ) diff --git a/tests/test_upload_parity_fixes.py b/tests/test_upload_parity_fixes.py new file mode 100644 index 0000000..cacea8e --- /dev/null +++ b/tests/test_upload_parity_fixes.py @@ -0,0 +1,512 @@ +"""Parity/ergonomics fixes surfaced by the sdk-rust review. + +Covers the two functional gaps (streaming-path selection, storage connect +timeout) and the ergonomic changes (top-level re-exports, unified UploadError +base, progress adapter, cancellation, exhausted marker, public timeout knob). +The seams are stubbed the same way as ``test_upload_multipart_resilient.py``: +create/finalize/mint via monkeypatch, storage PUTs via a sequencing fake pool. +""" + +from __future__ import annotations + +import threading +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional + +import pytest + +import hotdata.uploads as uploads_mod +from hotdata.uploads import ( + MIB, + STREAMING_THRESHOLD, + MalformedSessionError, + StorageError, + UploadsApi, +) +from hotdata.exceptions import ApiException +from hotdata.models.finalize_upload_response import FinalizeUploadResponse +from hotdata.models.minted_upload_part_response import MintedUploadPartResponse +from hotdata.models.mint_upload_parts_response import MintUploadPartsResponse +from hotdata.models.upload_session_response import UploadSessionResponse + + +# --- sequencing fake storage pool (mirrors test_upload_multipart_resilient) -- + + +class _Resp: + def __init__(self, status: int, *, etag: Optional[str] = None, body: bytes = b""): + self.status = status + self._headers = {} + if etag is not None: + self._headers["ETag"] = etag + self.data = body + + @property + def headers(self): + class _H: + def __init__(self, d): + self._d = {k.lower(): v for k, v in d.items()} + + def get(self, k, default=None): + return self._d.get(k.lower(), default) + + return _H(self._headers) + + +class _SeqPool: + def __init__(self) -> None: + self.responses: Dict[str, List[_Resp]] = {} + self.default = _Resp(200, etag='"default"') + self.calls: List[Dict[str, Any]] = [] + self._lock = threading.Lock() + + def request(self, method: str, url: str, *, body=None, headers=None, retries=None, **kw): + # The single-PUT path passes a file-like body (a _ProgressReader); drain it + # like real urllib3 would so its byte counter advances and the short-read + # guard is satisfied. Record the drained bytes so body assertions still work. + recorded = body + if hasattr(body, "read"): + recorded = body.read() + with self._lock: + self.calls.append({"method": method, "url": url, "body": recorded}) + seq = self.responses.get(url) + if not seq: + return self.default + return seq.pop(0) if len(seq) > 1 else seq[0] + + +@pytest.fixture +def seq_pool(monkeypatch): + pool = _SeqPool() + monkeypatch.setattr(uploads_mod, "_STORAGE_POOL", pool) + monkeypatch.setattr(uploads_mod, "install_watchdog_connection_classes", lambda p: None) + return pool + + +def _make_api() -> UploadsApi: + from hotdata import ApiClient, Configuration + + cfg = Configuration(host="https://api.hotdata.test", api_key="k", workspace_id="ws") + return UploadsApi(ApiClient(cfg)) + + +def _capture_create(monkeypatch, api, session, captured): + def fake_create(create_upload_request, **kw): + captured.append(create_upload_request) + return session + + monkeypatch.setattr(api, "create_upload_session_handler", fake_create) + + +def _patch_finalize(monkeypatch, api, capture): + def fake_finalize(upload_id, finalize_token, body=None, _request_timeout=None): + capture.append(("finalize", upload_id, body)) + return FinalizeUploadResponse( + content_type=None, + created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + size_bytes=0, + status="ready", + upload_id=upload_id, + ) + + monkeypatch.setattr(api, "_finalize_handler", fake_finalize) + + +def _streaming_session(part_size): + return UploadSessionResponse( + finalize_token="tok", + headers={}, + mode="multipart", + part_size=part_size, + part_urls=None, + upload_id="up_stream", + ) + + +def _single_session(): + return UploadSessionResponse( + finalize_token="tok", + headers={}, + mode="single", + upload_id="up_single", + url="https://storage.test/single", + ) + + +def _eager_session(part_urls, part_size): + return UploadSessionResponse( + finalize_token="tok", + headers={}, + mode="multipart", + part_size=part_size, + part_urls=part_urls, + upload_id="up_eager", + ) + + +# === GAP 1: large files omit declared_size_bytes to force streaming ========== + + +def test_large_file_omits_declared_size_bytes(seq_pool, monkeypatch): + """A file larger than STREAMING_THRESHOLD must OMIT declared_size_bytes so + the server opens a streaming (per-part on-demand mint) session — the only + way a presigned URL can't expire mid-transfer on a slow large upload. + """ + part_size = 5 * MIB + content = b"x" * (STREAMING_THRESHOLD + 1) + api = _make_api() + + def fake_mint(upload_id, x_upload_finalize_token, mint_upload_parts_request, **kw): + parts = [ + MintedUploadPartResponse(part_number=n, url=f"https://storage.test/s/{n}") + for n in mint_upload_parts_request.part_numbers + ] + return MintUploadPartsResponse(parts=parts) + + monkeypatch.setattr(api, "mint_upload_parts_handler", fake_mint) + captured: List[Any] = [] + _capture_create(monkeypatch, api, _streaming_session(part_size), captured) + _patch_finalize(monkeypatch, api, []) + + api.upload_file(content, max_concurrency=2, round_base_delay=0.0) + + # The field must be ABSENT on the wire, not sent as null — the server only + # opens a streaming session when declared_size_bytes is omitted entirely. + assert captured[0].declared_size_bytes is None + assert "declared_size_bytes" not in captured[0].to_dict() + + +def test_small_file_declares_size_bytes(seq_pool, monkeypatch): + """A file at or below STREAMING_THRESHOLD keeps the known-size path (single + PUT / eager fast path) by declaring its size. + """ + content = b"y" * STREAMING_THRESHOLD + api = _make_api() + captured: List[Any] = [] + _capture_create(monkeypatch, api, _single_session(), captured) + _patch_finalize(monkeypatch, api, []) + + api.upload_file(content) + + assert captured[0].declared_size_bytes == len(content) + assert captured[0].to_dict()["declared_size_bytes"] == len(content) + + +def test_streaming_part_bytes_keyed_by_part_number(seq_pool, monkeypatch): + """Byte ranges are keyed to the server part_number, never to completion or + mint-arrival order: each minted URL receives exactly its own slice even when + parts finish out of order under concurrency (coverage Gap 1). + """ + part_size = 5 * MIB + # Three distinguishable regions so a mis-keyed slice is detectable. + content = b"A" * part_size + b"B" * part_size + b"C" * (part_size // 2) + api = _make_api() + + def fake_mint(upload_id, x_upload_finalize_token, mint_upload_parts_request, **kw): + # Return the requested parts in REVERSED order to prove URL↔part matching + # is by number, not array position. + nums = list(mint_upload_parts_request.part_numbers) + parts = [ + MintedUploadPartResponse(part_number=n, url=f"https://storage.test/s/{n}") + for n in reversed(nums) + ] + return MintUploadPartsResponse(parts=parts) + + for n in (1, 2, 3): + seq_pool.responses[f"https://storage.test/s/{n}"] = [_Resp(200, etag=f'"e{n}"')] + monkeypatch.setattr(api, "mint_upload_parts_handler", fake_mint) + _capture_create(monkeypatch, api, _streaming_session(part_size), []) + _patch_finalize(monkeypatch, api, []) + + api.upload_file(content, max_concurrency=3, round_base_delay=0.0) + + by_url = {c["url"]: c["body"] for c in seq_pool.calls} + assert by_url["https://storage.test/s/1"] == content[0:part_size] + assert by_url["https://storage.test/s/2"] == content[part_size:2 * part_size] + assert by_url["https://storage.test/s/3"] == content[2 * part_size:] + + +# === GAP 2: storage pool has a bounded connect timeout ======================= + + +def test_storage_pool_has_connect_timeout(monkeypatch): + """The bare storage pool must bound TCP+TLS establishment so a dead endpoint + fails fast into the re-sweep, while leaving the read/transfer unbounded (a + large upload legitimately takes minutes; the watchdog owns stall detection). + """ + monkeypatch.setattr(uploads_mod, "_STORAGE_POOL", None) + pool = uploads_mod._storage_pool() + timeout = pool.connection_pool_kw["timeout"] + assert timeout.connect_timeout == uploads_mod.STORAGE_CONNECT_TIMEOUT + assert timeout.read_timeout is None + + +# === Ergo #1: error hierarchy + helpers re-exported at top level ============= + + +def test_error_hierarchy_importable_from_top_level(): + import hotdata + + for name in ( + "UploadError", + "SessionCreateError", + "MalformedSessionError", + "SizeLimitError", + "StorageError", + "StorageTransportError", + "MissingETagError", + "MintPartError", + "FinalizeError", + "UploadCancelledError", + "UploadOptions", + "UploadProgress", + "UploadSource", + "tqdm_progress", + ): + assert hasattr(hotdata, name), f"hotdata.{name} not re-exported" + assert name in hotdata.__all__, f"{name} missing from hotdata.__all__" + assert getattr(hotdata, name) is getattr(uploads_mod, name) + + +# === Ergo #2: every phase failure is catchable as UploadError ================ + + +def test_session_create_error_wraps_api_exception(seq_pool, monkeypatch): + from hotdata.uploads import SessionCreateError, UploadError + + api = _make_api() + + def boom(create_upload_request, **kw): + raise ApiException(status=501, reason="PRESIGN_UNSUPPORTED") + + monkeypatch.setattr(api, "create_upload_session_handler", boom) + + with pytest.raises(UploadError) as ei: # catchable via the unified base + api.upload_file(b"hi") + assert isinstance(ei.value, SessionCreateError) + assert ei.value.status == 501 # specific status stays detectable + assert isinstance(ei.value.__cause__, ApiException) + assert seq_pool.calls == [] + + +def test_finalize_error_wraps_api_exception(seq_pool, monkeypatch): + from hotdata.uploads import FinalizeError, UploadError + + api = _make_api() + _capture_create(monkeypatch, api, _single_session(), []) + + # Let the real _finalize_handler run (that's where the wrapping lives) and make + # the underlying generated call raise, so we exercise the ApiException→FinalizeError mapping. + def boom(self, upload_id, finalize_token, body=None, _request_timeout=None): + raise ApiException(status=409, reason="already finalized") + + monkeypatch.setattr(uploads_mod._GeneratedUploadsApi, "finalize_upload_handler", boom) + + with pytest.raises(UploadError) as ei: + api.upload_file(b"hello") + assert isinstance(ei.value, FinalizeError) + assert ei.value.status == 409 + assert isinstance(ei.value.__cause__, ApiException) + + +def test_mint_error_wraps_api_exception_and_is_retryable(seq_pool, monkeypatch): + from hotdata.uploads import MintPartError, UploadError + + part_size = 5 * MIB + content = b"z" * part_size + api = _make_api() + + def always_fail_mint(upload_id, token, req, **kw): + raise ApiException(status=500, reason="mint boom") + + monkeypatch.setattr(api, "mint_upload_parts_handler", always_fail_mint) + _capture_create(monkeypatch, api, _streaming_session(part_size), []) + _patch_finalize(monkeypatch, api, []) + + with pytest.raises(UploadError) as ei: + api.upload_file(content, max_concurrency=1, max_extra_rounds=1, round_base_delay=0.0) + assert isinstance(ei.value, MintPartError) + assert ei.value.status == 500 + # Retryable: it survived the round budget, so it is tagged exhausted (Ergo #6). + assert ei.value.exhausted is True + assert ei.value.rounds_attempted == 2 + + +# === Ergo #3: tqdm_progress adapter (cumulative -> delta, thread-safe) ======= + + +def test_tqdm_progress_converts_cumulative_to_delta(): + from hotdata.uploads import tqdm_progress + + class _Bar: + def __init__(self): + self.n = 0 + + def update(self, delta): + self.n += delta + + bar = _Bar() + cb = tqdm_progress(bar) + for done in (0, 8, 16, 20): # cumulative, as upload_file emits + cb(done, 20) + assert bar.n == 20 + + +def test_tqdm_progress_is_thread_safe(): + from hotdata.uploads import tqdm_progress + + class _Bar: + def __init__(self): + self.n = 0 + + def update(self, delta): + # Non-atomic read-modify-write; races would corrupt without the lock. + cur = self.n + self.n = cur + delta + + bar = _Bar() + cb = tqdm_progress(bar) + total = 1000 + # Many threads each advancing the cumulative counter by 1. + threads = [threading.Thread(target=cb, args=(i + 1, total)) for i in range(total)] + # Deliver in order so cumulative is monotone; the lock protects bar.n. + for t in threads: + t.start() + t.join() + assert bar.n == total + + +# === Ergo #4: cancellation ==================================================== + + +def test_cancel_event_preset_aborts_before_any_put(seq_pool, monkeypatch): + from hotdata.uploads import UploadCancelledError + + api = _make_api() + _capture_create(monkeypatch, api, _single_session(), []) + finalize_calls: List[Any] = [] + _patch_finalize(monkeypatch, api, finalize_calls) + + cancel = threading.Event() + cancel.set() + with pytest.raises(UploadCancelledError): + api.upload_file(b"data", cancel_event=cancel) + assert seq_pool.calls == [] # nothing PUT + assert finalize_calls == [] # never finalized + + +def test_cancel_event_midflight_aborts_multipart(seq_pool, monkeypatch): + from hotdata.uploads import UploadCancelledError + + part_size = 5 * MIB + content = b"c" * (3 * part_size) + api = _make_api() + urls = [f"https://storage.test/p{n}" for n in (1, 2, 3)] + for u in urls: + seq_pool.responses[u] = [_Resp(200, etag='"e"')] + session = UploadSessionResponse( + finalize_token="tok", + headers={}, + mode="multipart", + part_size=part_size, + part_urls=urls, + upload_id="up", + ) + _capture_create(monkeypatch, api, session, []) + finalize_calls: List[Any] = [] + _patch_finalize(monkeypatch, api, finalize_calls) + + cancel = threading.Event() + # Cancel as soon as the first part is PUT; serial so it takes effect on part 2. + orig_request = seq_pool.request + + def request_then_cancel(method, url, **kw): + resp = orig_request(method, url, **kw) + cancel.set() + return resp + + monkeypatch.setattr(seq_pool, "request", request_then_cancel) + + with pytest.raises(UploadCancelledError): + api.upload_file( + content, max_concurrency=1, cancel_event=cancel, round_base_delay=0.0 + ) + # Aborted before all three parts were PUT, and never finalized. + assert len([c for c in seq_pool.calls if c["method"] == "PUT"]) < 3 + assert finalize_calls == [] + + +# === Ergo #6: exhausted vs terminal marker on StorageError =================== + + +def test_storage_error_marked_exhausted_after_resweep(seq_pool, monkeypatch): + part_size = 5 * MIB + content = b"d" * part_size + api = _make_api() + url = "https://storage.test/p1" + seq_pool.responses[url] = [_Resp(500, body=b"boom")] # persistent 500 + _capture_create(monkeypatch, api, _eager_session([url], part_size), []) + _patch_finalize(monkeypatch, api, []) + + with pytest.raises(StorageError) as ei: + api.upload_file(content, max_concurrency=1, max_extra_rounds=1, round_base_delay=0.0) + assert ei.value.status == 500 + assert ei.value.exhausted is True + assert ei.value.rounds_attempted == 2 + + +def test_terminal_error_not_marked_exhausted(seq_pool, monkeypatch): + # A short-read is a deterministic, non-retryable (terminal) failure: it fails + # fast on the first round and must NOT be tagged exhausted. + part_size = 5 * MIB + api = _make_api() + + class _ShortSource: + size = 2 * part_size + + def read_range(self, offset, length): + return b"x" * (length - 1) # always one byte short + + def reader(self): + raise AssertionError("multipart path does not use reader()") + + monkeypatch.setattr(uploads_mod, "_make_source", lambda source, size: _ShortSource()) + url1, url2 = "https://storage.test/p1", "https://storage.test/p2" + _capture_create(monkeypatch, api, _eager_session([url1, url2], part_size), []) + _patch_finalize(monkeypatch, api, []) + + from hotdata.uploads import UploadError + + with pytest.raises(UploadError) as ei: + api.upload_file(b"ignored", max_concurrency=1, round_base_delay=0.0) + assert ei.value.exhausted is False + + +# === Ergo #7: public request_timeout forwarded to control-plane calls ======== + + +def test_request_timeout_forwarded_to_create_and_finalize(seq_pool, monkeypatch): + api = _make_api() + seen: Dict[str, Any] = {} + + def fake_create(create_upload_request, _request_timeout=None): + seen["create"] = _request_timeout + return _single_session() + + def fake_finalize(upload_id, finalize_token, body, _request_timeout): + seen["finalize"] = _request_timeout + return FinalizeUploadResponse( + content_type=None, + created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + size_bytes=0, + status="ready", + upload_id=upload_id, + ) + + monkeypatch.setattr(api, "create_upload_session_handler", fake_create) + monkeypatch.setattr(api, "_finalize_handler", fake_finalize) + + api.upload_file(b"hi", request_timeout=7.5) + + assert seen["create"] == 7.5 + assert seen["finalize"] == 7.5 diff --git a/tests/test_uploads.py b/tests/test_uploads.py index b40441a..6e46242 100644 --- a/tests/test_uploads.py +++ b/tests/test_uploads.py @@ -582,11 +582,13 @@ def test_exported_uploads_api_is_enhanced() -> None: # --- Session/finalize API error paths ------------------------------------- -def test_create_session_501_propagates_as_api_exception( +def test_create_session_501_wrapped_as_session_create_error( monkeypatch: pytest.MonkeyPatch, fake_pool: _FakeStoragePool, tmp_path: Any ) -> None: - # A 501 PRESIGN_UNSUPPORTED from POST /v1/uploads surfaces as ApiException - # (the documented fallback signal), and no storage PUT is attempted. + # A 501 PRESIGN_UNSUPPORTED from POST /v1/uploads surfaces as a + # SessionCreateError (an UploadError, so `except UploadError` catches the whole + # flow) that still exposes .status == 501 (the documented fallback signal) and + # chains the ApiException as __cause__. No storage PUT is attempted. src = tmp_path / "f.bin" src.write_bytes(b"x" * 10) api = _make_api() @@ -595,9 +597,11 @@ def boom(create_upload_request: Any, **kwargs: Any) -> Any: raise ApiException(status=501, reason="PRESIGN_UNSUPPORTED") monkeypatch.setattr(api, "create_upload_session_handler", boom) - with pytest.raises(ApiException) as ei: + with pytest.raises(uploads_mod.SessionCreateError) as ei: api.upload_file(str(src)) + assert isinstance(ei.value, uploads_mod.UploadError) assert ei.value.status == 501 + assert isinstance(ei.value.__cause__, ApiException) assert fake_pool.calls == [] # never reached storage @@ -667,9 +671,13 @@ def test_size_limit_guard_on_part_size_hint( assert ei.value.what == "part_size" -def test_size_limit_guard_on_declared_size( +def test_huge_declared_size_streams_instead_of_declaring( monkeypatch: pytest.MonkeyPatch, fake_pool: _FakeStoragePool, tmp_path: Any ) -> None: + # A file larger than STREAMING_THRESHOLD takes the streaming path, which OMITS + # declared_size_bytes — so an over-i64 size is never sent (and so never an + # obstacle), mirroring the Rust SDK. The part-size hint is still range-guarded + # (see test_size_limit_guard_on_part_size_hint). src = tmp_path / "f.bin" src.write_bytes(b"x" * 10) api = _make_api() @@ -684,9 +692,20 @@ class _HugeStat: "stat", lambda p: _HugeStat() if str(p) == str(src) else real_stat(p), ) - with pytest.raises(SizeLimitError) as ei: + + captured: List[Any] = [] + + class _Stop(Exception): + pass + + def capture_create(create_upload_request: Any, **kwargs: Any) -> Any: + captured.append(create_upload_request) + raise _Stop() # short-circuit; we only care about the create request + + monkeypatch.setattr(api, "create_upload_session_handler", capture_create) + with pytest.raises(_Stop): api.upload_file(str(src)) - assert ei.value.what == "declared_size_bytes" + assert captured[0].declared_size_bytes is None def test_custom_part_retry_is_forwarded(