diff --git a/docs/guide/cli-reference.md b/docs/guide/cli-reference.md index 741a9e2f..ed2c8777 100644 --- a/docs/guide/cli-reference.md +++ b/docs/guide/cli-reference.md @@ -171,6 +171,7 @@ Keep only certain track types, or skip certain track types. Attachments are alwa | Flag | Default | Description | |---|---|---| | `--workers` | downloader default | Per-track download threads. | +| `--adaptive-workers` | off | Scale per-track segment workers dynamically (up to the `--workers` cap) based on measured CDN throughput and errors. | | `--downloads` | `1` | Number of tracks downloaded concurrently. | | `--no-cache` | off | Bypass the title cache. | | `--reset-cache` | off | Clear the title cache. | diff --git a/docs/guide/downloading.md b/docs/guide/downloading.md index 1fd25e67..8826261a 100644 --- a/docs/guide/downloading.md +++ b/docs/guide/downloading.md @@ -462,6 +462,7 @@ Two related flags: | Flag | Purpose | | --- | --- | | `--workers N` | Threads used per track for segment downloads. Default depends on the downloader. | +| `--adaptive-workers` | Opt-in: start with a moderate per-track worker count and ramp it up or back off based on measured CDN throughput and errors, capped at `--workers`. Off by default (fixed worker count). | | `--downloads N` | Number of tracks downloaded concurrently. Default `1`. | | `--slow [MIN-MAX]` | Add a delay between titles to look more like a real device. `--slow` alone means 60-120s; `--slow 20-40` sets a custom range. Minimum 20s. | | `--no-cache` | Bypass the title cache for this download. | diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 4f34e023..3d22f2c6 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -465,8 +465,16 @@ TLS-fingerprinting and HTTP client settings for the rnet-based session. | `browser` | str | `"Chrome131"` | Impersonation preset. Must be an exact preset name (e.g. `Chrome131`, `Firefox135`, `Edge101`, `Safari18`, `OkHttp4_12`, `OkHttp5`, `Opera118`); unknown names raise an error. | | `http1_only` | bool | *(unset)* | Force HTTP/1.1. | | `http2_only` | bool | *(unset)* | Force HTTP/2. | -| `pool_max_idle_per_host` | int | *(unset)* | Connection-pool tuning. | -| `pool_max_size` | int | *(unset)* | Connection-pool tuning. | +| `connect_timeout` | int | `10` | Connect timeout in seconds. | +| `read_timeout` | int | `30` | Read timeout in seconds (max gap between response chunks). | +| `timeout` | int | *(unset)* | Total request timeout in seconds. | +| `allow_redirects` | bool | *(unset)* | Follow redirects automatically (rnet client default when unset). | +| `pool_idle_timeout` | int | `55` | Seconds an idle pooled connection is kept (default stays under typical ~60s CDN idle kills). | +| `pool_max_idle_per_host` | int | `16` | Max idle connections kept per host. | +| `pool_max_size` | int | *(unset)* | Max total connections in the pool. | +| `tcp_keepalive` | int | `30` | TCP keepalive time in seconds. | +| `tcp_keepalive_interval` | int | *(unset)* | TCP keepalive probe interval in seconds. | +| `tcp_keepalive_retries` | int | *(unset)* | TCP keepalive probe retry count. | | `tcp_nodelay` | bool | *(unset)* | Disable Nagle's algorithm. | ```yaml diff --git a/scripts/bench_subtitle_backends.py b/scripts/bench_subtitle_backends.py deleted file mode 100644 index 5333099f..00000000 --- a/scripts/bench_subtitle_backends.py +++ /dev/null @@ -1,98 +0,0 @@ -#!/usr/bin/env python3 -""" -Benchmark subtitle conversion backends to (re-)tune the preference ranks in -``unshackle/core/tracks/subtitle_convert.py``. - -Runs every backend that can read each input file, converting to a target format (default -SRT), and reports cue count, leaked ASS override tags, and output size — so you can compare -fidelity per (source, target) pair on real files. Read-only: copies inputs to a temp dir. - -Usage: - uv run python scripts/bench_subtitle_backends.py [ ...] [--target SRT] - -Example: - uv run python scripts/bench_subtitle_backends.py downloads/ -""" - -from __future__ import annotations - -import argparse -import re -import shutil -import tempfile -from pathlib import Path - -from unshackle.core.tracks import subtitle_convert as sc -from unshackle.core.tracks.subtitle import Subtitle - -Codec = Subtitle.Codec - -EXT_TO_CODEC = { - ".srt": Codec.SubRip, - ".vtt": Codec.WebVTT, - ".ass": Codec.SubStationAlphav4, - ".ssa": Codec.SubStationAlpha, - ".ttml": Codec.TimedTextMarkupLang, - ".smi": Codec.SAMI, - ".sami": Codec.SAMI, -} - - -def gather(paths: list[str]) -> list[Path]: - files: list[Path] = [] - for p in paths: - path = Path(p) - if path.is_dir(): - files.extend(f for f in path.rglob("*") if f.suffix.lower() in EXT_TO_CODEC) - elif path.suffix.lower() in EXT_TO_CODEC: - files.append(path) - return sorted(files) - - -def metrics(text: str) -> tuple[int, int, int]: - cues = len(re.findall(r"-->", text)) - ass_residue = len(re.findall(r"\{\\", text)) - return cues, ass_residue, len(text) - - -def main() -> None: - ap = argparse.ArgumentParser() - ap.add_argument("paths", nargs="+", help="subtitle files or directories") - ap.add_argument("--target", default="SRT", help="target codec value (SRT, VTT, ASS, ...)") - args = ap.parse_args() - - target = Codec(args.target.upper()) - files = gather(args.paths) - if not files: - print("No subtitle files found.") - return - - tmp = Path(tempfile.mkdtemp(prefix="subbench_")) - print(f"{'file':40} {'source':10} {'backend':12} {'ok':3} {'cues':>5} {'resid':>5} {'bytes':>7}") - for f in files: - source = EXT_TO_CODEC[f.suffix.lower()] - if source == target: - continue - for backend in sc.REGISTRY: - if not (backend.is_available() and backend.can_convert(source, target)): - continue - work = tmp / f"{f.stem}.{backend.name}{f.suffix}" - shutil.copy2(f, work) - sub = Subtitle(url="x", language="en", codec=source) - sub.path = work - try: - # Call the backend directly so each row reflects only that backend (no fallback). - out = work.with_suffix(f".{target.value.lower()}") - backend.convert(sub, target, out) - cues, resid, size = metrics(out.read_text("utf8", errors="replace")) - print( - f"{f.name[:40]:40} {source.name[:10]:10} {backend.name:12} {'Y':3} {cues:>5} {resid:>5} {size:>7}" - ) - except Exception as e: # noqa: BLE001 - benchmark reports failures, does not raise - print( - f"{f.name[:40]:40} {source.name[:10]:10} {backend.name:12} {'N':3} {'-':>5} {'-':>5} {'-':>7} {type(e).__name__}" - ) - - -if __name__ == "__main__": - main() diff --git a/tests/core/test_adaptive_workers.py b/tests/core/test_adaptive_workers.py new file mode 100644 index 00000000..e1a2cb03 --- /dev/null +++ b/tests/core/test_adaptive_workers.py @@ -0,0 +1,237 @@ +"""Unit tests for the opt-in adaptive worker controller (AIMD hill-climb over segment +concurrency) and the plumbing that carries the flag into the downloader. Pure controller +tests only. No threads, sockets, or real downloads; the clock is supplied explicitly.""" + +from __future__ import annotations + +import inspect +from functools import partial +from pathlib import Path + +import pytest + +from unshackle.core.downloaders.requests import (TAIL_BOOST_MIN_SEGMENT_SIZE, TAIL_BOOST_PART_SIZE, + AdaptiveWorkerController, _has_range_header, _plan_tail_parts, + _tail_boost_engages, requests) +from unshackle.core.tracks.track import DownloadContext + +pytestmark = pytest.mark.unit + +TICK = 4.0 +WINDOW = 10.0 +DT = 0.5 + + +def _fill_window(c: AdaptiveWorkerController, rate: float, now: float) -> float: + """Record ``WINDOW`` seconds of samples at ``rate`` bytes/sec ending on the 0.5s grid. + + Records a sample at ``now`` and every 0.5s up to ``now + WINDOW`` so the oldest kept + sample sits exactly ``WINDOW`` old. That is the condition ``update`` needs to start acting. + """ + c.record_bytes(int(rate * DT), now) + end = now + WINDOW + while now < end - 1e-9: + now += DT + c.record_bytes(int(rate * DT), now) + return now + + +def _run_tick(c: AdaptiveWorkerController, rate: float, now: float, errors: int = 0) -> tuple[int, float]: + """Advance one tick recording samples at ``rate``, optionally inject ``errors``, then update.""" + end = now + TICK + while now < end - 1e-9: + now += DT + c.record_bytes(int(rate * DT), now) + for _ in range(errors): + c.record_error(now) + return c.update(now), now + + +def test_starts_at_min_six_cap() -> None: + assert AdaptiveWorkerController(cap=16).target == 6 + assert AdaptiveWorkerController(cap=8).target == 6 + assert AdaptiveWorkerController(cap=4).target == 4 # min(6, cap) + assert AdaptiveWorkerController(cap=1).target == 2 # clamped up to ADAPTIVE_MIN + + +def test_no_op_before_full_window() -> None: + c = AdaptiveWorkerController(cap=16, tick=TICK, window=WINDOW) + assert c.update(0.0) == 6 # arms last_tick before any samples are recorded + c.record_bytes(500_000, 0.0) + c.record_bytes(500_000, 5.0) + # a full tick has elapsed, but only 5s of window exists -> hold + assert c.update(5.0) == 6 + c.record_bytes(500_000, 8.0) + assert c.update(9.0) == 6 + + +def test_ramps_up_while_speed_improves() -> None: + c = AdaptiveWorkerController(cap=16, tick=TICK, window=WINDOW) + now = 0.0 + c.update(now) # arm + rate = 1_000_000.0 + now = _fill_window(c, rate, now) + targets = [c.update(now)] # first evaluation: 6 -> 8 + for _ in range(6): + rate *= 2 # sustained >10% windowed gain each tick + target, now = _run_tick(c, rate, now) + targets.append(target) + assert targets == sorted(targets) # monotonic non-decreasing + assert targets[-1] == 16 # climbed to the cap + assert max(targets) <= 16 + + +def test_reverts_on_plateau() -> None: + c = AdaptiveWorkerController(cap=16, tick=TICK, window=WINDOW) + now = 0.0 + c.update(now) + rate = 1_000_000.0 + now = _fill_window(c, rate, now) + assert c.update(now) == 8 # first increase 6 -> 8 + # hold the rate flat: <10% gain -> the increase is reverted and held + target, now = _run_tick(c, rate, now) + assert target == 6 + + +def test_halves_and_cooldown_on_error_burst() -> None: + c = AdaptiveWorkerController(cap=16, tick=TICK, window=WINDOW) + now = 0.0 + c.update(now) + now = _fill_window(c, 1_000_000.0, now) + for _ in range(3): + c.record_error(now) + assert c.update(now) == 3 # 3 errors -> multiplicative decrease 6 // 2 + + # cooldown tick holds even with improving throughput + target, now = _run_tick(c, 2_000_000.0, now) + assert target == 3 + + # cooldown consumed -> climbing resumes + target, now = _run_tick(c, 4_000_000.0, now) + assert target == 5 + + +def test_clamps_to_bounds_and_never_exceeds_cap() -> None: + c = AdaptiveWorkerController(cap=4, tick=TICK, window=WINDOW) + assert c.target == 4 # already at cap (min(6, 4)) + now = 0.0 + c.update(now) + rate = 1_000_000.0 + now = _fill_window(c, rate, now) + seen = [c.update(now)] + for _ in range(10): + rate *= 2 + target, now = _run_tick(c, rate, now) + seen.append(target) + assert all(2 <= s <= 4 for s in seen) + assert max(seen) == 4 # never exceeds the cap + + +def test_error_burst_floors_at_min() -> None: + c = AdaptiveWorkerController(cap=16, start=2, tick=TICK, window=WINDOW) + now = 0.0 + c.update(now) + now = _fill_window(c, 1_000_000.0, now) + for _ in range(5): + c.record_error(now) + assert c.update(now) == 2 # max(ADAPTIVE_MIN, 2 // 2) + + +def test_tail_guard_suppresses_increase() -> None: + # with fewer units of work than the target, the tick must hold instead of probing up + c = AdaptiveWorkerController(cap=16, tick=TICK, window=WINDOW) + now = 0.0 + c.update(now) # arm at target 6 + now = _fill_window(c, 1_000_000.0, now) + assert c.update(now, 3) == 6 # 3 < target 6 -> holds at 6 rather than climbing to 8 + + +def test_tail_guard_none_path_unchanged() -> None: + # without a tail count, the update still probes upward on the first evaluation + c = AdaptiveWorkerController(cap=16, tick=TICK, window=WINDOW) + now = 0.0 + c.update(now) + now = _fill_window(c, 1_000_000.0, now) + assert c.update(now) == 8 # unchanged: first evaluation 6 -> 8 + + +def test_tail_guard_still_halves_on_error_burst() -> None: + # the guard suppresses growth but never blocks an error-burst decrease + c = AdaptiveWorkerController(cap=16, tick=TICK, window=WINDOW) + now = 0.0 + c.update(now) + now = _fill_window(c, 1_000_000.0, now) + for _ in range(3): + c.record_error(now) + assert c.update(now, 2) == 3 # guard active (2 < 6) but 3 errors -> 6 // 2 + + +def test_plan_tail_parts_full_coverage_no_gaps() -> None: + size = 20 * 1024 * 1024 + parts = _plan_tail_parts(size, spare=8) + assert parts, "expected a split for a large segment with spare capacity" + assert parts[0][0] == 0 + assert parts[-1][1] == size - 1 + for (_, prev_end), (start, _) in zip(parts, parts[1:]): + assert start == prev_end + 1 # each part begins exactly where the previous one ended + covered = sum(end - start + 1 for start, end in parts) + assert covered == size + + +def test_plan_tail_parts_respects_spare_cap() -> None: + size = 100 * 1024 * 1024 # would want ~25 parts at 4MB, but spare caps it + parts = _plan_tail_parts(size, spare=3) + assert len(parts) <= 3 + assert parts[0][0] == 0 and parts[-1][1] == size - 1 + + +def test_plan_tail_parts_part_size_bounded() -> None: + size = 40 * 1024 * 1024 + parts = _plan_tail_parts(size, spare=16) + # never more parts than ~ceil(size / TAIL_BOOST_PART_SIZE) + assert len(parts) <= -(-size // TAIL_BOOST_PART_SIZE) + + +def test_plan_tail_parts_skips_small_or_no_spare() -> None: + assert _plan_tail_parts(TAIL_BOOST_MIN_SEGMENT_SIZE - 1, spare=8) == [] # too small + assert _plan_tail_parts(20 * 1024 * 1024, spare=1) == [] # not enough spare workers + assert _plan_tail_parts(20 * 1024 * 1024, spare=0) == [] + + +def test_tail_boost_engages_on_idle_capacity() -> None: + # engages at the tail: idle workers outnumber the remaining whole segments + assert _tail_boost_engages(remaining=4, pending=2, target=8) # spare 6 >= 4 + assert _tail_boost_engages(remaining=6, pending=6, target=16) # spare 10 >= 6 + + +def test_tail_boost_engages_when_remaining_is_a_full_stride() -> None: + # remaining sits at 6 (one worker stride) while the target has climbed to 16, so 14 workers + # are idle. idle capacity far exceeds the remaining segments, so the boost must engage. + assert _tail_boost_engages(remaining=6, pending=2, target=16) + + +def test_tail_boost_does_not_engage_early_or_saturated() -> None: + assert not _tail_boost_engages(remaining=54, pending=2, target=6) # early: 54 > spare 4 + assert not _tail_boost_engages(remaining=1, pending=8, target=8) # saturated: spare 0 + assert not _tail_boost_engages(remaining=0, pending=2, target=8) # nothing left + assert not _tail_boost_engages(remaining=5, pending=10, target=8) # over-subscribed: spare < 0 + + +def test_has_range_header_detects_byte_range_slices() -> None: + # items carrying a Range header are byte-range slices; ranged-parallel must skip them + assert _has_range_header({"url": "u", "headers": {"Range": "bytes=0-99"}}) + assert _has_range_header({"url": "u", "headers": {"range": "bytes=0-99"}}) # case-insensitive + assert not _has_range_header({"url": "u", "headers": {"Authorization": "Bearer x"}}) + assert not _has_range_header({"url": "u", "headers": {}}) + assert not _has_range_header({"url": "u"}) + + +def test_download_context_adaptive_default_false() -> None: + ctx = DownloadContext(save_path=Path("x.mp4"), save_dir=Path("."), progress=partial(lambda **_: None)) + assert ctx.adaptive_workers is False + + +def test_requests_accepts_adaptive_kwarg() -> None: + sig = inspect.signature(requests) + assert "adaptive" in sig.parameters + assert sig.parameters["adaptive"].default is False diff --git a/tests/core/test_downloader_resume.py b/tests/core/test_downloader_resume.py new file mode 100644 index 00000000..3dca5766 --- /dev/null +++ b/tests/core/test_downloader_resume.py @@ -0,0 +1,229 @@ +"""Regression tests for downloader resume/completeness robustness. + +Each test drives the real ``download()`` generator against a local HTTP server +that injects a specific fault the fault-injection benchmark surfaced: + +- truncated segment (clean FIN mid-body) must not finalize as complete +- an interrupted download must resume from the bytes actually written, not a + pre-allocated full size +- an oversized/stale resume that provokes a 416 must restart clean +- byte-range-slice segments (DASH SegmentBase / HLS EXT-X-BYTERANGE style) must + never be range-probed or tail-boosted: that would clobber their Range header + and fetch the whole resource into the segment file + +All localhost + small payloads, so they run in the default (not live) suite. +""" + +from __future__ import annotations + +import importlib +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Optional + +import pytest +from requests import Session + +# the downloaders package re-exports a `requests` function that shadows this submodule +# attribute, so fetch the real module object from sys.modules via importlib +dl = importlib.import_module("unshackle.core.downloaders.requests") + + +def _consume(gen) -> list: + return list(gen) + + +class _FaultServer: + """Serves a fixed body with per-request fault injection and Range support.""" + + def __init__(self, body: bytes) -> None: + self.body = body + self.attempts = 0 # GETs seen for the single resource + self.ranges: list[Optional[str]] = [] # Range header per GET, in order + self.truncate_first = 0 # send only half the body + close, for the first N attempts + self._lock = threading.Lock() + server = self + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *_a: object) -> None: # silence + pass + + def do_GET(self) -> None: # noqa: N802 (stdlib name) + with server._lock: + server.attempts += 1 + attempt = server.attempts + rng = self.headers.get("Range") + server.ranges.append(rng) + + total = len(server.body) + start = 0 + if rng and rng.startswith("bytes="): + start = int(rng.removeprefix("bytes=").split("-", 1)[0]) + + if start >= total: # unsatisfiable: real CDNs answer 416 + self.send_response(416) + self.send_header("Content-Range", f"bytes */{total}") + self.end_headers() + return + + chunk = server.body[start:] + partial = attempt <= server.truncate_first + declared = len(chunk) # honest Content-Length... + sent = chunk[: declared // 2] if partial else chunk # ...but under-deliver the body + + if rng: + self.send_response(206) + self.send_header("Content-Range", f"bytes {start}-{total - 1}/{total}") + else: + self.send_response(200) + self.send_header("Accept-Ranges", "bytes") + self.send_header("Content-Length", str(declared)) + self.end_headers() + self.wfile.write(sent) # short write + return => connection closes (clean FIN) + + self._httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self._thread = threading.Thread(target=self._httpd.serve_forever, daemon=True) + + def __enter__(self) -> "_FaultServer": + self._thread.start() + return self + + def __exit__(self, *_exc: object) -> None: + self._httpd.shutdown() + self._httpd.server_close() + + @property + def url(self) -> str: + host, port = self._httpd.server_address + return f"http://{host}:{port}/file" + + +@pytest.fixture(autouse=True) +def _fast_retries(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(dl, "RETRY_WAIT", 0) + + +def test_truncated_segment_is_not_finalized(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """A segment cut short mid-body must raise, never finalize a truncated file.""" + monkeypatch.setattr(dl, "MAX_ATTEMPTS", 2) + body = b"A" * 2000 + save = tmp_path / "seg.bin" + with _FaultServer(body) as srv: + srv.truncate_first = 99 # every attempt under-delivers + with pytest.raises(Exception): + _consume(dl.download(url=srv.url, save_path=save, session=Session(), segmented=True)) + assert not save.exists(), "truncated segment must not be finalized as complete" + + +def test_interrupted_download_resumes_from_bytes_written(tmp_path: Path) -> None: + """After a mid-body interruption the retry must Range-resume from the real + byte count (not a pre-allocated full size) and produce the correct file.""" + body = b"".join(bytes([i % 256]) for i in range(4000)) + save = tmp_path / "file.bin" + with _FaultServer(body) as srv: + srv.truncate_first = 1 # first attempt delivers half then closes; retry completes + _consume(dl.download(url=srv.url, save_path=save, session=Session())) + assert save.read_bytes() == body + # attempt 1 had no Range; attempt 2 must resume partway in, not from 0 or past the end + assert srv.ranges[0] is None + resume_start = int(srv.ranges[1].removeprefix("bytes=").split("-", 1)[0]) + assert 0 < resume_start < len(body), f"resume offset {resume_start} not derived from bytes written" + + +def test_resumes_from_preexisting_dev_partial(tmp_path: Path) -> None: + """A .!dev left on disk by a killed process (Ctrl+C, outage, crash) must resume on the + next run via a Range request from the partial's size, not restart from zero.""" + body = b"".join(bytes([i % 256]) for i in range(4000)) + save = tmp_path / "file.bin" + tmp = save.with_name("file.bin.!dev") + tmp.write_bytes(body[:1500]) # correct prefix the interrupted run flushed to disk + with _FaultServer(body) as srv: + _consume(dl.download(url=srv.url, save_path=save, session=Session())) + assert save.read_bytes() == body + assert not tmp.exists() + assert srv.ranges[0] == "bytes=1500-", f"expected resume from 1500, got {srv.ranges[0]}" + + +def test_byte_range_slice_segments_never_probed_or_boosted(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Adaptive tail boost must skip segments whose per-item headers carry a Range. + + Such segments are byte-range slices of one larger resource; the boost probe and + part-mode download would overwrite that Range and pull the entire resource into + the segment file. A server-side barrier clusters the first six completions so the + main loop reliably sees spare workers while the two tail slices are still queued, + the exact window where an unguarded boost would corrupt output. + """ + seg = 2 * 1024 * 1024 + nseg = 8 + body = b"".join(bytes([i]) * seg for i in range(nseg)) # 16 MiB, distinct byte per slice + + barrier_count = [0] + barrier_lock = threading.Lock() + barrier = threading.Event() + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *_a: object) -> None: + pass + + def do_GET(self) -> None: # noqa: N802 (stdlib name) + rng = self.headers.get("Range") or "" + with barrier_lock: + barrier_count[0] += 1 + if barrier_count[0] >= 6: + barrier.set() + barrier.wait(5) # cluster the first six completions; auto-release as a safety net + start, _, end_s = rng.removeprefix("bytes=").partition("-") + start = int(start or 0) + end = int(end_s) if end_s else len(body) - 1 + chunk = body[start : end + 1] + self.send_response(206 if rng else 200) + if rng: + self.send_header("Content-Range", f"bytes {start}-{end}/{len(body)}") + self.send_header("Accept-Ranges", "bytes") + self.send_header("Content-Length", str(len(chunk))) + self.end_headers() + self.wfile.write(chunk) + + httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=httpd.serve_forever, daemon=True).start() + host, port = httpd.server_address + + probed: list[str] = [] + real_probe = dl._probe_ranged + monkeypatch.setattr(dl, "_probe_ranged", lambda url, session, **kw: probed.append(url) or real_probe(url, session, **kw)) + + url = f"http://{host}:{port}/media.mp4" + urls = [{"url": url, "headers": {"Range": f"bytes={i * seg}-{(i + 1) * seg - 1}"}} for i in range(nseg)] + try: + _consume( + dl.requests( + urls=urls, + output_dir=tmp_path, + filename="seg_{i:02}{ext}", + session=Session(), + max_workers=16, + adaptive=True, + ) + ) + finally: + httpd.shutdown() + httpd.server_close() + + assert probed == [], "byte-range slice segments must never be range-probed" + for i in range(nseg): + data = (tmp_path / f"seg_{i:02}.mp4").read_bytes() + assert data == body[i * seg : (i + 1) * seg], f"segment {i} corrupted: got {len(data)} bytes" + + +def test_oversized_resume_416_restarts_clean(tmp_path: Path) -> None: + """A stale .!dev larger than the resource provokes a 416; the downloader must + discard it and re-download cleanly rather than burn retries or corrupt.""" + body = b"Z" * 1500 + save = tmp_path / "file.bin" + tmp = save.with_name(f"{save.name}.!dev") + tmp.write_bytes(b"X" * 4096) # poisoned partial, bigger than the server's body + with _FaultServer(body) as srv: + _consume(dl.download(url=srv.url, save_path=save, session=Session())) + assert save.read_bytes() == body + assert srv.ranges[0] is not None and "4096" in srv.ranges[0] # first request over-ranged → 416 diff --git a/tests/remote/unit/test_errors.py b/tests/remote/unit/test_errors.py index b5906431..51f37534 100644 --- a/tests/remote/unit/test_errors.py +++ b/tests/remote/unit/test_errors.py @@ -32,6 +32,7 @@ def test_api_error_default_http_status_per_code() -> None: APIErrorCode.JOB_NOT_FOUND: 404, APIErrorCode.SESSION_NOT_FOUND: 404, APIErrorCode.TRACK_NOT_FOUND: 404, + APIErrorCode.CONFLICT: 409, APIErrorCode.RATE_LIMITED: 429, APIErrorCode.INTERNAL_ERROR: 500, APIErrorCode.DOWNLOAD_ERROR: 500, diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 024e5a5e..e310554c 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -665,6 +665,15 @@ class dl: default=None, help="Max workers/threads to download with per-track. Default depends on the downloader.", ) + @click.option( + "--adaptive-workers", + is_flag=True, + default=False, + help=( + "Dynamically scale per-track segment workers up to the --workers cap (or the built-in cap " + "when --workers is unset) based on measured CDN throughput and errors. Off by default." + ), + ) @click.option("--downloads", type=int, default=1, help="Amount of tracks to download concurrently.") @click.option( "-o", @@ -1201,6 +1210,7 @@ class dl: no_source: bool, no_mux: bool, workers: Optional[int], + adaptive_workers: bool, downloads: int, worst: bool, best_available: bool, @@ -1741,6 +1751,7 @@ class dl: ), cdm=self.cdm, max_workers=workers, + adaptive_workers=adaptive_workers, progress=progress_call, ) for song, track, progress_call in music_items @@ -2850,6 +2861,7 @@ class dl: ), cdm=self.cdm, max_workers=workers, + adaptive_workers=adaptive_workers, progress=tracks_progress_callables[i], ) # DRM-free and HLS-ClearKey tracks never reach prepare_drm, so export here. diff --git a/unshackle/core/downloaders/requests.py b/unshackle/core/downloaders/requests.py index 56b51335..fdaa3b13 100644 --- a/unshackle/core/downloaders/requests.py +++ b/unshackle/core/downloaders/requests.py @@ -3,6 +3,7 @@ import os import statistics import threading import time +from collections import deque from concurrent.futures import FIRST_COMPLETED, wait from concurrent.futures.thread import ThreadPoolExecutor from http.cookiejar import CookieJar @@ -38,6 +39,13 @@ SPEED_ROLLING_WINDOW = 10 # seconds of history to keep for speed calculation RANGE_PARALLEL_MIN_SIZE = 64 * 1024 * 1024 RANGE_PARALLEL_PART_SIZE = 16 * 1024 * 1024 +# Tail-end boost (opt-in, adaptive only): when only a few segments remain and workers +# would otherwise idle, split each remaining segment into intra-segment range parts so +# aggregate throughput holds to the end instead of collapsing to (few)x(per-connection). +TAIL_BOOST_MAX_PER_CYCLE = 4 # cap on segments probed/split per drain cycle so blocking probes can't stall the loop +TAIL_BOOST_MIN_SEGMENT_SIZE = 8 * 1024 * 1024 # don't probe+split segments smaller than this +TAIL_BOOST_PART_SIZE = 4 * 1024 * 1024 # target size per tail part (smaller than the whole-file part size) + def _adaptive_chunk_size(content_length: int) -> int: """Pick chunk size based on content length. Benchmarked sweet spot: 512KB-4MB.""" @@ -46,6 +54,140 @@ def _adaptive_chunk_size(content_length: int) -> int: return min(MAX_CHUNK, max(MIN_CHUNK, content_length // 4)) +# Adaptive worker controller (opt-in): AIMD hill-climb over segment concurrency +ADAPTIVE_TICK = 4.0 # seconds between target re-evaluations +ADAPTIVE_START = 6 # initial worker target, clamped to cap +ADAPTIVE_STEP = 2 # additive increase per tick +ADAPTIVE_MIN = 2 # never drop below this many workers +ADAPTIVE_ERROR_BURST = 3 # errors within a tick that trigger multiplicative decrease +ADAPTIVE_PLATEAU_GAIN = 1.10 # min speed ratio to justify keeping an increase + + +class AdaptiveWorkerController: + """CDN-aware segment concurrency governor. + + Pure and synchronous (no threads or sockets), so the AIMD policy is unit-testable. + ``update`` is evaluated once per ``tick`` and only after a full ``window`` of + throughput samples exists: it ramps the worker target up while the CDN keeps + rewarding extra concurrency, reverts the last increase on a plateau, and halves + (with a one-tick cooldown) on an error burst. Target is always clamped to + ``[ADAPTIVE_MIN, cap]``. + """ + + def __init__( + self, cap: int, start: Optional[int] = None, tick: float = ADAPTIVE_TICK, window: float = SPEED_ROLLING_WINDOW + ) -> None: + self.cap = max(ADAPTIVE_MIN, cap) + base = ADAPTIVE_START if start is None else start + self.target = max(ADAPTIVE_MIN, min(base, self.cap)) + self.tick = tick + self.window = window + self._samples: deque[tuple[float, int]] = deque() + self._first_sample: Optional[float] = None # first-ever sample time (survives pruning) + self._errors = 0 # errors observed since the last tick evaluation + self._last_tick: Optional[float] = None + self._last_action: Optional[str] = None # "increase" once we have probed upward + self._speed_before_increase = 0.0 + self._cooldown = False + + def _prune(self, now: float) -> None: + cutoff = now - self.window + while self._samples and self._samples[0][0] < cutoff: + self._samples.popleft() + + def record_bytes(self, n: int, now: float) -> None: + if self._first_sample is None: + self._first_sample = now + self._samples.append((now, n)) + self._prune(now) + + def record_error(self, now: float) -> None: + self._errors += 1 + + def speed(self, now: float) -> float: + """Rolling bytes/sec over the window.""" + self._prune(now) + if not self._samples: + return 0.0 + span = now - self._samples[0][0] + if span <= 0: + return 0.0 + return sum(n for _, n in self._samples) / span + + def _full_window(self, now: float) -> bool: + # true once we have been observing for at least a full window; measured from the + # first-ever sample so continuous pruning (which keeps span just under window) does + # not perpetually starve the controller of a "full" window in real timing + return self._first_sample is not None and (now - self._first_sample) >= self.window + + def update(self, now: float, inflight_plus_remaining: Optional[int] = None) -> int: + """Evaluate the policy at most once per tick; return the current target. + + ``inflight_plus_remaining`` (when given) is the count of segments still in flight or + not yet started. If it is below the target there is not enough work to saturate the + target, so the low measured speed reflects starvation rather than a plateau, and the + probe/revert step is skipped this tick. Error-burst decreases still apply. Default ``None`` keeps the + original behaviour for existing callers. + """ + if self._last_tick is None: + self._last_tick = now + return self.target + if now - self._last_tick < self.tick: + return self.target + self._last_tick = now + + # hold until a full window of throughput history has accumulated + if not self._full_window(now): + self._errors = 0 + return self.target + + errors = self._errors + self._errors = 0 + speed_now = self.speed(now) + old = self.target + + if self._cooldown: + # one tick of quiet after a decrease before probing again + self._cooldown = False + self._last_action = None + return self.target + + reason: Optional[str] = None + if errors >= ADAPTIVE_ERROR_BURST: + self.target = max(ADAPTIVE_MIN, self.target // 2) + self._cooldown = True + self._last_action = None + reason = "error_burst" + elif inflight_plus_remaining is not None and inflight_plus_remaining < self.target: + # tail guard: too little work to saturate the target, so a low measured speed + # here reflects starvation rather than a plateau, so hold and skip the probe/revert this tick + pass + elif self._last_action == "increase" and speed_now < ADAPTIVE_PLATEAU_GAIN * self._speed_before_increase: + # last probe upward did not pay off, so revert it and hold + self.target = max(ADAPTIVE_MIN, self.target - ADAPTIVE_STEP) + self._last_action = None + reason = "plateau" + elif self.target < self.cap: + self._speed_before_increase = speed_now + self.target = min(self.cap, self.target + ADAPTIVE_STEP) + self._last_action = "increase" + reason = "increase" + else: + self._last_action = None + + self.target = max(ADAPTIVE_MIN, min(self.cap, self.target)) + if self.target != old: + debug_logger = get_debug_logger() + if debug_logger: + debug_logger.log( + level="DEBUG", + operation="adaptive_workers", + message=f"worker target {old} -> {self.target}", + context={"old": old, "new": self.target, "speed_bps": round(speed_now), "reason": reason}, + ) + return self.target + + def _is_requests_session(session: Any) -> bool: """Check if the session is a standard requests.Session (supports resp.raw).""" return isinstance(session, Session) @@ -63,6 +205,8 @@ def _probe_ranged(url: str, session: Any, **kwargs: Any) -> tuple[int, bool]: rest = {k: v for k, v in kwargs.items() if k != "headers"} if _is_rnet_session(session): rest.setdefault("read_timeout", READ_TIMEOUT) + # let download() own segment retries, with no nested session-level retry loop + rest.setdefault("max_retries", 0) else: rest.setdefault("timeout", (CONNECT_TIMEOUT, READ_TIMEOUT)) try: @@ -85,6 +229,48 @@ def _probe_ranged(url: str, session: Any, **kwargs: Any) -> tuple[int, bool]: pass +def _has_range_header(item: dict[str, Any]) -> bool: + """True when a URL item's per-request headers already carry a Range. + + Such an item is itself a byte-range slice of a larger resource (DASH SegmentBase, + HLS EXT-X-BYTERANGE); probing or part-mode would clobber that Range and fetch the + wrong bytes, so ranged-parallel strategies must skip these items entirely. + """ + return any(k.lower() == "range" for k in (item.get("headers") or {})) + + +def _plan_tail_parts(size: int, spare: int) -> list[tuple[int, int]]: + """Plan intra-segment range parts for a tail segment. + + Returns inclusive ``[start, end]`` byte ranges covering ``[0, size)`` with no gaps or + overlaps, split into at most ``spare`` parts of ~``TAIL_BOOST_PART_SIZE``. Returns an + empty list when the segment is smaller than ``TAIL_BOOST_MIN_SEGMENT_SIZE`` (not worth a + probe+split) or there are fewer than two spare workers. Pure, with no I/O, so testable. + """ + if size < TAIL_BOOST_MIN_SEGMENT_SIZE or spare < 2: + return [] + 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 + ] + + +def _tail_boost_engages(remaining: int, pending: int, target: int) -> bool: + """True when idle workers outnumber the remaining whole segments. + + That is the tail condition where running one worker per remaining segment would leave + workers idle, so splitting those segments into range parts keeps throughput up. Gating on + idle capacity (not a fixed segment count) is what lets the boost fire despite ``remaining`` + draining in strides of the worker target. Pure, so the gate is unit-testable. + """ + spare = target - pending + return remaining > 0 and spare >= 2 and remaining <= spare + + def _dispatch_parts( url: str, save_path: Path, @@ -110,6 +296,9 @@ def _dispatch_parts( f.truncate(total_size) events: Queue[dict[str, Any]] = Queue() + # local to this ranged download: a failed part must not poison the process-global + # DOWNLOAD_CANCELLED (which would kill the sequential fallback and sibling tracks) + abort = threading.Event() def _worker(start: int, end: int) -> None: for ev in download( @@ -118,6 +307,7 @@ def _dispatch_parts( session=session, part_offset=start, part_end=end, + abort=abort, **kwargs, ): events.put(ev) @@ -158,7 +348,7 @@ def _dispatch_parts( exc = fut.exception() if exc: worker_error = True - DOWNLOAD_CANCELLED.set() + abort.set() raise exc advance = 0 @@ -194,6 +384,8 @@ def download( part_offset: Optional[int] = None, part_end: Optional[int] = None, claimed: Optional[Callable[[], bool]] = None, + abort: Optional[threading.Event] = None, + on_retry: Optional[Callable[[Exception], None]] = None, **kwargs: Any, ) -> Generator[dict[str, Any], None, None]: """ @@ -224,6 +416,12 @@ 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 local cancel signal. When set, the download stops like a cancel + (keeps its partial, returns silently, does not raise). Used by ranged-parallel + downloads to abort sibling parts without touching the process-global cancel. + on_retry: Optional callback invoked with the raised exception before each retry + wait (not on cancel or final exhaustion). Used to feed CDN error signals to + the adaptive worker controller. Callback errors are swallowed. 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"}`. @@ -268,6 +466,8 @@ def download( request_kwargs = dict(kwargs) if use_rnet: request_kwargs.setdefault("read_timeout", READ_TIMEOUT) + # own segment retries here, suppressing RnetSession's nested retry loop + request_kwargs.setdefault("max_retries", 0) else: request_kwargs.setdefault("timeout", (CONNECT_TIMEOUT, READ_TIMEOUT)) if part_mode: @@ -280,6 +480,19 @@ def download( 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: + # 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: + stream.close() + except Exception: + pass + tmp_file.unlink(missing_ok=True) + resume_offset = 0 + written = 0 + continue + stream.raise_for_status() resumed = (not part_mode) and resume_offset > 0 and stream.status_code == 206 @@ -314,12 +527,14 @@ def download( else: file_mode = "ab" if resumed else "wb" file_buffering = 1_048_576 + # non-part downloads write sequentially and must NOT pre-allocate: the resume + # offset is derived from the tmp file's size, so preallocating to content_length + # would make an interrupted write look fully downloaded and poison the resume + # (unsatisfiable Range on retry, silent zero-padded corruption). part_mode still + # pre-truncates in _dispatch_parts because its workers seek to fixed offsets. 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 @@ -334,7 +549,7 @@ def download( _bytes_since_yield = 0 emit_progress = (not segmented) or part_mode for chunk in chunks: - if DOWNLOAD_CANCELLED.is_set(): + if DOWNLOAD_CANCELLED.is_set() or (abort is not None and abort.is_set()): break _write(chunk) download_size = len(chunk) @@ -362,18 +577,21 @@ 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 + aborted = abort is not None and abort.is_set() + if (DOWNLOAD_CANCELLED.is_set() or aborted) and (not content_length or written < content_length): + # cancelled/aborted mid-stream: keep the partial for resume, never finalize, + # and skip the size checks below (a short read here is expected, not an error) return if part_mode: expected = part_end - part_offset + 1 if written < expected: raise IOError(f"Failed to read part {part_offset}-{part_end}: got {written}/{expected}") - elif not segmented and content_length and written < content_length: + 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 + # finalize as complete and corrupt the merge. content_length==0 (chunked/gzip) + # is unknown-length and correctly skips the check. raise IOError(f"Failed to read {content_length} bytes from the track URI.") if not part_mode: @@ -382,15 +600,26 @@ def download( if segmented: yield dict(advance=1) break - except Exception: + except Exception as exc: try: stream.close() except Exception: pass - if DOWNLOAD_CANCELLED.is_set() or attempts == MAX_ATTEMPTS: - if part_mode and not DOWNLOAD_CANCELLED.is_set(): + cancelled = DOWNLOAD_CANCELLED.is_set() or (abort is not None and abort.is_set()) + if cancelled or attempts == MAX_ATTEMPTS: + # a hedged loser whose file was already delivered must not kill the batch + if claimed is not None and claimed(): + return + # cancel/abort is not an error, but a genuine retry-exhaustion is, so surface it + # so the caller fails fast instead of hitting a missing segment at merge time + if not cancelled: raise return + if on_retry is not None: + try: + on_retry(exc) + except Exception: + pass if not part_mode: resume_offset = tmp_file.stat().st_size if tmp_file.exists() else 0 time.sleep(RETRY_WAIT) @@ -406,6 +635,7 @@ def requests( proxy: Optional[str] = None, max_workers: Optional[int] = None, session: Optional[Any] = None, + adaptive: bool = False, ) -> Generator[dict[str, Any], None, None]: """ Download files with optimized I/O and adaptive chunk sizing. @@ -441,6 +671,11 @@ def requests( session: An optional requests.Session or RnetSession to use. If provided, it will be used directly (preserving TLS fingerprinting). If None, a new requests.Session with HTTPAdapter connection pooling will be created. + adaptive: Opt-in CDN-aware dynamic segment concurrency for segmented (multi-URL) + downloads. When True the per-track worker count starts moderate and ramps up + or backs off based on measured throughput and CDN errors, capped at + max_workers. When False (default) all segments are submitted upfront using a + fixed worker count, matching the non-adaptive behaviour exactly. """ if not urls: raise ValueError("urls must be provided and not empty") @@ -522,6 +757,7 @@ def requests( "max_workers": max_workers, "has_proxy": bool(proxy), "session_type": type(session).__name__, + "adaptive": adaptive, }, ) @@ -531,7 +767,7 @@ def requests( url_item = urls[0] try: ranged_used = False - if max_workers > 1: + if max_workers > 1 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: @@ -581,6 +817,10 @@ def requests( seg_durations: list[float] = [] hedged: set[int] = set() + controller = AdaptiveWorkerController(cap=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 + def _download_worker(index: int, url_item: dict[str, Any], hedge: bool = False) -> None: item = dict(url_item) save_path = item.pop("save_path") @@ -596,12 +836,16 @@ def requests( # 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() + # under seg_lock: the hedge scan iterates seg_start.items() on the main + # thread, and a first-time key insert during that iteration raises + with seg_lock: + seg_start[index] = time.time() for event in download( session=session, segmented=segmented_batch, save_path=target, claimed=lambda: index in seg_done, + on_retry=on_retry_cb, **item, ): if "file_downloaded" in event: @@ -616,13 +860,159 @@ def requests( event = dict(event, file_downloaded=save_path) event_queue.put(event) - futures = [pool.submit(_download_worker, i, url) for i, url in enumerate(urls)] - pending = set(futures) + # not-yet-submitted segments; adaptive mode meters submission against the controller + # target, while fixed mode (adaptive=False) submits every segment upfront + remaining: deque[tuple[int, dict[str, Any]]] = deque(enumerate(urls)) + pending: set = set() + + # tail boost bookkeeping (adaptive only): indices whose parts are in flight vs + # indices probed and left for the normal single-worker path, each decided once + tail_boosted: set[int] = set() + tail_skipped: set[int] = set() + + def _submit(count: int, only: Optional[set[int]] = None) -> None: + # only= submits just those indices (used in the tail window to release + # segments the boost declined while holding back boost candidates); with no + # `only`, submits the leading `count` segments in FIFO order. + if count <= 0 or not remaining: + return + picked: list[tuple[int, dict[str, Any]]] = [] + for item in remaining: + if count <= 0: + break + if only is not None and item[0] not in only: + continue + pending.add(pool.submit(_download_worker, item[0], item[1])) + picked.append(item) + count -= 1 + for item in picked: + remaining.remove(item) + + def _tail_part_worker( + index: int, + url: str, + part_target: Path, + save_path: Path, + start: int, + end: int, + size: int, + abort: threading.Event, + parts_left: list[int], + part_lock: threading.Lock, + req_kwargs: dict[str, Any], + ) -> None: + try: + # part_mode download writes [start, end] into the pre-truncated target and + # emits only byte-`advance`; swallow those, since tail segments report at segment + # granularity (advance=1) like every other segment, once on finalize below + for _ev in download( + url=url, + save_path=part_target, + session=session, + part_offset=start, + part_end=end, + abort=abort, + on_retry=on_retry_cb, + **req_kwargs, + ): + pass + except BaseException: + abort.set() # stop sibling parts of this segment; a real failure fails the batch + raise + with part_lock: + parts_left[0] -= 1 + finalize = parts_left[0] == 0 + if not finalize: + return + if DOWNLOAD_CANCELLED.is_set() or abort.is_set(): + return # cancelled/aborted: keep the partial target for the stray sweep, don't finalize + with seg_lock: + if index in seg_done: + part_target.unlink(missing_ok=True) + return + seg_done.add(index) + seg_durations.append(time.time() - seg_start[index]) + os.replace(part_target, save_path) + event_queue.put(dict(file_downloaded=save_path, written=size)) + event_queue.put(dict(advance=1)) + + def _maybe_tail_boost(target: int) -> None: + # Tail boost: near the end, split the few not-yet-started segments across idle + # workers with intra-segment range parts so aggregate throughput doesn't collapse + # to (few segments)x(per-connection speed). Only not-yet-started segments in + # `remaining` are boosted (never mid-flight ones) to avoid entangling the main + # loop; gated behind adaptive, though it could be decoupled from the controller. + # Engage on idle capacity so the boost fires despite `remaining` draining in strides + # of the worker target; probe at most TAIL_BOOST_MAX_PER_CYCLE per cycle so a run of + # 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 + boosted = 0 + for index, url_item in list(remaining): + if boosted >= TAIL_BOOST_MAX_PER_CYCLE: + break + spare = target - len(pending) + if spare < 2: + break + if index in tail_boosted or index in tail_skipped: + continue + item = dict(url_item) + save_path = item.pop("save_path") + if save_path.exists(): # finished in a previous run; let the normal fast-path handle it + tail_skipped.add(index) + continue + if _has_range_header(item): # byte-range slice; part-mode would clobber its Range + tail_skipped.add(index) + continue + req_kwargs = {k: v for k, v in item.items() if k != "url"} + size, supports_ranges = _probe_ranged(item["url"], session, **req_kwargs) + parts = _plan_tail_parts(size, spare) if supports_ranges else [] + if not parts: + tail_skipped.add(index) # no ranges / too small / no capacity -> normal path + continue + try: + remaining.remove((index, url_item)) + except ValueError: + tail_skipped.add(index) + continue + # distinct .tp.!dev target so an interrupted boost never collides with the + # normal worker's .p.!dev; the .!dev suffix keeps the stray visible to cleanup + part_target = save_path.with_name(f"{save_path.name}.tp.!dev") + with open(part_target, "wb") as f: + f.truncate(size) + seg_start[index] = time.time() + tail_boosted.add(index) + abort = threading.Event() + parts_left = [len(parts)] + part_lock = threading.Lock() + for start, end in parts: + pending.add( + pool.submit( + _tail_part_worker, + index, + item["url"], + part_target, + save_path, + start, + end, + size, + abort, + parts_left, + part_lock, + req_kwargs, + ) + ) + boosted += 1 + + if controller: + _submit(controller.update(time.time())) + else: + _submit(len(remaining)) pending_advance = 0 try: - while pending: + while pending or remaining: # Drain queued events — batch advances, track bytes for speed while True: try: @@ -638,6 +1028,8 @@ def requests( written = event.get("written") if written: total_bytes += written + if controller: + controller.record_bytes(written, time.time()) # Pass through other events (file_downloaded, total, etc.) yield event @@ -646,25 +1038,50 @@ def requests( yield dict(advance=pending_advance) pending_advance = 0 - # Yield speed every 0.5s (throttled to avoid spamming Rich) now = time.time() - if now - last_speed_report > 0.5 and total_bytes > 0: - elapsed = now - start_time - if elapsed > 0: - download_speed = math.ceil(total_bytes / elapsed) - yield dict(downloaded=f"{filesize.decimal(download_speed)}/s") - last_speed_report = now - # hedge stuck segments once spare workers exist (pending < max_workers) - if len(pending) < max_workers and seg_durations and not DOWNLOAD_CANCELLED.is_set(): + # Adaptive: re-evaluate the worker target and top up in-flight futures + target = max_workers + if controller: + target = controller.update(now, len(pending) + len(remaining)) + _maybe_tail_boost(target) # spend spare workers on the tail before topping up + if len(pending) < target: + # In the final stride (unstarted segments <= one worker target), hold back + # boost candidates so idle workers accumulate for a range-split; only + # release segments the boost already declined (too small / no range) so + # nothing stalls. Otherwise top up in FIFO order as usual. + if remaining and len(remaining) <= target: + _submit(target - len(pending), only=tail_skipped) + else: + _submit(target - len(pending)) + + # Yield speed every 0.5s (throttled to avoid spamming Rich) + if now - last_speed_report > 0.5: + if controller: + rolling = controller.speed(now) + if rolling > 0: + yield dict(downloaded=f"{filesize.decimal(math.ceil(rolling))}/s") + last_speed_report = now + elif total_bytes > 0: + elapsed = now - start_time + if elapsed > 0: + download_speed = math.ceil(total_bytes / elapsed) + yield dict(downloaded=f"{filesize.decimal(download_speed)}/s") + last_speed_report = now + + # hedge stuck segments once spare workers exist within the active target + if len(pending) < target and seg_durations and not DOWNLOAD_CANCELLED.is_set(): threshold = max(HEDGE_FACTOR * statistics.median(seg_durations), HEDGE_MIN_WAIT) with seg_lock: stuck = [ i for i, started in seg_start.items() - if i not in seg_done and i not in hedged and now - started > threshold + if i not in seg_done + and i not in hedged + and i not in tail_boosted # boosted segments are already range-split + and now - started > threshold ] - for i in stuck[: max_workers - len(pending)]: + for i in stuck[: target - len(pending)]: hedged.add(i) pending.add(pool.submit(_download_worker, i, urls[i], True)) diff --git a/unshackle/core/manifests/dash.py b/unshackle/core/manifests/dash.py index 1d15cf5b..96c7f361 100644 --- a/unshackle/core/manifests/dash.py +++ b/unshackle/core/manifests/dash.py @@ -281,6 +281,7 @@ class DASH: progress = ctx.progress proxy = ctx.proxy max_workers = ctx.max_workers + adaptive = ctx.adaptive_workers license_widevine = ctx.license_widevine cdm = ctx.cdm @@ -458,6 +459,7 @@ class DASH: proxy=proxy, max_workers=max_workers, session=session, + adaptive=adaptive, ) log_event( diff --git a/unshackle/core/manifests/hls.py b/unshackle/core/manifests/hls.py index 2b1c2eec..fddabca1 100644 --- a/unshackle/core/manifests/hls.py +++ b/unshackle/core/manifests/hls.py @@ -523,6 +523,7 @@ class HLS: progress = ctx.progress proxy = ctx.proxy max_workers = ctx.max_workers + adaptive = ctx.adaptive_workers license_widevine = ctx.license_widevine cdm = ctx.cdm @@ -672,6 +673,7 @@ class HLS: proxy=proxy, max_workers=max_workers, session=session, + adaptive=adaptive, ) log_event( diff --git a/unshackle/core/manifests/ism.py b/unshackle/core/manifests/ism.py index 0c8913ef..2fb45b45 100644 --- a/unshackle/core/manifests/ism.py +++ b/unshackle/core/manifests/ism.py @@ -370,6 +370,7 @@ class ISM: progress = ctx.progress proxy = ctx.proxy max_workers = ctx.max_workers + adaptive = ctx.adaptive_workers license_widevine = ctx.license_widevine cdm = ctx.cdm @@ -411,6 +412,7 @@ class ISM: proxy=proxy, max_workers=max_workers, session=session, + adaptive=adaptive, ) log_event( diff --git a/unshackle/core/service.py b/unshackle/core/service.py index 109cd77e..f124c950 100644 --- a/unshackle/core/service.py +++ b/unshackle/core/service.py @@ -4,7 +4,7 @@ from collections.abc import Callable, Generator from dataclasses import dataclass, field from http.cookiejar import CookieJar from pathlib import Path -from typing import TYPE_CHECKING, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union from urllib.parse import urlparse, urlunparse if TYPE_CHECKING: @@ -25,12 +25,55 @@ from unshackle.core.constants import AnyTrack from unshackle.core.credential import Credential from unshackle.core.drm import DRM_T from unshackle.core.search_result import SearchResult +from unshackle.core.session import (BACKOFF_FACTOR, CONNECT_TIMEOUT, MAX_BACKOFF, MAX_RETRIES, POOL_MAX_SIZE, + READ_TIMEOUT, RETRY_METHODS, STATUS_FORCELIST) from unshackle.core.title_cacher import TitleCacher, get_account_hash, get_region_from_proxy from unshackle.core.titles import Title_T, Titles_T, remap_titles from unshackle.core.tracks import Chapters, Tracks from unshackle.core.tracks.video import Video from unshackle.core.utils.ip_info import get_ip_info +# Default (connect, read) timeout for the requests path, mirroring RnetSession's +# connect_timeout / read_timeout construction defaults. A per-request timeout= wins. +DEFAULT_TIMEOUT = (CONNECT_TIMEOUT, READ_TIMEOUT) + + +class TimeoutSession(requests.Session): + """requests.Session applying DEFAULT_TIMEOUT when the caller passes none. + + requests has no native default timeout; without one a stalled connect or + read hangs forever. RnetSession bounds every request via its client's + connect_timeout/read_timeout, so this mirrors that on the requests path. + A per-request non-None ``timeout=`` wins; an explicit ``timeout=None`` is + still replaced with the default by :class:`TimeoutHTTPAdapter` on the + mounted adapters, so there is no unbounded read, matching RnetSession where + the client-level timeouts always apply. Pass a large timeout instead. + """ + + def request(self, *args: Any, **kwargs: Any) -> requests.Response: + kwargs.setdefault("timeout", DEFAULT_TIMEOUT) + return super().request(*args, **kwargs) + + +class TimeoutHTTPAdapter(HTTPAdapter): + """HTTPAdapter applying DEFAULT_TIMEOUT when the caller passes none. + + Backstops :class:`TimeoutSession` for the ``session.send(prepared)`` path, + which bypasses ``Session.request``; RnetSession bounds those too via its + client, so this keeps parity. A per-request non-None ``timeout=`` wins; + ``None`` (unset, or explicitly passed) gets the default, because the adapter + cannot distinguish the two, and rnet has no unbounded mode either. + """ + + def __init__(self, *args: Any, timeout: Any = DEFAULT_TIMEOUT, **kwargs: Any) -> None: + self.default_timeout = timeout + super().__init__(*args, **kwargs) + + def send(self, request: Any, **kwargs: Any) -> Any: + if kwargs.get("timeout") is None: + kwargs["timeout"] = self.default_timeout + return super().send(request, **kwargs) + @dataclass class TrackRequest: @@ -322,12 +365,27 @@ class Service(metaclass=ABCMeta): from config, cookies, retry handler, and a proxy if available. :returns: Prepared Python-requests Session """ - session = requests.Session() + session = TimeoutSession() session.headers.update(config.headers) + # Retry / pool policy mirrors RnetSession via the shared constants in session.py: + # same total, forcelist, backoff cap, allowed_methods (incl. POST, so license + # requests retry on both paths), and pool size. Accepted divergence: urllib3's + # backoff_jitter adds an absolute 0..N seconds whereas rnet jitters ±10% of the + # backoff, the only intentional difference, not worth hand-rolling a retry loop. session.mount( "https://", - HTTPAdapter( - max_retries=Retry(total=5, backoff_factor=0.2, status_forcelist=[429, 500, 502, 503, 504]), + TimeoutHTTPAdapter( + max_retries=Retry( + total=MAX_RETRIES, + backoff_factor=BACKOFF_FACTOR, + backoff_max=MAX_BACKOFF, + backoff_jitter=0.2, + status_forcelist=STATUS_FORCELIST, + allowed_methods=RETRY_METHODS, + respect_retry_after_header=True, + ), + pool_connections=POOL_MAX_SIZE, + pool_maxsize=POOL_MAX_SIZE, pool_block=True, ), ) diff --git a/unshackle/core/session.py b/unshackle/core/session.py index 66e0a7d0..71ad8b70 100644 --- a/unshackle/core/session.py +++ b/unshackle/core/session.py @@ -19,6 +19,24 @@ from requests.structures import CaseInsensitiveDict from unshackle.core.config import config +# --------------------------------------------------------------------------- +# Shared retry / pool / timeout policy: single source of truth +# --------------------------------------------------------------------------- +# rnet's only distinguishing role is TLS fingerprinting; retry, backoff, +# pooling, and timeouts must match the plain-requests path. Both RnetSession +# (below) and Service.get_session import these so the two can't drift. +# The unified downloader keeps its own copies (downloaders/requests.py) by design. + +MAX_RETRIES = 5 +BACKOFF_FACTOR = 0.2 +MAX_BACKOFF = 60.0 # shared backoff cap (rnet max_backoff == requests Retry backoff_max) +STATUS_FORCELIST = [429, 500, 502, 503, 504] +RETRY_METHODS = frozenset({"GET", "POST", "HEAD", "OPTIONS", "PUT", "DELETE", "TRACE"}) +POOL_MAX_SIZE = 16 # rnet pool_max_idle_per_host == requests pool_maxsize/pool_connections +CONNECT_TIMEOUT = 10 +READ_TIMEOUT = 30 + + # --------------------------------------------------------------------------- # Impersonate preset mapping — rnet uses named presets (no custom JA3/Akamai) # --------------------------------------------------------------------------- @@ -483,9 +501,9 @@ class RnetSession: def __init__( self, - max_retries: int = 5, - backoff_factor: float = 0.2, - max_backoff: float = 60.0, + max_retries: int = MAX_RETRIES, + backoff_factor: float = BACKOFF_FACTOR, + max_backoff: float = MAX_BACKOFF, status_forcelist: Optional[list[int]] = None, allowed_methods: Optional[set[str]] = None, catch_exceptions: Optional[tuple[type[Exception], ...]] = None, @@ -494,10 +512,11 @@ class RnetSession: self.max_retries = max_retries self.backoff_factor = backoff_factor self.max_backoff = max_backoff - self.status_forcelist = status_forcelist or [429, 500, 502, 503, 504] - self.allowed_methods = allowed_methods or {"GET", "POST", "HEAD", "OPTIONS", "PUT", "DELETE", "TRACE"} + self.status_forcelist = status_forcelist or list(STATUS_FORCELIST) + self.allowed_methods = allowed_methods or set(RETRY_METHODS) self.catch_exceptions = catch_exceptions or ( rnet.ConnectionError, + rnet.ConnectionResetError, # sibling of ConnectionError; mid-request resets are retryable rnet.TimeoutError, rnet.RequestError, ) @@ -511,9 +530,13 @@ class RnetSession: "read_timeout", "proxies", "verify", - "redirect", + "allow_redirects", + "pool_idle_timeout", "pool_max_idle_per_host", "pool_max_size", + "tcp_keepalive", + "tcp_keepalive_interval", + "tcp_keepalive_retries", "http1_only", "http2_only", "tcp_nodelay", @@ -597,6 +620,12 @@ class RnetSession: client = self._ensure_client() method_upper = method.upper() if isinstance(method, str) else str(method).upper() + # Per-request override of the session retry budget (rnet.request() would reject the kwarg). + # max_retries=0 means a single attempt with no session-level retries. + max_retries = kwargs.pop("max_retries", None) + if max_retries is None: + max_retries = self.max_retries + # Build URL with params url = self._build_url(url, kwargs.pop("params", None)) @@ -637,7 +666,7 @@ class RnetSession: last_exception: Optional[Exception] = None response: Optional[RnetResponse] = None - for attempt in range(self.max_retries + 1): + for attempt in range(max_retries + 1): try: raw_resp = client.request(rnet_method, url, **kwargs) response = RnetResponse(raw_resp) @@ -662,17 +691,17 @@ class RnetSession: last_exception = HTTPError(f"Received status code: {response.status_code}") self.log.warning( f"{response.status_code} {response.reason}({urlparse(url).path}). Retrying... " - f"({attempt + 1}/{self.max_retries})" + f"({attempt + 1}/{max_retries})" ) except self.catch_exceptions as e: last_exception = e response = None self.log.warning( - f"{e.__class__.__name__}({urlparse(url).path}). Retrying... ({attempt + 1}/{self.max_retries})" + f"{e.__class__.__name__}({urlparse(url).path}). Retrying... ({attempt + 1}/{max_retries})" ) - if attempt < self.max_retries: + if attempt < max_retries: if sleep_duration := self.get_sleep_time(response, attempt + 1): if sleep_duration > 0: time.sleep(sleep_duration) @@ -774,11 +803,36 @@ def session( session_kwargs: dict[str, Any] = {"impersonate": impersonate} # optional rnet client knobs, see docs/NETWORK_CONFIG.md - for key in ("http1_only", "http2_only", "pool_max_idle_per_host", "pool_max_size", "tcp_nodelay"): + for key in ( + "http1_only", + "http2_only", + "pool_max_idle_per_host", + "pool_max_size", + "tcp_nodelay", + "connect_timeout", + "read_timeout", + "timeout", + "pool_idle_timeout", + "tcp_keepalive", + "tcp_keepalive_interval", + "tcp_keepalive_retries", + "allow_redirects", + ): if key in config.network: session_kwargs[key] = config.network[key] session_kwargs.update(kwargs) + # Connection-pool / timeout defaults applied only when neither config.network nor the caller set them. + # connect_timeout + read_timeout mirror the requests-path default timeout (CONNECT_TIMEOUT, READ_TIMEOUT) + # so an unset config still gets a bounded connect and read like requests; pool_idle_timeout stays under + # the typical ~60s CDN idle kill; pool_max_idle_per_host matches the downloader worker cap to avoid + # HTTP/1.1 connection churn; tcp_keepalive keeps long idle segments warm. + session_kwargs.setdefault("connect_timeout", CONNECT_TIMEOUT) + session_kwargs.setdefault("read_timeout", READ_TIMEOUT) + session_kwargs.setdefault("pool_idle_timeout", 55) + session_kwargs.setdefault("pool_max_idle_per_host", POOL_MAX_SIZE) + session_kwargs.setdefault("tcp_keepalive", 30) + session_obj = RnetSession(**session_kwargs) session_obj.headers.update(config.headers) return session_obj diff --git a/unshackle/core/tracks/subtitle.py b/unshackle/core/tracks/subtitle.py index 203c271d..c7190295 100644 --- a/unshackle/core/tracks/subtitle.py +++ b/unshackle/core/tracks/subtitle.py @@ -246,8 +246,17 @@ class Subtitle(Track): *, cdm: Optional[object] = None, no_proxy_download: bool = False, + adaptive_workers: bool = False, ): - super().download(session, prepare_drm, max_workers, progress, cdm=cdm, no_proxy_download=no_proxy_download) + super().download( + session, + prepare_drm, + max_workers, + progress, + cdm=cdm, + no_proxy_download=no_proxy_download, + adaptive_workers=adaptive_workers, + ) if not self.path: return diff --git a/unshackle/core/tracks/track.py b/unshackle/core/tracks/track.py index 11e09c78..35cd6bd3 100644 --- a/unshackle/core/tracks/track.py +++ b/unshackle/core/tracks/track.py @@ -58,6 +58,7 @@ class DownloadContext: session: Optional[Union[Session, "RnetSession"]] = None proxy: Optional[str] = None max_workers: Optional[int] = None + adaptive_workers: bool = False license_widevine: Optional[Callable] = None cdm: Optional[object] = None @@ -232,6 +233,7 @@ class Track: *, cdm: Optional[object] = None, no_proxy_download: bool = False, + adaptive_workers: bool = False, ): """Download and optionally Decrypt this Track.""" from unshackle.core.manifests import DASH, HLS, ISM @@ -293,6 +295,7 @@ class Track: session=dl_session, proxy=proxy, max_workers=max_workers, + adaptive_workers=adaptive_workers, license_widevine=prepare_drm, cdm=cdm, ) @@ -362,6 +365,7 @@ class Track: proxy=proxy, max_workers=max_workers, session=dl_session, + adaptive=adaptive_workers, ): file_downloaded = status_update.get("file_downloaded") if not file_downloaded: