diff --git a/tests/downloaders/__init__.py b/tests/downloaders/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/downloaders/test_download_hardening.py b/tests/downloaders/test_download_hardening.py new file mode 100644 index 00000000..fe46cb91 --- /dev/null +++ b/tests/downloaders/test_download_hardening.py @@ -0,0 +1,122 @@ +"""Downloader hardening invariants. + +Pins the adaptive start-at-cap wiring, the tail-probe gate for small segments, +and the byte-range-slice protections: a server ignoring an item's Range (200 +instead of 206) must fail the attempt rather than corrupt the merge, a retried +slice must be rewritten whole instead of resumed with a clobbered Range, and +media requests must ask for identity encoding. +""" + +import importlib +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest + +dl = importlib.import_module("unshackle.core.downloaders.requests") + +PARENT = bytes(range(256)) * 256 # 64KiB parent resource, distinctive bytes +SLICE_START = 1024 +SLICE_END = 5119 # inclusive; 4KiB slice + + +class _Handler(BaseHTTPRequestHandler): + def do_GET(self): + server = self.server + with server.lock: + server.requests.append(dict(self.headers)) + nth = len(server.requests) + rng = self.headers.get("Range") + if server.honor_range and rng: + start, end = (int(x) for x in rng.removeprefix("bytes=").split("-")) + body = PARENT[start : end + 1] + self.send_response(206) + self.send_header("Content-Range", f"bytes {start}-{end}/{len(PARENT)}") + else: + body = PARENT + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + if server.truncate_first and nth == 1: + self.wfile.write(body[: len(body) // 2]) + return # connection closes mid-body -> client sees a short read + self.wfile.write(body) + + def log_message(self, *args): + pass + + +@pytest.fixture +def server(): + srv = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) + srv.lock = threading.Lock() + srv.requests = [] + srv.honor_range = True + srv.truncate_first = False + thread = threading.Thread(target=srv.serve_forever, daemon=True) + thread.start() + yield srv + srv.shutdown() + thread.join(timeout=5) + + +def _url(srv, path="/parent.bin"): + return f"http://127.0.0.1:{srv.server_address[1]}{path}" + + +def _run(srv, tmp_path, urls, **kwargs): + for _ in dl.requests(urls, output_dir=tmp_path, filename="seg_{i:04}.bin", **kwargs): + pass + return sorted(tmp_path.glob("seg_*.bin")) + + +def test_adaptive_controller_starts_at_cap(server, tmp_path, monkeypatch): + seen = {} + real = dl.AdaptiveWorkerController + + def recorder(*args, **kw): + ctl = real(*args, **kw) + seen["cap"], seen["start"] = ctl.cap, ctl.target + return ctl + + monkeypatch.setattr(dl, "AdaptiveWorkerController", recorder) + files = _run(server, tmp_path, [{"url": _url(server)} for _ in range(8)], max_workers=4, adaptive=True) + assert len(files) == 8 + assert seen == {"cap": 4, "start": 4} + + +def test_small_segments_never_probed(server, tmp_path, monkeypatch): + def no_probe(*args, **kw): + raise AssertionError("tail boost probed a small-segment batch") + + monkeypatch.setattr(dl, "_probe_ranged", no_probe) + files = _run(server, tmp_path, [{"url": _url(server)} for _ in range(30)], max_workers=4, adaptive=True) + assert len(files) == 30 + assert all(f.read_bytes() == PARENT for f in files) + + +def test_range_item_ignored_by_server_fails_not_corrupts(server, tmp_path, monkeypatch): + server.honor_range = False + monkeypatch.setattr(dl, "RETRY_WAIT", 0.01) + urls = [{"url": _url(server), "headers": {"Range": f"bytes={SLICE_START}-{SLICE_END}"}}] + with pytest.raises(Exception): + _run(server, tmp_path, urls, max_workers=1) + # nothing finalized: a 200 full-parent body must never become the segment + assert not list(tmp_path.glob("seg_*.bin")) + + +def test_range_item_retry_rewrites_slice(server, tmp_path, monkeypatch): + server.truncate_first = True + monkeypatch.setattr(dl, "RETRY_WAIT", 0.01) + urls = [{"url": _url(server), "headers": {"Range": f"bytes={SLICE_START}-{SLICE_END}"}}] + files = _run(server, tmp_path, urls, max_workers=1) + assert files[0].read_bytes() == PARENT[SLICE_START : SLICE_END + 1] + # the retry must re-send the item's own Range, never a resume Range from the partial + ranges = [r.get("Range") for r in server.requests] + assert ranges == [f"bytes={SLICE_START}-{SLICE_END}"] * len(ranges) + assert len(ranges) >= 2 + + +def test_segment_requests_ask_identity_encoding(server, tmp_path): + _run(server, tmp_path, [{"url": _url(server)}], max_workers=1) + assert all(r.get("Accept-Encoding") == "identity" for r in server.requests) diff --git a/tests/downloaders/test_hedge_handles.py b/tests/downloaders/test_hedge_handles.py new file mode 100644 index 00000000..7c7e633c --- /dev/null +++ b/tests/downloaders/test_hedge_handles.py @@ -0,0 +1,111 @@ +"""Hedged segment racers must release their file handles before the batch completes. + +A superseded loser that keeps its .!dev handle open makes the manifest merge's +stray sweep raise PermissionError (WinError 32) on Windows. This test forces a +hedge with a trickling first response and checks the invariant at the exact +moment merge would run. +""" + +import importlib +import os +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import pytest + +from unshackle.core.constants import DOWNLOAD_CANCELLED + +# the downloaders package re-exports the requests() function, shadowing the submodule +dl = importlib.import_module("unshackle.core.downloaders.requests") + +SEGMENTS = 10 +SLOW_INDEX = 4 +FAST_BODY = b"F" * 4096 +SLOW_BODY = b"S" * 40960 +TRICKLE_CHUNK = 2048 +TRICKLE_DELAY = 0.05 + + +class _Handler(BaseHTTPRequestHandler): + def do_GET(self): + with self.server.count_lock: + self.server.counts[self.path] = self.server.counts.get(self.path, 0) + 1 + nth = self.server.counts[self.path] + slow_path = self.path == f"/seg/{SLOW_INDEX}" + body = SLOW_BODY if slow_path else FAST_BODY + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + if slow_path and nth == 1: + # trickle only the first attempt so the hedge retry wins instantly + for i in range(0, len(body), TRICKLE_CHUNK): + try: + self.wfile.write(body[i : i + TRICKLE_CHUNK]) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + return + time.sleep(TRICKLE_DELAY) + else: + self.wfile.write(body) + + def log_message(self, *args): + pass + + +@pytest.fixture +def trickle_server(): + server = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) + server.counts = {} + server.count_lock = threading.Lock() + threading.Thread(target=server.serve_forever, daemon=True).start() + yield server + server.shutdown() + + +@pytest.mark.unit +def test_hedge_loser_releases_handles_before_merge(tmp_path, trickle_server, monkeypatch): + monkeypatch.setattr(dl, "HEDGE_MIN_WAIT", 0.3) + monkeypatch.setattr(dl, "RETRY_WAIT", 0.1) + DOWNLOAD_CANCELLED.clear() + threads_before = {t.ident for t in threading.enumerate()} + + port = trickle_server.server_address[1] + urls = [f"http://127.0.0.1:{port}/seg/{i}" for i in range(SEGMENTS)] + start = time.monotonic() + for _event in dl.requests(urls=urls, output_dir=tmp_path, filename="{i:02}.mp4", max_workers=4): + pass + wall = time.monotonic() - start + + with trickle_server.count_lock: + slow_hits = trickle_server.counts[f"/seg/{SLOW_INDEX}"] + assert slow_hits == 2, "hedge never fired; the race under test was not exercised" + + # the batch must not stall draining the loser's stream (trickle alone takes ~1s) + assert wall < 5 + + for i in range(SEGMENTS): + expected = SLOW_BODY if i == SLOW_INDEX else FAST_BODY + assert (tmp_path / f"{i:02}.mp4").read_bytes() == expected + + strays = list(tmp_path.glob("*.!dev")) + fd_dir = Path("/proc/self/fd") + if fd_dir.exists(): + open_targets = set() + for fd in fd_dir.iterdir(): + try: + open_targets.add(os.readlink(fd)) + except OSError: + continue + held = [s.name for s in strays if str(s) in open_targets] + assert not held, f"racer still holds open handles at completion: {held}" + for stray in strays: + stray.unlink() # raises PermissionError (WinError 32) on Windows if a handle is open + + leaked = [ + t.name + for t in threading.enumerate() + if t.ident not in threads_before and t.name.startswith("ThreadPoolExecutor") + ] + assert not leaked, f"racer threads still alive at completion: {leaked}" diff --git a/unshackle/core/downloaders/requests.py b/unshackle/core/downloaders/requests.py index f2b97d97..8f0ada2b 100644 --- a/unshackle/core/downloaders/requests.py +++ b/unshackle/core/downloaders/requests.py @@ -262,10 +262,7 @@ def _plan_tail_parts(size: int, spare: int) -> list[tuple[int, int]]: n_parts = max(1, min(spare, math.ceil(size / TAIL_BOOST_PART_SIZE))) part_size = math.ceil(size / n_parts) return [ - (s, e) - for i in range(n_parts) - for s, e in [(i * part_size, min(size - 1, (i + 1) * part_size - 1))] - if s <= e + (s, e) for i in range(n_parts) for s, e in [(i * part_size, min(size - 1, (i + 1) * part_size - 1))] if s <= e ] @@ -460,6 +457,10 @@ def download( _time = time.time use_raw = _is_requests_session(session) + # item carries its own Range (DASH SegmentBase / HLS EXT-X-BYTERANGE slice): never + # overwrite it with a resume Range (that would fetch the parent's tail, not the slice); + # retries rewrite the whole slice in "wb" mode instead + item_range = _has_range_header(kwargs) attempts = 1 written = 0 @@ -480,18 +481,19 @@ def download( request_kwargs.setdefault("max_retries", 0) else: request_kwargs.setdefault("timeout", (CONNECT_TIMEOUT, READ_TIMEOUT)) + req_headers = dict(request_kwargs.get("headers", {}) or {}) + # media bytes must arrive unrecoded: transparent CDN (de)compression breaks + # Content-Length accounting and is meaningless on ranged requests + req_headers.setdefault("Accept-Encoding", "identity") if part_mode: - req_headers = dict(request_kwargs.get("headers", {}) or {}) req_headers["Range"] = f"bytes={part_offset + written}-{part_end}" - request_kwargs["headers"] = req_headers - elif resume_offset > 0: - req_headers = dict(request_kwargs.get("headers", {}) or {}) + elif resume_offset > 0 and not item_range: req_headers["Range"] = f"bytes={resume_offset}-" - request_kwargs["headers"] = req_headers + request_kwargs["headers"] = req_headers stream = session.get(url, stream=True, **request_kwargs) - if (not part_mode) and resume_offset > 0 and stream.status_code == 416: + if (not part_mode) and (not item_range) and resume_offset > 0 and stream.status_code == 416: # our Range started past the end (a stale or oversized .!dev from a prior run). # discard it and restart clean rather than raise_for_status → burn retries. try: @@ -505,13 +507,21 @@ def download( stream.raise_for_status() - resumed = (not part_mode) and resume_offset > 0 and stream.status_code == 206 + # item_range 206s reflect the slice's own Range, not a resume; those retries rewrite in "wb" + resumed = (not part_mode) and (not item_range) and resume_offset > 0 and stream.status_code == 206 if (not part_mode) and resume_offset > 0 and not resumed: resume_offset = 0 if part_mode and stream.status_code != 206: raise IOError(f"expected 206 for ranged part, got {stream.status_code}") + if item_range and not part_mode and stream.status_code != 206: + # server ignored the slice's Range and sent the whole parent; writing that + # as the segment would silently corrupt the merge, so fail the attempt + raise IOError(f"expected 206 for byte-range segment, got {stream.status_code}") if use_rnet: content_length = stream.content_length or 0 + ce = (stream.headers.get("Content-Encoding") or stream.headers.get("content-encoding") or "").lower() + if ce in ("gzip", "deflate", "br"): + content_length = 0 else: try: content_length = int(stream.headers.get("Content-Length", "0")) @@ -838,7 +848,10 @@ def requests( seg_durations: list[float] = [] hedged: set[int] = set() - controller = AdaptiveWorkerController(cap=max_workers) if adaptive else None + # start at the cap and adapt downward under CDN pressure: ramping up from a low target + # can at best converge to the throughput the cap already gives, so there is nothing to + # gain by starting low. error-burst backoff and AIMD recovery still apply from there. + controller = AdaptiveWorkerController(cap=max_workers, start=max_workers) if adaptive else None # feed CDN error signals (429/timeout/reset) from each retrying segment into the controller on_retry_cb = (lambda _exc: controller.record_error(time.time())) if controller else None @@ -968,6 +981,12 @@ def requests( # blocking range probes can't stall the drain loop. if DOWNLOAD_CANCELLED.is_set() or not _tail_boost_engages(len(remaining), len(pending), target): return + # completed segments predict the tail's sizes: below the min every probe would + # decline anyway, and at real-CDN RTT those sequential probes cost seconds on the + # main loop, so skip probing and release the tail through the normal path + if seg_done and (total_bytes / len(seg_done)) < TAIL_BOOST_MIN_SEGMENT_SIZE: + tail_skipped.update(index for index, _ in remaining) + return boosted = 0 for index, url_item in list(remaining): if boosted >= TAIL_BOOST_MAX_PER_CYCLE: