From d4690e5c6ba9cb4b01dc6d294e455dc042fb83dc Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Mon, 24 Aug 2026 15:10:36 -0600 Subject: [PATCH] fix(downloader): bound requests after a cut and stop retry nesting Request 4 MiB sub-ranges once a host cuts a body, skip the session's own urllib3 Retry inside download(), fail the ranged path at once on a 4xx, and cap resumes at 50 of 1 MiB each. Adds a deny fault profile. --- scripts/fault_server.py | 14 ++ tests/downloaders/test_bounded_requests.py | 93 ++++++++++++ tests/downloaders/test_download_hardening.py | 24 ++++ tests/downloaders/test_ranged_parallel.py | 64 ++++++++- unshackle/core/downloaders/requests.py | 142 +++++++++++++++++-- 5 files changed, 316 insertions(+), 21 deletions(-) create mode 100644 tests/downloaders/test_bounded_requests.py diff --git a/scripts/fault_server.py b/scripts/fault_server.py index eeebca0c..40863439 100644 --- a/scripts/fault_server.py +++ b/scripts/fault_server.py @@ -17,6 +17,7 @@ Fault profiles: tail-slow heavily throttle the last K segments' bodies throttle-cut throttle every body and close it (FIN) after N seconds or N bytes flaky-first fail attempt 1 per path (503), succeed on retry + deny reply 403 to every request (a permanent error no retry can fix) loopback no faults; payloads served from memory (overhead ceiling) Range requests are honored (206 + Content-Range) so byte-range / hedge / tail-boost @@ -26,6 +27,7 @@ Standalone: uv run python scripts/fault_server.py --profile rate-limit --segments 32 uv run python scripts/fault_server.py --profile stall --stall-secs 4 --port 8080 uv run python scripts/fault_server.py --profile throttle-cut --segments 1 --seg-size 71303168 + uv run python scripts/fault_server.py --profile deny --deny-after 1 Importable: build a ``FaultServer``, serve it on a thread, and hash output against ``segment_payload(path, size)``. @@ -78,6 +80,9 @@ class FaultProfile: flaky_status: int = 503 # status returned by a flaky-first failure cut_secs: float = 0.0 # close the body cleanly (FIN) after this long (0 = never) cut_bytes: int = 0 # close the body cleanly (FIN) after this many bytes (0 = never) + deny: bool = False # reply deny_status to every request past deny_after + deny_status: int = 403 # status returned by a deny + deny_after: int = 0 # serve this many attempts per path before denying (0 = deny all) PROFILES: dict[str, FaultProfile] = { @@ -96,6 +101,10 @@ PROFILES: dict[str, FaultProfile] = { # and then ends short, so each retry makes progress but no single attempt finishes. Not # seed-gated like stall and reset, because the real host cuts every attempt, not a share. "throttle-cut": FaultProfile("throttle-cut", body_rate_kib=480, cut_secs=6.5, cut_bytes=5 * MIB // 2), + # a permanently refused resource (expired token, geo block): no retry and no fallback + # recovers, so a client that keeps trying is burning attempts for nothing. --deny-after + # serves the first N attempts per path, which lets a range probe through first. + "deny": FaultProfile("deny", deny=True), } @@ -231,6 +240,9 @@ class FaultHandler(BaseHTTPRequestHandler): attempt = server.note_attempt(self.path) prof = server.profile + if prof.deny and attempt > prof.deny_after: + self._send_status(prof.deny_status) + return if prof.flaky_first and attempt == 1: # first attempt per path always fails, retry recovers self._send_status(prof.flaky_status) return @@ -388,6 +400,7 @@ def _build_arg_parser() -> argparse.ArgumentParser: p.add_argument("--body-rate-kib", type=int, help="Override: throttle every body to this KiB/s (0 = off).") p.add_argument("--tail-slow", type=int, help="Override: throttle the last K segments.") p.add_argument("--tail-rate-kib", type=int, help="Override: tail-slow throttle rate (KiB/s).") + p.add_argument("--deny-after", type=int, help="Override: serve this many attempts per path before denying.") p.add_argument("--cut-secs", type=float, help="Override: close each body (FIN) after this long (0 = never).") p.add_argument("--cut-bytes", type=int, help="Override: close each body (FIN) after this many bytes (0 = never).") return p @@ -407,6 +420,7 @@ def main() -> None: "tail_rate_kib": args.tail_rate_kib, "cut_secs": args.cut_secs, "cut_bytes": args.cut_bytes, + "deny_after": args.deny_after, } profile = resolve_profile(args.profile, overrides) server = FaultServer(profile, args.seg_size, args.segments, args.seed, port=args.port) diff --git a/tests/downloaders/test_bounded_requests.py b/tests/downloaders/test_bounded_requests.py new file mode 100644 index 00000000..9b49e725 --- /dev/null +++ b/tests/downloaders/test_bounded_requests.py @@ -0,0 +1,93 @@ +import sys +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest +import requests as rq + +from unshackle.core.constants import DOWNLOAD_CANCELLED +from unshackle.core.downloaders.requests import download + +# the package re-exports the `requests` function, shadowing the module of the same name +downloader = sys.modules[download.__module__] + +BODY = bytes(range(256)) * 256 # 64 KiB +CUT_AFTER = 10 * 1024 +CHUNK = 16 * 1024 + + +class _CutServer(ThreadingHTTPServer): + """Serves BODY but ends any body longer than CUT_AFTER early, like a per-request-limited CDN.""" + + allow_reuse_address = True + daemon_threads = True + + def __init__(self): + self.ranges: list[str] = [] + self.cuts = 0 + self.connections = 0 + super().__init__(("127.0.0.1", 0), _CutHandler) + self.url = f"http://127.0.0.1:{self.server_address[1]}/seg.mp4" + + def handle_error(self, request, client_address): # the client hanging up is expected here + pass + + +class _CutHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, *_): + pass + + def setup(self): + super().setup() + self.server.connections += 1 + + def do_GET(self): + header = self.headers.get("Range") + self.server.ranges.append(header) + start, end = 0, len(BODY) - 1 + if header: + start_s, _, end_s = header.removeprefix("bytes=").partition("-") + start, end = int(start_s), (int(end_s) if end_s else len(BODY) - 1) + body = BODY[start : end + 1] + self.send_response(206 if header else 200) + if header: + self.send_header("Content-Range", f"bytes {start}-{end}/{len(BODY)}") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + if len(body) > CUT_AFTER: + self.server.cuts += 1 + self.wfile.write(body[:CUT_AFTER]) + self.close_connection = True + return + self.wfile.write(body) + + +@pytest.fixture +def server(monkeypatch): + DOWNLOAD_CANCELLED.clear() + monkeypatch.setattr(downloader, "RETRY_WAIT", 0) + monkeypatch.setattr(downloader, "REQUEST_CHUNK_SIZE", CHUNK) + monkeypatch.setattr(downloader, "MIN_REQUEST_CHUNK", 4 * 1024) + monkeypatch.setattr(downloader, "MIN_RESUME_PROGRESS", 1024) + server = _CutServer() + threading.Thread(target=server.serve_forever, daemon=True).start() + yield server + server.shutdown() + server.server_close() + DOWNLOAD_CANCELLED.clear() + + +def test_a_cut_body_bounds_the_following_requests(tmp_path, server): + save_path = tmp_path / "0.mp4" + list(download(url=server.url, save_path=save_path, session=rq.Session())) + + assert save_path.read_bytes() == BODY + # the open-ended first request is cut, the chunk-sized one after it is cut too, and the + # halved chunk fits under the host's limit, so every later request completes whole + assert server.cuts == 2 + spans = [r.removeprefix("bytes=").split("-") for r in server.ranges if r and r.split("-")[-1]] + assert spans and all(int(end) - int(start) + 1 <= CHUNK for start, end in spans) + # the bounded requests are cheap because they share one pooled connection + assert server.connections < len(server.ranges) diff --git a/tests/downloaders/test_download_hardening.py b/tests/downloaders/test_download_hardening.py index a79cd518..04ad331c 100644 --- a/tests/downloaders/test_download_hardening.py +++ b/tests/downloaders/test_download_hardening.py @@ -201,3 +201,27 @@ def test_transient_429_with_retry_after_does_not_abort_track(server, tmp_path, m assert all(f.read_bytes() == PARENT for f in files) assert not dl.DOWNLOAD_CANCELLED.is_set() assert len(server.requests) == 5 # 3 successes + 2 rate-limited attempts + + +def test_session_retry_does_not_nest_inside_an_attempt(server, tmp_path, monkeypatch): + # a service session mounts urllib3 Retry with a 429 forcelist; download() owns segment + # retries, so the two must not multiply (MAX_ATTEMPTS requests, not MAX_ATTEMPTS x 6) + from requests import Session + from requests.adapters import HTTPAdapter, Retry + + dl.DOWNLOAD_CANCELLED.clear() + server.fail_429_first = 10**6 + monkeypatch.setattr(dl, "RETRY_WAIT", 0.01) + session = Session() + adapter = HTTPAdapter( + max_retries=Retry(total=5, backoff_factor=0, status_forcelist=[429, 500, 502, 503, 504]), + ) + session.mount("http://", adapter) + session.mount("https://", adapter) + + with pytest.raises(Exception): + run(server, tmp_path, [{"url": url(server)}], max_workers=1, session=session) + + assert len(server.requests) == dl.MAX_ATTEMPTS + # the caller's session keeps its own retry policy + assert session.get_adapter("http://x/").max_retries.total == 5 diff --git a/tests/downloaders/test_ranged_parallel.py b/tests/downloaders/test_ranged_parallel.py index b6372594..1f7f0ad5 100644 --- a/tests/downloaders/test_ranged_parallel.py +++ b/tests/downloaders/test_ranged_parallel.py @@ -10,12 +10,14 @@ Pins the byte-range fan-out that splits one large file across parts: - that failure sets only the ranged download's LOCAL abort event, never the process-global DOWNLOAD_CANCELLED (which would poison sibling tracks); - ``probe_ranged`` parses the Content-Range total on a 206 and declines a 200 - or a content-encoded 206. + or a content-encoded 206; +- the public single-URL path falls back to a sequential download only when the + server does not really support ranges, and re-raises a 4xx instead of paying + MAX_ATTEMPTS a second time from byte 0. -Tests 2/3 exercise ``dispatch_parts`` directly: the public single-URL path in -``requests()`` deliberately swallows a ranged failure and falls back to a plain -sequential download, so the raise the invariant guarantees is only observable at -the dispatch layer. +Tests 2/3 exercise ``dispatch_parts`` directly: the public path swallows a +range-support failure and falls back, so the raise that the invariant guarantees +is only observable at the dispatch layer. """ import importlib @@ -23,13 +25,18 @@ import threading from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import pytest -from requests import Session +from requests import HTTPError, Session dl = importlib.import_module("unshackle.core.downloaders.requests") PAYLOAD = bytes(range(256)) * 128 # 32 KiB, distinctive repeating byte pattern +class _Resp: + def __init__(self, status_code): + self.status_code = status_code + + class _RangeHandler(BaseHTTPRequestHandler): def do_GET(self): srv = self.server @@ -48,6 +55,13 @@ class _RangeHandler(BaseHTTPRequestHandler): self.wfile.write(b"\x00") return + if mode == "deny_parts" and rng != "bytes=0-0": + # probe passes, every part window is refused for good (expired token, geo block) + self.send_response(403) + self.send_header("Content-Length", "0") + self.end_headers() + return + # always_200 ignores Range entirely; probe_only_206 answers the 0-0 probe # with a real 206 but hands every wider part window a 200 full body serve_full = mode == "always_200" or not rng or (mode == "probe_only_206" and rng != "bytes=0-0") @@ -162,3 +176,41 @@ def test_probe_ranged_parses_content_range(server): assert dl.probe_ranged(url(server), session) == (0, False) server.mode = "gzip_probe" assert dl.probe_ranged(url(server), session) == (0, False) + + +def test_permanent_part_failure_is_not_retried_sequentially(server, tmp_path, monkeypatch): + server.mode = "deny_parts" + monkeypatch.setattr(dl, "RETRY_WAIT", 0.01) + monkeypatch.setattr(dl, "RANGE_PARALLEL_MIN_SIZE", 8 * 1024) + monkeypatch.setattr(dl, "RANGE_PARALLEL_PART_SIZE", 8 * 1024) + with pytest.raises(HTTPError): + run(server, tmp_path, [{"url": url(server)}], max_workers=4) + # a sequential fallback would issue a plain unranged GET; only the probe and the + # part windows may appear, each part capped at MAX_ATTEMPTS + assert [r for r in server.requests if r is None] == [] + parts = len({r for r in server.requests if r and r != "bytes=0-0"}) + assert len(server.requests) <= 1 + parts * dl.MAX_ATTEMPTS + # nothing half-written is left behind for a later run to mistake for a finished file + assert sorted(tmp_path.glob("seg_*")) == [] + + +def test_range_unsupported_falls_back_to_sequential(server, tmp_path, monkeypatch): + server.mode = "probe_only_206" + monkeypatch.setattr(dl, "RETRY_WAIT", 0.01) + monkeypatch.setattr(dl, "RANGE_PARALLEL_MIN_SIZE", 8 * 1024) + monkeypatch.setattr(dl, "RANGE_PARALLEL_PART_SIZE", 8 * 1024) + files = run(server, tmp_path, [{"url": url(server)}], max_workers=4) + assert len(files) == 1 + assert files[0].read_bytes() == PAYLOAD + assert not files[0].with_name(files[0].name + ".!dev").exists() + + +def test_permanent_status_unwraps_causes(): + assert dl.permanent_status(HTTPError(response=_Resp(403))) == 403 + assert dl.permanent_status(HTTPError(response=_Resp(429))) is None + assert dl.permanent_status(HTTPError(response=_Resp(503))) is None + assert dl.permanent_status(IOError("short read")) is None + # RnetSession wraps the HTTPError one level down as __cause__ + wrapped = RuntimeError("max retries") + wrapped.__cause__ = HTTPError(response=_Resp(404)) + assert dl.permanent_status(wrapped) == 404 diff --git a/unshackle/core/downloaders/requests.py b/unshackle/core/downloaders/requests.py index 7df1c78d..8288f4c7 100644 --- a/unshackle/core/downloaders/requests.py +++ b/unshackle/core/downloaders/requests.py @@ -9,9 +9,10 @@ import sys import threading import time import traceback -from collections import deque +from collections import OrderedDict, deque from concurrent.futures import FIRST_COMPLETED, wait from concurrent.futures.thread import ThreadPoolExecutor +from copy import copy from datetime import datetime, timezone from email.utils import parsedate_to_datetime from http.cookiejar import CookieJar @@ -21,7 +22,7 @@ from string import Formatter from typing import Any, Callable, Generator, MutableMapping, Optional, Union, cast from requests import Session -from requests.adapters import HTTPAdapter +from requests.adapters import HTTPAdapter, Retry from rich import filesize from unshackle.core.constants import DOWNLOAD_CANCELLED @@ -31,8 +32,16 @@ MAX_ATTEMPTS = 5 RETRY_WAIT = 2 # a cut that still delivered bytes is resumable and refreshes the attempt budget; bound it # so a server that sends a little and closes, over and over, still terminates -MAX_RESUMES = 500 -MIN_RESUME_PROGRESS = 64 * 1024 +MAX_RESUMES = 50 +MIN_RESUME_PROGRESS = 1024 * 1024 + +# a host that cuts a body it already started (per-request time limit, per-flow throttle) will +# cut the next one at the same place. After such a cut each following request asks for a +# bounded sub-range instead of the whole remainder, halving it on every further cut until the +# requests fit under the limit: a sub-range that completes is a success, so it costs neither +# an attempt nor a backoff nap. +REQUEST_CHUNK_SIZE = 4 * 1024 * 1024 +MIN_REQUEST_CHUNK = 256 * 1024 PROGRESS_WINDOW = 2 # read timeout bounds the gap between chunks so a quiet connection errors instead of hanging @@ -302,6 +311,23 @@ class AdaptiveWorkerController: return self.target +def permanent_status(exc: BaseException) -> Optional[int]: + """HTTP status of ``exc`` when it is a client error that no retry can fix. + + Both session types raise requests' HTTPError; RnetSession wraps it in MaxRetriesError + as ``__cause__``, so unwrap one level like retry_sleep does. 408 and 429 do not count: + they are transient, and a fan-out that trips a per-client rate limit can still succeed + over one sequential connection. + """ + response = getattr(exc, "response", None) + if response is None: + response = getattr(getattr(exc, "__cause__", None), "response", None) + status = getattr(response, "status_code", None) + if isinstance(status, int) and 400 <= status < 500 and status not in (408, 429): + return status + return None + + def retry_sleep(exc: Exception, attempts: int) -> float: """Seconds to wait before retry number ``attempts``. @@ -344,6 +370,24 @@ def is_requests_session(session: Any) -> bool: return isinstance(session, Session) +def no_retry_session(session: Session) -> Session: + """A shallow copy of ``session`` whose adapters never retry, sharing its state and pools. + + download() owns segment retries, so a mounted urllib3 ``Retry`` nests inside a single + attempt: one 429 carrying ``Retry-After: 60`` costs five one-minute sleeps before + download() even sees the failure. requests takes no per-call retry override, and track + threads share the caller's session, so remounting it is not an option. A shallow adapter + copy keeps the poolmanager object, so TLS and pooling behaviour does not change. + """ + view = copy(session) + view.adapters = OrderedDict() + for prefix, adapter in session.adapters.items(): + twin = copy(adapter) + twin.max_retries = Retry(0, read=False) + view.adapters[prefix] = twin + return view + + def is_rnet_session(session: Any) -> bool: """Whether the HTTP session is an RnetSession (it uses ``resp.stream()``).""" from unshackle.core.session import RnetSession @@ -372,6 +416,7 @@ def probe_ranged(url: str, session: Any, **kwargs: Any) -> tuple[int, bool]: rest.setdefault("max_retries", 0) else: rest.setdefault("timeout", (CONNECT_TIMEOUT, READ_TIMEOUT)) + session = no_retry_session(session) try: resp = session.get(url, stream=True, headers=headers, **rest) except Exception: @@ -684,6 +729,9 @@ def download( to request Byte-ranges use e.g., `headers={"Range": "bytes=0-128"}`. """ session = session or Session() + if is_requests_session(session): + # same reason as the rnet max_retries=0 below: download() owns segment retries + session = no_retry_session(session) part_mode = part_offset is not None and part_end is not None # partial data lives at the .!dev name and is renamed into place on completion: @@ -715,6 +763,11 @@ def download( attempts = 1 written = 0 resumes = 0 + # None until a cut proves the host bounds one request; the sequential path also needs a + # known complete length, or a short body could not be told apart from a finished download + request_limit: Optional[int] = None + known_total: Optional[int] = None + progress_seeded = False while True: if claimed is not None and claimed(): return @@ -742,10 +795,24 @@ def download( # Content-Length accounting and is meaningless on ranged requests if not any(str(k).lower() == "accept-encoding" for k in req_headers): req_headers["Accept-Encoding"] = "identity" - if part_mode: - req_headers["Range"] = f"bytes={part_offset + written}-{part_end}" + # inclusive end of this request when it covers only part of what is left, else None + chunk_end: Optional[int] = None + if part_mode and part_offset is not None and part_end is not None: # spelled so mypy narrows + req_start = part_offset + written + chunk_end = part_end + if request_limit is not None and part_end - req_start + 1 > request_limit: + chunk_end = req_start + request_limit - 1 + req_headers["Range"] = f"bytes={req_start}-{chunk_end}" elif resume_offset > 0 and not item_range: - req_headers["Range"] = f"bytes={resume_offset}-" + if ( + request_limit is not None + and known_total is not None + and known_total - resume_offset > request_limit + ): + chunk_end = resume_offset + request_limit - 1 + req_headers["Range"] = f"bytes={resume_offset}-{chunk_end}" + else: + req_headers["Range"] = f"bytes={resume_offset}-" request_kwargs["headers"] = req_headers stream = session.get(url, stream=True, **request_kwargs) @@ -789,6 +856,8 @@ def download( resume_offset = 0 written = 0 continue + elif content_range[2] is not None: + known_total = content_range[2] if (not part_mode) and resume_offset > 0 and not resumed: resume_offset = 0 if part_mode and part_offset is not None and part_end is not None: # spelled so mypy narrows @@ -800,13 +869,13 @@ def download( if ( content_range is None or content_range[0] != part_offset + written - or content_range[1] > part_end + or (chunk_end is not None and content_range[1] > chunk_end) or (content_range[2] is not None and content_range[2] <= part_end) ): # a wrong start would land bytes at the wrong offset in the shared file, - # an end past part_end would spill into a sibling part's region, and a - # total at or below part_end means the resource changed under the part - # plan; fail the attempt before any byte is written + # an end past what this request asked for would spill into a sibling + # part's region, and a total at or below part_end means the resource + # changed under the part plan; fail the attempt before any byte is written raise IOError(f"part {part_offset}-{part_end} got Content-Range {content_range!r}") if use_rnet: content_encoded = is_content_encoded( @@ -822,6 +891,10 @@ def download( except ValueError: content_length = 0 + if known_total is None and not part_mode and not item_range and stream.status_code == 200: + # a 200 carries the whole resource, so its length is the complete length + known_total = content_length or None + if item_range and not part_mode and stream.status_code != 206: # server ignored the slice's Range and sent the whole parent (RFC 9110 allows # a 200 here); writing that as the segment would silently corrupt the merge, so @@ -870,7 +943,10 @@ def download( secured = written if part_mode else resume_offset total_size = (resume_offset + content_length) if resumed and content_length > 0 else content_length - if not segmented and not part_mode: + if not segmented and not part_mode and not progress_seeded: + # only the first response of this call sees the true total, and the bytes a + # previous run left on disk are progress that the loop below never reports + progress_seeded = True if total_size > 0: yield dict(total=total_size) else: @@ -961,10 +1037,16 @@ def download( # and skip the size checks below (a short read here is expected, not an error) return - if part_mode: + if part_mode and part_offset is not None and part_end is not None: # spelled so mypy narrows expected = part_end - part_offset + 1 - if written < expected: + target = expected if chunk_end is None else chunk_end - part_offset + 1 + if written < target: raise IOError(f"Failed to read part {part_offset}-{part_end}: got {written}/{expected}") + if written < expected: + # this request's sub-range arrived whole; the rest of the part is the next + # request, not a retry, so the attempt budget starts over for it + attempts = 1 + continue elif content_length and written < content_length: # applies to segments too: a connection FIN'd mid-body is a clean EOF, so the # chunk loop ends without raising, so without this check a truncated segment would @@ -973,6 +1055,10 @@ def download( raise IOError(f"Failed to read {content_length} bytes from the track URI.") if not part_mode: + if chunk_end is not None and known_total is not None and resume_offset + written < known_total: + resume_offset += written + attempts = 1 + continue os.replace(tmp_file, save_path) yield dict(file_downloaded=save_path, written=resume_offset + written) if segmented: @@ -1003,6 +1089,11 @@ def download( if progress - secured >= MIN_RESUME_PROGRESS and resumes < MAX_RESUMES: resumes += 1 attempts = 0 + if not item_range: + # the host cut a body it was already delivering, so ask for less next time + request_limit = ( + REQUEST_CHUNK_SIZE if request_limit is None else max(MIN_REQUEST_CHUNK, request_limit // 2) + ) elif on_retry is not None: try: on_retry(exc) @@ -1500,7 +1591,9 @@ def requests( ranged_used = True except KeyboardInterrupt: raise - except Exception: + except Exception as exc: + # the pre-truncated save_path is full-size, so the sequential path would + # read it as already complete: drop it either way, no resume from it save_path = url_item.get("save_path") if save_path: sp = Path(save_path) @@ -1509,6 +1602,25 @@ def requests( stray.unlink(missing_ok=True) except OSError: pass + if DOWNLOAD_CANCELLED.is_set(): + return + status = permanent_status(exc) + fatal = status is not None + if debug_logger: + debug_logger.log( + level="DEBUG", + operation="ranged_parallel_failed", + message=f"ranged download failed: {exc.__class__.__name__}", + context={ + "error": exc.__class__.__name__, + "status": status, + "action": "raise" if fatal else "sequential_fallback", + }, + ) + if fatal: + # retrying a 4xx from byte 0 only burns MAX_ATTEMPTS again at + # twice the wall clock + raise if not ranged_used: yield from download( session=session,