feat(downloader): adaptive worker scaling, tail boost, session parity, resume fixes

Opt-in --adaptive-workers (AIMD scaling + tail range-split), request/rnet session parity, and fixes for resume poisoning, truncated-segment acceptance, 416 handling, and byte-range-slice corruption. Tests + docs.

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

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