From d7e08e979b4f9a0f34bae4de9b19b035b45da457 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Fri, 7 Aug 2026 11:31:11 -0600 Subject: [PATCH] fix(downloader): surface failed downloads instead of returning silently --- tests/downloaders/test_failure_surfacing.py | 164 ++++++++++++++++++++ unshackle/core/downloaders/requests.py | 157 +++++++++++++------ 2 files changed, 273 insertions(+), 48 deletions(-) create mode 100644 tests/downloaders/test_failure_surfacing.py diff --git a/tests/downloaders/test_failure_surfacing.py b/tests/downloaders/test_failure_surfacing.py new file mode 100644 index 00000000..2ac346e7 --- /dev/null +++ b/tests/downloaders/test_failure_surfacing.py @@ -0,0 +1,164 @@ +import socket +import sys +import threading + +import pytest +import requests as rq + +from unshackle.core.constants import DOWNLOAD_CANCELLED +from unshackle.core.downloaders.requests import _has_range_header, download +from unshackle.core.downloaders.requests import requests as requests_downloader + +# the package re-exports the `requests` function, shadowing the module of the same name +downloader = sys.modules[download.__module__] + +BODY = b"".join(f"{i:05d} payload\n".encode() for i in range(500)) + + +class _Server: + """Serves BODY, with a switch for the two failure modes retries have to handle.""" + + def __init__(self, mode): + self.mode = mode + self.sock = socket.socket() + self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self.sock.bind(("127.0.0.1", 0)) + self.sock.listen(8) + self.url = f"http://127.0.0.1:{self.sock.getsockname()[1]}/seg.mp4" + self.requests = [] + threading.Thread(target=self._serve, daemon=True).start() + + def _serve(self): + while True: + try: + conn, _ = self.sock.accept() + except OSError: + return + try: + request = conn.recv(65536).decode("latin1") + headers = dict( + (name.strip().lower(), value.strip()) + for name, _, value in (line.partition(":") for line in request.split("\r\n")) + if value + ) + self.requests.append(headers) + ranged = "range" in headers + if self.mode == "unavailable": + conn.sendall(b"HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") + elif ranged: + conn.sendall( + b"HTTP/1.1 416 Range Not Satisfiable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ) + else: + conn.sendall( + f"HTTP/1.1 200 OK\r\nContent-Length: {len(BODY)}\r\nConnection: close\r\n\r\n".encode() + BODY + ) + except Exception: + pass + finally: + conn.close() + + def close(self): + self.sock.close() + + +@pytest.fixture(autouse=True) +def _clear_cancel(): + # the batch failure path sets the process-global cancel; leaking it makes every later + # test that waits on the event (the speed limiter) return instantly + DOWNLOAD_CANCELLED.clear() + yield + DOWNLOAD_CANCELLED.clear() + + +@pytest.fixture +def no_retry_wait(monkeypatch): + monkeypatch.setattr(downloader, "RETRY_WAIT", 0) + + +@pytest.fixture +def server(request, no_retry_wait): + server = _Server(request.param) + yield server + server.close() + + +@pytest.mark.parametrize("server", ["unavailable"], indirect=True) +def test_exhausted_retries_raise_instead_of_reporting_success(tmp_path, server): + save_path = tmp_path / "0.mp4" + with pytest.raises(rq.HTTPError): + list(download(url=server.url, save_path=save_path, session=rq.Session())) + assert len(server.requests) == downloader.MAX_ATTEMPTS + assert not save_path.exists() + + +@pytest.mark.parametrize("server", ["unavailable"], indirect=True) +def test_cancelled_download_stays_silent(tmp_path, server, monkeypatch): + # one attempt only, so a cancel check ordered after the exhaustion raise would fail here + monkeypatch.setattr(downloader, "MAX_ATTEMPTS", 1) + save_path = tmp_path / "0.mp4" + DOWNLOAD_CANCELLED.set() + list(download(url=server.url, save_path=save_path, session=rq.Session())) + assert not save_path.exists() + + +@pytest.mark.parametrize("server", ["range_unsatisfiable"], indirect=True) +def test_stale_partial_restarts_clean_on_416(tmp_path, server): + save_path = tmp_path / "0.mp4" + save_path.with_name("0.mp4.!dev").write_bytes(b"\x00" * len(BODY)) + list(download(url=server.url, save_path=save_path, session=rq.Session())) + assert save_path.read_bytes() == BODY + assert [r.get("range") for r in server.requests] == [f"bytes={len(BODY)}-", None] + + +@pytest.mark.parametrize("server", ["range_unsatisfiable"], indirect=True) +def test_slice_segment_never_gets_a_resume_range(tmp_path, server): + save_path = tmp_path / "0.mp4" + save_path.with_name("0.mp4.!dev").write_bytes(b"\x00" * 128) + with pytest.raises(OSError): + list(download(url=server.url, save_path=save_path, session=rq.Session(), headers={"Range": "bytes=0-127"})) + assert [r["range"] for r in server.requests] == ["bytes=0-127"] * downloader.MAX_ATTEMPTS + + +@pytest.mark.parametrize("server", ["range_unsatisfiable"], indirect=True) +def test_identity_encoding_requested_unless_the_caller_set_one(tmp_path, server): + save_path = tmp_path / "0.mp4" + list(download(url=server.url, save_path=save_path, session=rq.Session())) + assert server.requests[-1]["accept-encoding"] == "identity" + + list( + download( + url=server.url, save_path=tmp_path / "1.mp4", session=rq.Session(), headers={"accept-encoding": "gzip"} + ) + ) + assert server.requests[-1]["accept-encoding"] == "gzip" + + +@pytest.mark.parametrize("server", ["unavailable"], indirect=True) +def test_failed_segment_fails_the_batch(tmp_path, server): + # a segment the batch never delivers must raise, not merge short of the segment count + with pytest.raises(rq.HTTPError): + list( + requests_downloader( + urls=[server.url, server.url], + output_dir=tmp_path, + filename="{i:01}.mp4", + max_workers=2, + ) + ) + assert sorted(p.name for p in tmp_path.iterdir() if p.suffix == ".mp4") == [] + + +@pytest.mark.parametrize( + ("item", "expected"), + [ + ({}, False), + ({"headers": None}, False), + ({"headers": {}}, False), + ({"headers": {"Accept": "*/*"}}, False), + ({"headers": {"Range": "bytes=0-127"}}, True), + ({"headers": {"range": "bytes=0-127"}}, True), + ], +) +def test_has_range_header(item, expected): + assert _has_range_header(item) is expected diff --git a/unshackle/core/downloaders/requests.py b/unshackle/core/downloaders/requests.py index 8f793b1c..495fb1fd 100644 --- a/unshackle/core/downloaders/requests.py +++ b/unshackle/core/downloaders/requests.py @@ -173,6 +173,17 @@ def _is_content_encoded(value: Optional[str]) -> bool: return any(coding.strip() not in ("", "identity") for coding in (value or "").lower().split(",")) +def _has_range_header(item: dict[str, Any]) -> bool: + """ + Whether the caller asked for a byte-range slice of a parent resource + (DASH SegmentBase, HLS EXT-X-BYTERANGE). + + Resume and ranged-parallel requests must leave such a Range alone. Replacing it + fetches the parent's bytes instead of the slice's, which corrupts the merge. + """ + return any(str(k).lower() == "range" for k in (item.get("headers") or {})) + + def _probe_ranged(url: str, session: Any, **kwargs: Any) -> tuple[int, bool]: headers = {**(kwargs.get("headers") or {}), "Range": "bytes=0-0"} rest = {k: v for k, v in kwargs.items() if k != "headers"} @@ -224,6 +235,9 @@ def _dispatch_parts( f.truncate(total_size) events: Queue[dict[str, Any]] = Queue() + # scoped to this file's parts only: setting the process-global cancel here would + # poison the plain-download fallback and silently stand down sibling tracks + abort = threading.Event() def _worker(start: int, end: int) -> None: for ev in download( @@ -232,6 +246,7 @@ def _dispatch_parts( session=session, part_offset=start, part_end=end, + abort=abort, **kwargs, ): events.put(ev) @@ -275,7 +290,7 @@ def _dispatch_parts( exc = fut.exception() if exc: worker_error = True - DOWNLOAD_CANCELLED.set() + abort.set() raise exc advance = 0 @@ -311,6 +326,7 @@ def download( part_offset: Optional[int] = None, part_end: Optional[int] = None, claimed: Optional[Callable[[], bool]] = None, + abort: Optional[threading.Event] = None, **kwargs: Any, ) -> Generator[dict[str, Any], None, None]: """ @@ -342,6 +358,8 @@ def download( part_end: Inclusive end byte of the part. Required when `part_offset` is set. claimed: Optional predicate checked before every attempt; when it returns True the download stops silently (another worker already delivered this file). + abort: Optional event that silently stops this download when set, without + touching the process-global cancel. kwargs: Any extra keyword arguments to pass to the session.get() call. Use this for one-time request changes like a header, cookie, or proxy. For example, to request Byte-ranges use e.g., `headers={"Range": "bytes=0-128"}`. @@ -370,12 +388,15 @@ def download( _time = time.time use_raw = _is_requests_session(session) + item_range = _has_range_header(kwargs) attempts = 1 written = 0 while True: if claimed is not None and claimed(): return + if abort is not None and abort.is_set(): + return if not part_mode: written = 0 last_speed_refresh = _time() @@ -388,23 +409,35 @@ def download( request_kwargs.setdefault("read_timeout", READ_TIMEOUT) else: request_kwargs.setdefault("timeout", (CONNECT_TIMEOUT, READ_TIMEOUT)) + req_headers = dict(request_kwargs.get("headers", {}) or {}) + if not any(str(k).lower() == "accept-encoding" for k in req_headers): + req_headers["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 (not item_range) and resume_offset > 0 and stream.status_code == 416: + try: + stream.close() + except Exception: + pass + tmp_file.unlink(missing_ok=True) + resume_offset = 0 + continue + stream.raise_for_status() - resumed = (not part_mode) and resume_offset > 0 and stream.status_code == 206 + 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: + raise IOError(f"expected 206 for byte-range segment, got {stream.status_code}") content_encoded = False if use_rnet: content_length = stream.content_length or 0 @@ -449,9 +482,6 @@ def download( with open(save_path if part_mode else tmp_file, file_mode, buffering=file_buffering) as f: if part_mode: f.seek(part_offset + written) - elif not resumed and content_length > 0: - f.truncate(content_length) - f.seek(0) _write = f.write @@ -472,7 +502,7 @@ def download( for chunk in chunks: if DOWNLOAD_CANCELLED.is_set(): break - if claimed is not None and claimed(): + if (claimed is not None and claimed()) or (abort is not None and abort.is_set()): # close the handle or Windows can't delete the stray .!dev at merge (WinError 32) try: stream.close() @@ -507,9 +537,6 @@ def download( except Exception: pass - if not part_mode and not resumed and content_length > 0 and written != content_length: - f.truncate(written) - if DOWNLOAD_CANCELLED.is_set() and (not content_length or written < content_length): # cancelled mid-stream: keep the partial .!dev for resume, never finalize return @@ -535,11 +562,11 @@ def download( if claimed is not None and claimed(): # a superseded loser's error must not retry or kill the batch return - if DOWNLOAD_CANCELLED.is_set() or attempts == MAX_ATTEMPTS: - if part_mode and not DOWNLOAD_CANCELLED.is_set(): - raise + if DOWNLOAD_CANCELLED.is_set() or (abort is not None and abort.is_set()): return - if not part_mode: + if attempts == MAX_ATTEMPTS: + raise + if not part_mode and not item_range: resume_offset = tmp_file.stat().st_size if tmp_file.exists() else 0 time.sleep(RETRY_WAIT) attempts += 1 @@ -674,7 +701,7 @@ def requests( url_item = urls[0] try: ranged_used = False - if max_workers > 1 and _speed_limiter is None: + if max_workers > 1 and _speed_limiter is None and not _has_range_header(url_item): total_size, supports_ranges = _probe_ranged(url_item["url"], session) if supports_ranges and total_size >= RANGE_PARALLEL_MIN_SIZE: try: @@ -706,6 +733,20 @@ def requests( DOWNLOAD_CANCELLED.set() yield dict(downloaded="[yellow]CANCELLED") raise + except Exception as e: + yield dict(downloaded="[red]FAILED") + if debug_logger: + debug_logger.log( + level="ERROR", + operation="downloader_failed", + message=f"Download failed: {e}", + error=e, + context={ + "url_count": len(urls), + "output_dir": str(output_dir), + }, + ) + raise else: # Segmented download with thread pool # Speed is tracked here on the main thread, not in workers @@ -725,43 +766,63 @@ def requests( seg_lock = threading.Lock() seg_start: dict[int, float] = {} seg_done: set[int] = set() + seg_active: dict[int, int] = {} seg_durations: list[float] = [] hedged: set[int] = set() def _download_worker(index: int, url_item: dict[str, Any], hedge: bool = False) -> None: - item = dict(url_item) - save_path = item.pop("save_path") - if save_path.exists(): # finished in a previous run - with seg_lock: - if index in seg_done: - return - seg_done.add(index) - event_queue.put(dict(file_downloaded=save_path, written=save_path.stat().st_size)) - event_queue.put(dict(advance=1)) - return - # each racer writes to its own target so a loser never touches the final file; - # the .!dev suffix keeps strays visible to the manifest parsers' cleanup - target = save_path.with_name(f"{save_path.name}.{'h' if hedge else 'p'}.!dev") - if not hedge: - seg_start[index] = time.time() - for event in download( - session=session, - segmented=segmented_batch, - save_path=target, - claimed=lambda: index in seg_done, - **item, - ): - if "file_downloaded" in event: + with seg_lock: + seg_active[index] = seg_active.get(index, 0) + 1 + left = False + try: + item = dict(url_item) + save_path = item.pop("save_path") + if save_path.exists(): # finished in a previous run with seg_lock: if index in seg_done: - target.unlink(missing_ok=True) return seg_done.add(index) - if not hedge: - seg_durations.append(time.time() - seg_start[index]) - os.replace(target, save_path) - event = dict(event, file_downloaded=save_path) - event_queue.put(event) + event_queue.put(dict(file_downloaded=save_path, written=save_path.stat().st_size)) + event_queue.put(dict(advance=1)) + return + # each racer writes to its own target so a loser never touches the final file; + # the .!dev suffix keeps strays visible to the manifest parsers' cleanup + target = save_path.with_name(f"{save_path.name}.{'h' if hedge else 'p'}.!dev") + if not hedge: + seg_start[index] = time.time() + for event in download( + session=session, + segmented=segmented_batch, + save_path=target, + claimed=lambda: index in seg_done, + **item, + ): + if "file_downloaded" in event: + with seg_lock: + if index in seg_done: + target.unlink(missing_ok=True) + return + seg_done.add(index) + if not hedge: + seg_durations.append(time.time() - seg_start[index]) + os.replace(target, save_path) + event = dict(event, file_downloaded=save_path) + event_queue.put(event) + except Exception: + # leaving this racer and reading the survivor count must be one critical + # section, or two racers failing together each see the other as the survivor + # and the segment goes missing without an exception + with seg_lock: + seg_active[index] -= 1 + left = True + survivor = seg_active[index] > 0 or index in seg_done + if survivor: + return + raise + finally: + if not left: + with seg_lock: + seg_active[index] -= 1 futures = [pool.submit(_download_worker, i, url) for i, url in enumerate(urls)] pending = set(futures)