mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-14 04:55:07 +00:00
fix(downloader): honor Retry-After with cancel-safe backoff, add fault-injection bench and tests
This commit is contained in:
imSp4rky
+188
-1
@@ -43,6 +43,7 @@ from shutil import rmtree
|
||||
from typing import Any, Optional
|
||||
|
||||
import click
|
||||
import fault_server # local sibling (scripts/): fault-injection server + payload verifier
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
@@ -393,6 +394,174 @@ def print_report(rows: list[Row], runs: int, total_segments: int) -> None:
|
||||
)
|
||||
|
||||
|
||||
def fault_integrity(out_dir: Path, segments: int, seg_size: int) -> tuple[list[int], list[int]]:
|
||||
"""Verify every output file byte-for-byte against its seeded expected payload.
|
||||
|
||||
Returns (missing, corrupt). ``corrupt`` (present but not byte-identical) is the loud
|
||||
CORRUPTION signal: a fault path that finalizes wrong bytes must fail the run, not pass.
|
||||
"""
|
||||
missing: list[int] = []
|
||||
corrupt: list[int] = []
|
||||
for i in range(segments):
|
||||
p = out_dir / f"seg_{i}.bin"
|
||||
if not p.exists():
|
||||
missing.append(i)
|
||||
continue
|
||||
if p.read_bytes() != fault_server.segment_payload(f"/seg/{i}", seg_size):
|
||||
corrupt.append(i)
|
||||
return missing, corrupt
|
||||
|
||||
|
||||
# per worker-count under a fault profile: workers, median_s, MB/s, missing, corrupt, raised,
|
||||
# requests seen, amplification (requests/segments), retries (requests-segments), status_counts,
|
||||
# then worst-across-runs sweep safety: strays, locked handles, live pool threads
|
||||
FaultRow = tuple[int, float, float, list[int], list[int], bool, int, float, int, dict[int, int], int, int, int]
|
||||
|
||||
|
||||
def print_fault_report(rows: list[FaultRow], runs: int, segments: int, profile: str) -> None:
|
||||
plural = "s" if runs != 1 else ""
|
||||
table = Table(title=f"fault profile '{profile}': {segments} segments (median of {runs} run{plural})")
|
||||
table.add_column("workers", justify="right")
|
||||
table.add_column("seconds", justify="right")
|
||||
table.add_column("MB/s", justify="right")
|
||||
table.add_column("status", justify="left")
|
||||
table.add_column("requests", justify="right")
|
||||
table.add_column("amp", justify="right")
|
||||
table.add_column("sweep", justify="left")
|
||||
for w, secs, mbps, missing, corrupt, raised, requests_seen, amp, _retries, _codes, strays, locked, threads in rows:
|
||||
if corrupt:
|
||||
status = f"[red]CORRUPTION {len(corrupt)} segs[/]"
|
||||
elif missing:
|
||||
status = f"[red]FAIL {len(missing)} incomplete[/]"
|
||||
elif raised:
|
||||
status = "[yellow]raised (all present)[/]"
|
||||
else:
|
||||
status = "[green]ok[/]"
|
||||
if locked:
|
||||
sweep_status = f"[red]{locked} LOCKED (WinError 32)[/]"
|
||||
elif threads:
|
||||
sweep_status = f"[yellow]{threads} pool threads live[/]"
|
||||
elif strays:
|
||||
sweep_status = f"{strays} strays"
|
||||
else:
|
||||
sweep_status = "[green]clean[/]"
|
||||
table.add_row(str(w), f"{secs:.3f}", f"{mbps:.1f}", status, str(requests_seen), f"{amp:.2f}x", sweep_status)
|
||||
console.print(table)
|
||||
for w, _secs, _mbps, _missing, _corrupt, _raised, requests_seen, amp, retries, codes, *_rest in rows:
|
||||
code_str = ", ".join(f"{c}:{n}" for c, n in codes.items()) or "none"
|
||||
console.print(
|
||||
f"amplification workers={w:<3} {requests_seen} requests for {segments} segments "
|
||||
f"({amp:.2f}x, {retries} excess requests (retries/hedges/probes/parts)), statuses [{code_str}]"
|
||||
)
|
||||
|
||||
|
||||
def _procs_notice(procs: int, n_urls: int, session: Optional[Any]) -> None:
|
||||
"""Warn loudly when --procs cannot engage, so a run is never mistaken for a multiprocess test."""
|
||||
if procs <= 1:
|
||||
return
|
||||
if n_urls < dl_mod.MP_MIN_SEGMENTS:
|
||||
console.print(
|
||||
f"[yellow]--procs {procs} will NOT engage: {n_urls} segments < MP_MIN_SEGMENTS "
|
||||
f"({dl_mod.MP_MIN_SEGMENTS}); the single-process path runs instead[/]"
|
||||
)
|
||||
elif dl_mod._build_session_spec(session) is None:
|
||||
console.print(
|
||||
f"[yellow]--procs {procs} will NOT engage: session not rebuildable in child processes; "
|
||||
f"the single-process path runs instead[/]"
|
||||
)
|
||||
|
||||
|
||||
def run_fault_mode(
|
||||
profile_name: str,
|
||||
segments: int,
|
||||
seg_size: int,
|
||||
worker_counts: list[int],
|
||||
runs: int,
|
||||
session: Optional[Any],
|
||||
adaptive: bool,
|
||||
procs: int,
|
||||
seed: int,
|
||||
) -> None:
|
||||
"""Start a FaultServer in-process, sweep worker counts, and report reliability + overhead.
|
||||
|
||||
Reuses run_batch/tail_droop/sweep_strays (via run_batch) and the report conventions, adding
|
||||
request amplification and byte-exact integrity, the metrics that catch retry/hedge storms
|
||||
and silent corruption against a misbehaving host.
|
||||
"""
|
||||
profile = fault_server.resolve_profile(profile_name)
|
||||
server = fault_server.FaultServer(profile, seg_size, segments, seed)
|
||||
threading.Thread(target=server.serve_forever, daemon=True).start()
|
||||
host, port = str(server.server_address[0]), int(server.server_address[1])
|
||||
urls = [f"http://{host}:{port}/seg/{i}" for i in range(segments)]
|
||||
console.print(
|
||||
f"[cyan]Fault mode:[/] profile={profile_name} seed={seed} "
|
||||
f"{segments} x {seg_size / MIB:.2f} MiB, impersonate={getattr(session, 'impersonate_name', None) or 'none'}"
|
||||
)
|
||||
|
||||
rows: list[FaultRow] = []
|
||||
try:
|
||||
with tempfile.TemporaryDirectory(prefix="faultbench_") as root_str:
|
||||
root = Path(root_str)
|
||||
for w in worker_counts:
|
||||
samples: list[tuple[float, float, list[int], list[int], bool, int, dict[int, int], int, int, int]] = []
|
||||
for _ in range(runs):
|
||||
server.reset_counters()
|
||||
out_dir = Path(tempfile.mkdtemp(prefix="run_", dir=root))
|
||||
secs, _comps, err, strays, locked, pool_threads = run_batch(
|
||||
urls, out_dir, w, session, adaptive, procs
|
||||
)
|
||||
missing, corrupt = fault_integrity(out_dir, segments, seg_size)
|
||||
stats = server.stats()
|
||||
samples.append(
|
||||
(
|
||||
secs,
|
||||
downloaded_mib(out_dir),
|
||||
missing,
|
||||
corrupt,
|
||||
err is not None,
|
||||
stats["requests"],
|
||||
stats["status_counts"],
|
||||
strays,
|
||||
locked,
|
||||
pool_threads,
|
||||
)
|
||||
)
|
||||
rmtree(out_dir, ignore_errors=True)
|
||||
median_s = statistics.median(s[0] for s in samples)
|
||||
best = min(samples, key=lambda s: abs(s[0] - median_s)) # representative run near the median
|
||||
mbps = best[1] / best[0] if best[0] else 0.0
|
||||
requests_seen = best[5]
|
||||
amp = requests_seen / segments if segments else 0.0
|
||||
retries = max(0, requests_seen - segments)
|
||||
rows.append(
|
||||
(
|
||||
w,
|
||||
median_s,
|
||||
mbps,
|
||||
best[2],
|
||||
best[3],
|
||||
best[4],
|
||||
requests_seen,
|
||||
amp,
|
||||
retries,
|
||||
best[6],
|
||||
max(s[7] for s in samples),
|
||||
max(s[8] for s in samples),
|
||||
max(s[9] for s in samples),
|
||||
)
|
||||
)
|
||||
finally:
|
||||
server.stall_event.set() # release any stalled handler threads before shutdown
|
||||
server.shutdown()
|
||||
|
||||
print_fault_report(rows, runs, segments, profile_name)
|
||||
|
||||
broken = {row[0]: (row[3], row[4]) for row in rows if row[3] or row[4]}
|
||||
if broken:
|
||||
detail = "; ".join(f"workers={w}: missing {miss} corrupt {corr}" for w, (miss, corr) in broken.items())
|
||||
raise click.ClickException(f"incomplete/corrupt downloads under fault profile: {detail}")
|
||||
|
||||
|
||||
@click.command(context_settings={"help_option_names": ["-h", "--help"]})
|
||||
@click.option("--segments", default=64, show_default=True, help="Synthetic segment count.")
|
||||
@click.option("--seg-size", default=2 * MIB, show_default=True, help="Bytes per synthetic segment.")
|
||||
@@ -434,6 +603,12 @@ def print_report(rows: list[Row], runs: int, total_segments: int) -> None:
|
||||
@click.option(
|
||||
"--no-read1", is_flag=True, help="A/B: force racers back to blocking full-chunk reads (pre-fix behavior)."
|
||||
)
|
||||
@click.option(
|
||||
"--fault-profile",
|
||||
type=click.Choice(sorted(fault_server.PROFILES)),
|
||||
help="Reliability mode: serve via the fault-injection server (fault_server.py) and verify integrity + amplification.",
|
||||
)
|
||||
@click.option("--seed", default=fault_server.DEFAULT_SEED, show_default=True, help="Fault RNG seed (--fault-profile).")
|
||||
def main(
|
||||
segments: int,
|
||||
seg_size: int,
|
||||
@@ -452,6 +627,8 @@ def main(
|
||||
adaptive: bool,
|
||||
procs: int,
|
||||
no_read1: bool,
|
||||
fault_profile: Optional[str],
|
||||
seed: int,
|
||||
) -> None:
|
||||
"""Benchmark downloader throughput vs worker count, with optional fault injection."""
|
||||
worker_counts = [int(w) for w in workers.split(",") if w.strip()]
|
||||
@@ -481,7 +658,7 @@ def main(
|
||||
saved["RACER_READ1"] = dl_mod.RACER_READ1
|
||||
dl_mod.RACER_READ1 = False
|
||||
else:
|
||||
console.print("[yellow]--no-read1 ignored: this checkout has no RACER_READ1 (pre-fix behavior anyway)[/]")
|
||||
console.print("[yellow]--no-read1 ignored: only applies to older checkouts that still have RACER_READ1[/]")
|
||||
|
||||
if impersonate:
|
||||
from unshackle.core.session import session as make_session
|
||||
@@ -489,6 +666,14 @@ def main(
|
||||
session = make_session(impersonate)
|
||||
|
||||
try:
|
||||
if fault_profile:
|
||||
if urls_file:
|
||||
console.print("[yellow]--urls-file ignored in --fault-profile mode[/]")
|
||||
if any(faults.values()):
|
||||
console.print("[yellow]--fault-* flags ignored in --fault-profile mode (use the profile knobs)[/]")
|
||||
_procs_notice(procs, segments, session)
|
||||
run_fault_mode(fault_profile, segments, seg_size, worker_counts, runs, session, adaptive, procs, seed)
|
||||
return
|
||||
if urls_file:
|
||||
if any(faults.values()):
|
||||
console.print("[yellow]fault flags ignored in --urls-file mode[/]")
|
||||
@@ -496,6 +681,7 @@ def main(
|
||||
if not urls:
|
||||
raise click.ClickException("--urls-file contained no URLs")
|
||||
total = len(urls)
|
||||
_procs_notice(procs, len(urls), session)
|
||||
console.print(f"[cyan]URL-file mode:[/] {len(urls)} URLs, impersonate={impersonate or 'none'}")
|
||||
else:
|
||||
# deterministic shared buffer; generated once, served for every segment, and
|
||||
@@ -507,6 +693,7 @@ def main(
|
||||
host, port = str(server.server_address[0]), int(server.server_address[1])
|
||||
urls = [f"http://{host}:{port}/seg/{i}" for i in range(segments)]
|
||||
active = {k: v for k, v in faults.items() if v}
|
||||
_procs_notice(procs, segments, session)
|
||||
console.print(
|
||||
f"[cyan]Synthetic mode:[/] {segments} x {seg_size / MIB:.2f} MiB "
|
||||
f"(latency={latency_ms}ms, rate={rate_kib or 'unbounded'} KiB/s, "
|
||||
|
||||
@@ -73,6 +73,16 @@ run scripts/bench_downloader.py --fault-503 2 --fast-timeouts --segments 8
|
||||
run scripts/bench_downloader.py --fault-reset 2 --fast-timeouts --segments 8
|
||||
run scripts/bench_downloader.py --fault-stall 1 --fast-timeouts --segments 8 || true
|
||||
|
||||
# 2b. Fault-injection reliability + request amplification (fault_server.py; every run is
|
||||
# byte-integrity verified, so a CORRUPTION line here is a real finalize-wrong-bytes bug).
|
||||
# loopback is the zero-fault python-overhead ceiling; the rest exercise retry/hedge/tail paths.
|
||||
run scripts/bench_downloader.py --fault-profile loopback --segments 64 --workers 8,16 --runs 3
|
||||
run scripts/bench_downloader.py --fault-profile rate-limit --segments 32 --workers 8
|
||||
run scripts/bench_downloader.py --fault-profile flaky-first --segments 16 --workers 8 --fast-timeouts
|
||||
run scripts/bench_downloader.py --fault-profile reset --segments 16 --workers 8 --fast-timeouts
|
||||
run scripts/bench_downloader.py --fault-profile stall --segments 16 --workers 8 --fast-timeouts
|
||||
run scripts/bench_downloader.py --fault-profile tail-slow --segments 32 --workers 8 --adaptive --fast-timeouts
|
||||
|
||||
# 3. Multiprocess path (>=24 segments engages spawned children; strided chunk split).
|
||||
# big batch: at small sizes the fixed spawn cost drowns the throughput signal
|
||||
run scripts/bench_downloader.py --segments 128 --workers 8 --procs 2
|
||||
@@ -91,6 +101,18 @@ run scripts/bench_downloader.py --segments 32 --workers 8 --impersonate Chrome13
|
||||
echo; echo "==================== pytest tests/test_downloader_leaks.py ===================="
|
||||
uv run pytest tests/test_downloader_leaks.py -q || echo "(pytest failed or missing - run 'uv sync' to install dev deps)"
|
||||
|
||||
# 6b. Low-core variant: pin to two cores to emulate a small box, so the python-overhead
|
||||
# ceiling and one fault profile can be read against the full-core numbers above. taskset is
|
||||
# util-linux (Linux only); skipped cleanly elsewhere.
|
||||
if command -v taskset >/dev/null 2>&1 && taskset -c 0-1 true 2>/dev/null; then
|
||||
echo; echo "==================== low-core variant (taskset -c 0-1, 2 cores) ===================="
|
||||
run_lowcore() { echo; echo "-------------------- taskset -c 0-1 $* --------------------"; taskset -c 0-1 uv run python "$@"; }
|
||||
run_lowcore scripts/bench_downloader.py --fault-profile loopback --segments 64 --workers 8,16 --runs 3
|
||||
run_lowcore scripts/bench_downloader.py --fault-profile rate-limit --segments 32 --workers 8
|
||||
else
|
||||
echo; echo "==================== low-core variant skipped (taskset missing or CPUs 0-1 not schedulable) ===================="
|
||||
fi
|
||||
|
||||
# 7. Live network tests (opt-in)
|
||||
if [[ "${1:-}" == "--live" ]]; then
|
||||
run scripts/live_manifest_test.py --workers 16 --track largest
|
||||
|
||||
@@ -0,0 +1,379 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Localhost fault-injection HTTP server for downloader reliability benchmarking.
|
||||
|
||||
Serves deterministic pseudo-random segment payloads (seeded per URL path, so the
|
||||
bytes are reproducible and any output file can be hashed against them to prove
|
||||
integrity) and applies a configurable *fault profile* to those responses. Stall
|
||||
and reset faults are decided deterministically from a printed seed keyed on
|
||||
(path, attempt), so a retry can recover instead of failing forever; rate-limit
|
||||
windows are wall-clock relative to run start, so their timing is not seed-bound.
|
||||
|
||||
Fault profiles:
|
||||
stall sleep N seconds mid-body on X% of (path, attempt)s
|
||||
reset RST the connection mid-body on X% of (path, attempt)s
|
||||
rate-limit reply 429 + Retry-After during recurring throttle windows
|
||||
(window_open seconds serving / window_closed seconds throttled)
|
||||
tail-slow heavily throttle the last K segments' bodies
|
||||
flaky-first fail attempt 1 per path (503), succeed on retry
|
||||
loopback no faults; payloads served from memory (overhead ceiling)
|
||||
|
||||
Range requests are honored (206 + Content-Range) so byte-range / hedge / tail-boost
|
||||
download paths exercise correctly, and an unsatisfiable resume gets a 416.
|
||||
|
||||
Standalone:
|
||||
uv run python scripts/fault_server.py --profile rate-limit --segments 32
|
||||
uv run python scripts/fault_server.py --profile stall --stall-secs 4 --port 8080
|
||||
|
||||
Importable: build a ``FaultServer``, serve it on a thread, and hash output against
|
||||
``segment_payload(path, size)``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import random
|
||||
import socket
|
||||
import struct
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import asdict, dataclass
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from typing import Any, Optional
|
||||
|
||||
MIB = 1024 * 1024
|
||||
DEFAULT_SEED = 1729
|
||||
|
||||
|
||||
def segment_payload(path: str, size: int) -> bytes:
|
||||
"""Deterministic pseudo-random body for ``path``.
|
||||
|
||||
Pure and seed-stable: the bench regenerates this to byte-compare every output
|
||||
file, so it must never depend on the fault seed, request order, or attempt.
|
||||
"""
|
||||
seed = int.from_bytes(hashlib.sha256(path.encode("utf-8")).digest()[:8], "big")
|
||||
return random.Random(seed).randbytes(size)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FaultProfile:
|
||||
"""A named set of deterministic response faults. Presets live in ``PROFILES``."""
|
||||
|
||||
name: str
|
||||
stall_pct: float = 0.0 # % of (path, attempt)s that stall mid-body
|
||||
stall_secs: float = 3.0 # seconds to hold when a stall fires
|
||||
reset_pct: float = 0.0 # % of (path, attempt)s that RST mid-body
|
||||
rate_limit: bool = False # 429 + Retry-After during closed windows
|
||||
retry_after: int = 1 # Retry-After seconds advertised on a 429
|
||||
window_open: float = 5.0 # seconds serving normally before a throttle window
|
||||
window_closed: float = 5.0 # seconds replying 429 during a throttle window
|
||||
body_rate_kib: int = 0 # throttle every segment body to this KiB/s (0 = unthrottled)
|
||||
tail_slow: int = 0 # throttle the last K segments' bodies
|
||||
tail_rate_kib: int = 256 # throttle rate for tail-slow segments (KiB/s)
|
||||
flaky_first: bool = False # fail attempt 1 per path, succeed after
|
||||
flaky_status: int = 503 # status returned by a flaky-first failure
|
||||
|
||||
|
||||
PROFILES: dict[str, FaultProfile] = {
|
||||
"loopback": FaultProfile("loopback"),
|
||||
"stall": FaultProfile("stall", stall_pct=25.0, stall_secs=3.0),
|
||||
"reset": FaultProfile("reset", reset_pct=25.0),
|
||||
# short windows + a body throttle so a fast loopback batch actually spans several throttle
|
||||
# windows and trips the 429 + Retry-After retry path (a real 900 MB/s batch would otherwise
|
||||
# finish inside the first open window). Standalone use can widen the windows via overrides.
|
||||
"rate-limit": FaultProfile(
|
||||
"rate-limit", rate_limit=True, retry_after=1, window_open=1.0, window_closed=1.0, body_rate_kib=2048
|
||||
),
|
||||
"tail-slow": FaultProfile("tail-slow", tail_slow=4, tail_rate_kib=256),
|
||||
"flaky-first": FaultProfile("flaky-first", flaky_first=True, flaky_status=503),
|
||||
}
|
||||
|
||||
|
||||
class FaultServer(ThreadingHTTPServer):
|
||||
"""Threaded loopback server serving seeded payloads under a ``FaultProfile``.
|
||||
|
||||
Counters (total requests, per-status, per-path attempts) are the request-amplification
|
||||
signal the bench reads: a retry/hedge storm against a throttled host shows up as
|
||||
requests >> segment count. ``reset_counters`` re-baselines between benchmark runs.
|
||||
"""
|
||||
|
||||
allow_reuse_address = True
|
||||
daemon_threads = True
|
||||
request_queue_size = 128 # default 5 drops SYNs past 5 concurrent connects -> phantom 1s retransmit tail
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
profile: FaultProfile,
|
||||
seg_size: int,
|
||||
total_segments: int,
|
||||
seed: int = DEFAULT_SEED,
|
||||
host: str = "127.0.0.1",
|
||||
port: int = 0,
|
||||
) -> None:
|
||||
self.profile = profile
|
||||
self.seg_size = seg_size
|
||||
self.total_segments = total_segments
|
||||
self.seed = seed
|
||||
self.stall_event = threading.Event() # set at teardown to release stalled handlers
|
||||
self._lock = threading.Lock()
|
||||
self._request_count = 0
|
||||
self._status_counts: dict[int, int] = {}
|
||||
self._attempts: dict[str, int] = {}
|
||||
self._window_origin = time.monotonic()
|
||||
# precompute the known segment bodies so serving is a dict lookup (no per-request
|
||||
# generation), making the loopback profile a true python-overhead ceiling
|
||||
self._payloads: dict[str, bytes] = {
|
||||
f"/seg/{i}": segment_payload(f"/seg/{i}", seg_size) for i in range(total_segments)
|
||||
}
|
||||
super().__init__((host, port), FaultHandler)
|
||||
|
||||
def handle_error(self, request: Any, client_address: Any) -> None: # silence expected fault noise
|
||||
pass
|
||||
|
||||
def payload(self, path: str) -> bytes:
|
||||
cached = self._payloads.get(path)
|
||||
if cached is not None:
|
||||
return cached
|
||||
with self._lock: # unknown path (standalone use): generate + cache once
|
||||
body = self._payloads.get(path)
|
||||
if body is None:
|
||||
body = segment_payload(path, self.seg_size)
|
||||
self._payloads[path] = body
|
||||
return body
|
||||
|
||||
def note_request(self) -> None:
|
||||
with self._lock:
|
||||
self._request_count += 1
|
||||
|
||||
def note_status(self, code: int) -> None:
|
||||
with self._lock:
|
||||
self._status_counts[code] = self._status_counts.get(code, 0) + 1
|
||||
|
||||
def note_attempt(self, path: str) -> int:
|
||||
with self._lock:
|
||||
n = self._attempts.get(path, 0) + 1
|
||||
self._attempts[path] = n
|
||||
return n
|
||||
|
||||
def in_throttle_window(self) -> bool:
|
||||
p = self.profile
|
||||
period = p.window_open + p.window_closed
|
||||
if period <= 0:
|
||||
return False
|
||||
phase = (time.monotonic() - self._window_origin) % period
|
||||
return phase >= p.window_open
|
||||
|
||||
def decides(self, kind: str, path: str, attempt: int, pct: float) -> bool:
|
||||
"""Deterministic per-(path, attempt) fault gate: fires for ``pct`` percent of them."""
|
||||
if pct <= 0:
|
||||
return False
|
||||
digest = hashlib.sha256(f"{self.seed}:{kind}:{path}:{attempt}".encode("utf-8")).digest()
|
||||
return (int.from_bytes(digest[:4], "big") / 0xFFFFFFFF) * 100.0 < pct
|
||||
|
||||
def reset_counters(self) -> None:
|
||||
with self._lock:
|
||||
self._request_count = 0
|
||||
self._status_counts.clear()
|
||||
self._attempts.clear()
|
||||
self._window_origin = time.monotonic() # realign throttle windows to this run's start
|
||||
|
||||
def stats(self) -> dict[str, Any]:
|
||||
with self._lock:
|
||||
max_attempts = max(self._attempts.values(), default=0)
|
||||
return {
|
||||
"requests": self._request_count,
|
||||
"status_counts": dict(sorted(self._status_counts.items())),
|
||||
"max_attempts": max_attempts,
|
||||
}
|
||||
|
||||
|
||||
class FaultHandler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
def log_message(self, *_: Any) -> None: # silence per-request logging
|
||||
pass
|
||||
|
||||
def _parse_range(self, total: int) -> Optional[tuple[int, int]]:
|
||||
header = self.headers.get("Range")
|
||||
if not header or not header.startswith("bytes="):
|
||||
return None
|
||||
start_s, _, end_s = header[len("bytes=") :].partition("-")
|
||||
if not start_s: # RFC 7233 suffix range: bytes=-N means the last N bytes
|
||||
return max(0, total - int(end_s)), total - 1
|
||||
start = int(start_s)
|
||||
end = int(end_s) if end_s else total - 1
|
||||
if start > end: # inverted range: report unsatisfiable so _serve answers 416
|
||||
return total, total - 1
|
||||
return start, min(end, total - 1)
|
||||
|
||||
def do_GET(self) -> None:
|
||||
server: FaultServer = self.server # type: ignore[assignment]
|
||||
server.note_request()
|
||||
if not self.path.startswith("/seg/"):
|
||||
self.send_error(404)
|
||||
return
|
||||
try:
|
||||
idx = int(self.path.rsplit("/", 1)[-1])
|
||||
except ValueError:
|
||||
self.send_error(404)
|
||||
return
|
||||
|
||||
attempt = server.note_attempt(self.path)
|
||||
prof = server.profile
|
||||
|
||||
if prof.flaky_first and attempt == 1: # first attempt per path always fails, retry recovers
|
||||
self._send_status(prof.flaky_status)
|
||||
return
|
||||
if prof.rate_limit and server.in_throttle_window(): # recurring throttle window
|
||||
self._send_429(prof.retry_after)
|
||||
return
|
||||
|
||||
rate = prof.body_rate_kib
|
||||
if prof.tail_slow and idx >= server.total_segments - prof.tail_slow:
|
||||
rate = prof.tail_rate_kib
|
||||
stall = server.decides("stall", self.path, attempt, prof.stall_pct)
|
||||
reset = server.decides("reset", self.path, attempt, prof.reset_pct)
|
||||
self._serve(server.payload(self.path), stall=stall, reset=reset, rate_kib=rate)
|
||||
|
||||
def _send_status(self, code: int) -> None:
|
||||
server: FaultServer = self.server # type: ignore[assignment]
|
||||
server.note_status(code)
|
||||
self.send_response(code)
|
||||
self.send_header("Content-Length", "0")
|
||||
self.end_headers()
|
||||
|
||||
def _send_429(self, retry_after: int) -> None:
|
||||
server: FaultServer = self.server # type: ignore[assignment]
|
||||
server.note_status(429)
|
||||
self.send_response(429)
|
||||
self.send_header("Retry-After", str(retry_after))
|
||||
self.send_header("Content-Length", "0")
|
||||
self.end_headers()
|
||||
|
||||
def _serve(self, full: bytes, stall: bool = False, reset: bool = False, rate_kib: int = 0) -> None:
|
||||
server: FaultServer = self.server # type: ignore[assignment]
|
||||
rng = self._parse_range(len(full))
|
||||
if rng is not None and rng[0] >= len(full): # unsatisfiable resume: a real CDN answers 416
|
||||
server.note_status(416)
|
||||
self.send_response(416)
|
||||
self.send_header("Content-Range", f"bytes */{len(full)}")
|
||||
self.send_header("Content-Length", "0")
|
||||
self.end_headers()
|
||||
return
|
||||
if rng is not None:
|
||||
start, end = rng
|
||||
body = full[start : end + 1]
|
||||
server.note_status(206)
|
||||
self.send_response(206)
|
||||
self.send_header("Content-Range", f"bytes {start}-{end}/{len(full)}")
|
||||
else:
|
||||
body = full
|
||||
server.note_status(200)
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.send_header("Accept-Ranges", "bytes")
|
||||
self.send_header("Content-Type", "application/octet-stream")
|
||||
self.end_headers()
|
||||
|
||||
try:
|
||||
if reset: # send half, then abort the TCP connection (RST via SO_LINGER=0)
|
||||
half = max(1, len(body) // 2)
|
||||
self.wfile.write(body[:half])
|
||||
self.wfile.flush()
|
||||
sock = self.connection
|
||||
sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0))
|
||||
sock.close()
|
||||
self.close_connection = True
|
||||
elif stall: # send half, hold (a client READ_TIMEOUT shorter than stall_secs will retry), then finish
|
||||
half = max(1, len(body) // 2)
|
||||
self.wfile.write(body[:half])
|
||||
self.wfile.flush()
|
||||
server.stall_event.wait(server.profile.stall_secs)
|
||||
self.wfile.write(body[half:])
|
||||
self.wfile.flush()
|
||||
else:
|
||||
self._write_body(body, rate_kib)
|
||||
except (BrokenPipeError, ConnectionError, OSError):
|
||||
self.close_connection = True
|
||||
|
||||
def _write_body(self, body: bytes, rate_kib: int) -> None:
|
||||
if rate_kib <= 0:
|
||||
self.wfile.write(body)
|
||||
return
|
||||
chunk = 64 * 1024 # throttle: fixed-size chunks paced to rate_kib KiB/s
|
||||
per_chunk = chunk / (rate_kib * 1024)
|
||||
for off in range(0, len(body), chunk):
|
||||
self.wfile.write(body[off : off + chunk])
|
||||
self.wfile.flush()
|
||||
time.sleep(per_chunk)
|
||||
|
||||
|
||||
def resolve_profile(name: str, overrides: Optional[dict[str, Any]] = None) -> FaultProfile:
|
||||
"""Return a copy of the named preset with any non-None ``overrides`` applied."""
|
||||
if name not in PROFILES:
|
||||
raise KeyError(f"unknown profile {name!r}; choices: {', '.join(sorted(PROFILES))}")
|
||||
base = asdict(PROFILES[name])
|
||||
for key, value in (overrides or {}).items():
|
||||
if value is not None and key in base:
|
||||
base[key] = value
|
||||
return FaultProfile(**base)
|
||||
|
||||
|
||||
def _build_arg_parser() -> argparse.ArgumentParser:
|
||||
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
p.add_argument("--port", type=int, default=0, help="Listen port (0 = auto).")
|
||||
p.add_argument("--profile", choices=sorted(PROFILES), default="loopback", help="Fault profile to apply.")
|
||||
p.add_argument("--segments", type=int, default=32, help="Number of /seg/i payloads to precompute.")
|
||||
p.add_argument("--seg-size", type=int, default=2 * MIB, help="Bytes per segment payload.")
|
||||
p.add_argument("--seed", type=int, default=DEFAULT_SEED, help="Fault RNG seed (printed for reproducibility).")
|
||||
# profile knobs (override the chosen preset when given)
|
||||
p.add_argument("--stall-pct", type=float, help="Override: %% of (path, attempt)s that stall.")
|
||||
p.add_argument("--stall-secs", type=float, help="Override: seconds a stall holds.")
|
||||
p.add_argument("--reset-pct", type=float, help="Override: %% of (path, attempt)s that RST.")
|
||||
p.add_argument("--retry-after", type=int, help="Override: Retry-After seconds on a 429.")
|
||||
p.add_argument("--window-open", type=float, help="Override: seconds serving before a throttle window.")
|
||||
p.add_argument("--window-closed", type=float, help="Override: seconds replying 429 during a window.")
|
||||
p.add_argument("--body-rate-kib", type=int, help="Override: throttle every body to this KiB/s (0 = off).")
|
||||
p.add_argument("--tail-slow", type=int, help="Override: throttle the last K segments.")
|
||||
p.add_argument("--tail-rate-kib", type=int, help="Override: tail-slow throttle rate (KiB/s).")
|
||||
return p
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _build_arg_parser().parse_args()
|
||||
overrides = {
|
||||
"stall_pct": args.stall_pct,
|
||||
"stall_secs": args.stall_secs,
|
||||
"reset_pct": args.reset_pct,
|
||||
"retry_after": args.retry_after,
|
||||
"window_open": args.window_open,
|
||||
"window_closed": args.window_closed,
|
||||
"body_rate_kib": args.body_rate_kib,
|
||||
"tail_slow": args.tail_slow,
|
||||
"tail_rate_kib": args.tail_rate_kib,
|
||||
}
|
||||
profile = resolve_profile(args.profile, overrides)
|
||||
server = FaultServer(profile, args.seg_size, args.segments, args.seed, port=args.port)
|
||||
host, port = server.server_address[0], server.server_address[1]
|
||||
config = {
|
||||
"listen": f"http://{host}:{port}",
|
||||
"seed": args.seed,
|
||||
"segments": args.segments,
|
||||
"seg_size": args.seg_size,
|
||||
"base_url": f"http://{host}:{port}/seg/{{i}}",
|
||||
"profile": asdict(profile),
|
||||
}
|
||||
print(json.dumps(config, indent=2), flush=True)
|
||||
try:
|
||||
server.serve_forever()
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
finally:
|
||||
server.stall_event.set()
|
||||
server.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -11,9 +11,16 @@ 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.downloaders.requests import (
|
||||
TAIL_BOOST_MIN_SEGMENT_SIZE,
|
||||
TAIL_BOOST_PART_SIZE,
|
||||
AdaptiveWorkerController,
|
||||
_has_range_header,
|
||||
_plan_tail_parts,
|
||||
_split_ranges,
|
||||
_tail_boost_engages,
|
||||
requests,
|
||||
)
|
||||
from unshackle.core.tracks.track import DownloadContext
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
@@ -202,6 +209,22 @@ def test_plan_tail_parts_skips_small_or_no_spare() -> None:
|
||||
assert _plan_tail_parts(20 * 1024 * 1024, spare=0) == []
|
||||
|
||||
|
||||
def test_split_ranges_covers_exactly_no_gaps_no_overlaps() -> None:
|
||||
for size in (1, 100, 4 * 1024 * 1024, 16 * 1024 * 1024 + 7, 64 * 1024 * 1024):
|
||||
parts = _split_ranges(size, max_parts=8, part_target=4 * 1024 * 1024)
|
||||
assert parts[0][0] == 0 and parts[-1][1] == size - 1
|
||||
for (_, prev_end), (start, _) in zip(parts, parts[1:]):
|
||||
assert start == prev_end + 1
|
||||
assert all(start <= end for start, end in parts)
|
||||
assert sum(end - start + 1 for start, end in parts) == size
|
||||
|
||||
|
||||
def test_split_ranges_respects_max_parts_and_target() -> None:
|
||||
assert len(_split_ranges(100 * 1024 * 1024, max_parts=4, part_target=16 * 1024 * 1024)) == 4 # capped
|
||||
assert len(_split_ranges(10 * 1024 * 1024, max_parts=8, part_target=16 * 1024 * 1024)) == 1 # below one target
|
||||
assert _split_ranges(1, max_parts=8, part_target=16 * 1024 * 1024) == [(0, 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
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
"""Integration coverage for the CONSUMING side of the DASH single-URL byte-range collapse.
|
||||
|
||||
The pure predicate ``DASH._collapsible_single_url`` is unit-tested in
|
||||
``test_dash_collapse.py``; this file drives the real ``DASH.download_track`` code path that
|
||||
acts on it. A minimal MPD is parsed with ``DASH.from_text().to_tracks()`` (so the track's
|
||||
``data["dash"]`` is built by real parser code, not hand-faked) and downloaded against a
|
||||
localhost server, asserting that:
|
||||
|
||||
- a SegmentList whose ranges collapse fetches the parent resource whole in one direct
|
||||
(Range-less) request and writes a byte-correct file at ``save_path``, leaving no
|
||||
segment dir or ``.!dev`` artifacts, and
|
||||
- a track with mixed segment URLs still takes the per-segment download-and-merge path.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from functools import partial
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
import pytest
|
||||
|
||||
from unshackle.core.constants import DOWNLOAD_CANCELLED
|
||||
from unshackle.core.manifests import DASH
|
||||
from unshackle.core.tracks.track import DownloadContext
|
||||
|
||||
# distinctive parent resource for the collapse case; init is its [0, INIT_LEN) prefix
|
||||
PARENT = bytes(i % 256 for i in range(10240))
|
||||
INIT_LEN = 512
|
||||
|
||||
# separate resources for the non-collapse (mixed-URL) case
|
||||
INIT2 = bytes(range(100))
|
||||
PART0 = bytes([0xA0]) * 300
|
||||
PART1 = bytes([0xB1]) * 200
|
||||
|
||||
|
||||
class _Handler(BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
server = self.server
|
||||
rng = self.headers.get("Range")
|
||||
with server.lock:
|
||||
server.requests.append((self.path, rng))
|
||||
body = server.routes.get(self.path)
|
||||
if body is None:
|
||||
self.send_response(404)
|
||||
self.send_header("Content-Length", "0")
|
||||
self.end_headers()
|
||||
return
|
||||
if rng:
|
||||
start, end = (int(x) for x in rng.removeprefix("bytes=").split("-"))
|
||||
chunk = body[start : end + 1]
|
||||
self.send_response(206)
|
||||
self.send_header("Content-Range", f"bytes {start}-{end}/{len(body)}")
|
||||
body = chunk
|
||||
else:
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server():
|
||||
srv = ThreadingHTTPServer(("127.0.0.1", 0), _Handler)
|
||||
srv.lock = threading.Lock()
|
||||
srv.requests = []
|
||||
srv.routes = {}
|
||||
thread = threading.Thread(target=srv.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
yield srv
|
||||
srv.shutdown()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_cancel():
|
||||
was_set = DOWNLOAD_CANCELLED.is_set()
|
||||
DOWNLOAD_CANCELLED.clear()
|
||||
yield
|
||||
if was_set:
|
||||
DOWNLOAD_CANCELLED.set()
|
||||
else:
|
||||
DOWNLOAD_CANCELLED.clear()
|
||||
|
||||
|
||||
def _base(srv) -> str:
|
||||
return f"http://127.0.0.1:{srv.server_address[1]}/"
|
||||
|
||||
|
||||
def _collapse_mpd(srv) -> str:
|
||||
# SegmentList on one media.mp4: init is bytes [0, INIT_LEN), media ranges cover the rest
|
||||
# contiguously -> _collapsible_single_url returns the URL and download_track collapses.
|
||||
return f"""<?xml version="1.0"?>
|
||||
<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" profiles="urn:mpeg:dash:profile:isoff-on-demand:2011">
|
||||
<BaseURL>{_base(srv)}</BaseURL>
|
||||
<Period id="0">
|
||||
<AdaptationSet contentType="video" mimeType="video/mp4" lang="en">
|
||||
<Representation id="v0" codecs="avc1.640028" bandwidth="1000000" width="1920" height="1080">
|
||||
<BaseURL>media.mp4</BaseURL>
|
||||
<SegmentList timescale="1">
|
||||
<Initialization range="0-{INIT_LEN - 1}"/>
|
||||
<SegmentURL mediaRange="{INIT_LEN}-5119"/>
|
||||
<SegmentURL mediaRange="5120-{len(PARENT) - 1}"/>
|
||||
</SegmentList>
|
||||
</Representation>
|
||||
</AdaptationSet>
|
||||
</Period>
|
||||
</MPD>"""
|
||||
|
||||
|
||||
def _mixed_mpd(srv) -> str:
|
||||
# Two SegmentURLs on distinct paths -> mixed URLs -> predicate False -> segmented path.
|
||||
return f"""<?xml version="1.0"?>
|
||||
<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" profiles="urn:mpeg:dash:profile:isoff-on-demand:2011">
|
||||
<BaseURL>{_base(srv)}</BaseURL>
|
||||
<Period id="0">
|
||||
<AdaptationSet contentType="video" mimeType="video/mp4" lang="en">
|
||||
<Representation id="v0" codecs="avc1.640028" bandwidth="1000000" width="1920" height="1080">
|
||||
<SegmentList timescale="1">
|
||||
<Initialization sourceURL="init.bin"/>
|
||||
<SegmentURL media="part0.bin"/>
|
||||
<SegmentURL media="part1.bin"/>
|
||||
</SegmentList>
|
||||
</Representation>
|
||||
</AdaptationSet>
|
||||
</Period>
|
||||
</MPD>"""
|
||||
|
||||
|
||||
def _track(mpd: str, srv):
|
||||
return DASH.from_text(mpd, url=f"{_base(srv)}manifest.mpd").to_tracks(language="en").videos[0]
|
||||
|
||||
|
||||
def _ctx(tmp_path, name="out.mp4"):
|
||||
return DownloadContext(
|
||||
save_path=tmp_path / name,
|
||||
save_dir=tmp_path / "v0_segments",
|
||||
progress=partial(lambda **kw: None),
|
||||
)
|
||||
|
||||
|
||||
def test_collapse_downloads_parent_whole_in_one_direct_request(server, tmp_path):
|
||||
server.routes["/media.mp4"] = PARENT
|
||||
track = _track(_collapse_mpd(server), server)
|
||||
ctx = _ctx(tmp_path, "collapsed.mp4")
|
||||
|
||||
DASH.download_track(track, ctx)
|
||||
|
||||
# the whole parent lands at save_path, byte-for-byte
|
||||
assert ctx.save_path.read_bytes() == PARENT
|
||||
assert track.path == ctx.save_path
|
||||
# exactly one Range-less GET (the direct whole-resource download); the only other hit is
|
||||
# the ranged init probe done while parsing SegmentList, never a per-segment slice fetch.
|
||||
no_range = [path for path, rng in server.requests if rng is None]
|
||||
assert no_range == ["/media.mp4"]
|
||||
# collapse skips the segment dir + merge entirely: no dir, no leftover control files
|
||||
assert not ctx.save_dir.exists()
|
||||
assert not list(tmp_path.glob("*.!dev"))
|
||||
assert not list(tmp_path.glob("**/*.!dev"))
|
||||
|
||||
|
||||
def test_non_collapse_takes_segmented_merge_path(server, tmp_path):
|
||||
server.routes["/init.bin"] = INIT2
|
||||
server.routes["/part0.bin"] = PART0
|
||||
server.routes["/part1.bin"] = PART1
|
||||
track = _track(_mixed_mpd(server), server)
|
||||
ctx = _ctx(tmp_path, "merged.mp4")
|
||||
|
||||
DASH.download_track(track, ctx)
|
||||
|
||||
# merged output is init followed by both segments, in order
|
||||
assert ctx.save_path.read_bytes() == INIT2 + PART0 + PART1
|
||||
# both distinct segments were fetched -> genuinely the per-segment path, not a collapse
|
||||
fetched = {path for path, _ in server.requests}
|
||||
assert {"/part0.bin", "/part1.bin"} <= fetched
|
||||
# segment files are consumed and the segment dir is cleaned up after merge
|
||||
assert not list(ctx.save_dir.glob("*")) if ctx.save_dir.exists() else True
|
||||
@@ -0,0 +1,145 @@
|
||||
"""Cancel/abort hardening for the downloader.
|
||||
|
||||
Pins two teardown-safety invariants:
|
||||
|
||||
- ``_dispatch_parts`` must NOT finalize when the process-global ``DOWNLOAD_CANCELLED``
|
||||
fires mid-download: part workers return silently keeping partials, so every future
|
||||
completes without error, and finalizing would strip the ``.!dev`` control file and
|
||||
pass off a hole-filled pre-truncated file as complete on the next run.
|
||||
- A worker parked in the retry backoff wait (which can reach MAX_BACKOFF) must wake and
|
||||
exit as soon as the batch abort is set, instead of sleeping out the full delay and
|
||||
stalling shutdown or the merge.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import threading
|
||||
import time
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
import pytest
|
||||
from requests import Session
|
||||
|
||||
dl = importlib.import_module("unshackle.core.downloaders.requests")
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
BODY = b"x" * 4096
|
||||
|
||||
|
||||
class _Handler(BaseHTTPRequestHandler):
|
||||
def do_GET(self) -> None:
|
||||
server = self.server
|
||||
if server.always_status:
|
||||
self.send_response(server.always_status)
|
||||
self.send_header("Content-Length", "0")
|
||||
self.end_headers()
|
||||
return
|
||||
rng = self.headers.get("Range")
|
||||
if rng:
|
||||
start, end = (int(x) for x in rng.removeprefix("bytes=").split("-"))
|
||||
body = BODY[start : end + 1]
|
||||
self.send_response(206)
|
||||
self.send_header("Content-Range", f"bytes {start}-{end}/{len(BODY)}")
|
||||
else:
|
||||
body = BODY
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
if server.gate is not None:
|
||||
self.wfile.write(body[:64])
|
||||
self.wfile.flush()
|
||||
server.started.set()
|
||||
server.gate.wait(timeout=10)
|
||||
self.wfile.write(body[64:])
|
||||
else:
|
||||
self.wfile.write(body)
|
||||
|
||||
def log_message(self, *args): # keep the test output quiet
|
||||
pass
|
||||
|
||||
|
||||
class _Server(ThreadingHTTPServer):
|
||||
daemon_threads = True
|
||||
|
||||
def handle_error(self, request, client_address): # cancelled readers close mid-write
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server():
|
||||
srv = _Server(("127.0.0.1", 0), _Handler)
|
||||
srv.always_status = 0
|
||||
srv.gate = None
|
||||
srv.started = threading.Event()
|
||||
thread = threading.Thread(target=srv.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
yield srv
|
||||
if srv.gate is not None:
|
||||
srv.gate.set()
|
||||
srv.shutdown()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_cancel():
|
||||
dl.DOWNLOAD_CANCELLED.clear()
|
||||
yield
|
||||
dl.DOWNLOAD_CANCELLED.clear()
|
||||
|
||||
|
||||
def _url(server: _Server) -> str:
|
||||
host, port = server.server_address
|
||||
return f"http://{host}:{port}/seg.bin"
|
||||
|
||||
|
||||
def test_dispatch_parts_does_not_finalize_on_global_cancel(server, tmp_path):
|
||||
server.gate = threading.Event()
|
||||
save_path = tmp_path / "big.bin"
|
||||
|
||||
def _cancel_when_started() -> None:
|
||||
assert server.started.wait(timeout=10)
|
||||
dl.DOWNLOAD_CANCELLED.set()
|
||||
server.gate.set()
|
||||
|
||||
canceller = threading.Thread(target=_cancel_when_started)
|
||||
canceller.start()
|
||||
events = list(
|
||||
dl._dispatch_parts(
|
||||
url=_url(server),
|
||||
save_path=save_path,
|
||||
session=Session(),
|
||||
total_size=len(BODY),
|
||||
max_workers=2,
|
||||
)
|
||||
)
|
||||
canceller.join(timeout=10)
|
||||
|
||||
assert not any("file_downloaded" in ev for ev in events)
|
||||
# the control file must survive so the next run restarts instead of trusting the file
|
||||
assert save_path.with_name(f"{save_path.name}.!dev").exists()
|
||||
|
||||
|
||||
def test_backoff_wait_exits_promptly_on_batch_abort(server, tmp_path, monkeypatch):
|
||||
server.always_status = 503
|
||||
monkeypatch.setattr(dl, "_retry_sleep", lambda exc, attempts: 30.0)
|
||||
abort = threading.Event()
|
||||
|
||||
def _consume() -> None:
|
||||
for _ in dl.download(
|
||||
url=_url(server),
|
||||
save_path=tmp_path / "seg.bin",
|
||||
session=Session(),
|
||||
segmented=True,
|
||||
abort=abort,
|
||||
):
|
||||
pass
|
||||
|
||||
worker = threading.Thread(target=_consume)
|
||||
start = time.monotonic()
|
||||
worker.start()
|
||||
time.sleep(0.5) # let the first attempt fail and enter the backoff wait
|
||||
abort.set()
|
||||
worker.join(timeout=10)
|
||||
assert not worker.is_alive()
|
||||
# generous bound: far below the 30s nap, so the wait was interrupted, not slept out
|
||||
assert time.monotonic() - start < 10
|
||||
@@ -9,7 +9,10 @@ media requests must ask for identity encoding.
|
||||
|
||||
import importlib
|
||||
import threading
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from email.utils import format_datetime
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -26,6 +29,12 @@ class _Handler(BaseHTTPRequestHandler):
|
||||
with server.lock:
|
||||
server.requests.append(dict(self.headers))
|
||||
nth = len(server.requests)
|
||||
if nth <= server.fail_429_first:
|
||||
self.send_response(429)
|
||||
self.send_header("Retry-After", "0")
|
||||
self.send_header("Content-Length", "0")
|
||||
self.end_headers()
|
||||
return
|
||||
rng = self.headers.get("Range")
|
||||
if server.honor_range and rng:
|
||||
start, end = (int(x) for x in rng.removeprefix("bytes=").split("-"))
|
||||
@@ -53,6 +62,7 @@ def server():
|
||||
srv.requests = []
|
||||
srv.honor_range = True
|
||||
srv.truncate_first = False
|
||||
srv.fail_429_first = 0
|
||||
thread = threading.Thread(target=srv.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
yield srv
|
||||
@@ -120,3 +130,64 @@ def test_range_item_retry_rewrites_slice(server, tmp_path, monkeypatch):
|
||||
def test_segment_requests_ask_identity_encoding(server, tmp_path):
|
||||
_run(server, tmp_path, [{"url": _url(server)}], max_workers=1)
|
||||
assert all(r.get("Accept-Encoding") == "identity" for r in server.requests)
|
||||
|
||||
|
||||
def _exc_with_headers(headers):
|
||||
exc = Exception("boom")
|
||||
exc.response = SimpleNamespace(headers=headers)
|
||||
return exc
|
||||
|
||||
|
||||
def test_retry_sleep_honors_numeric_retry_after():
|
||||
assert dl._retry_sleep(_exc_with_headers({"Retry-After": "3"}), 1) == 3.0
|
||||
assert dl._retry_sleep(_exc_with_headers({"Retry-After": "9999"}), 1) == 60.0 # session.MAX_BACKOFF cap
|
||||
|
||||
|
||||
def test_retry_sleep_honors_http_date_retry_after():
|
||||
when = format_datetime(datetime.now(timezone.utc) + timedelta(seconds=10), usegmt=True)
|
||||
assert 5 <= dl._retry_sleep(_exc_with_headers({"Retry-After": when}), 1) <= 10
|
||||
|
||||
|
||||
def test_retry_sleep_exponential_backoff_with_jitter():
|
||||
# no response on the failure: exponential from RETRY_WAIT with up to 10% jitter
|
||||
for attempts, base in ((1, dl.RETRY_WAIT), (2, dl.RETRY_WAIT * 2), (3, dl.RETRY_WAIT * 4)):
|
||||
assert base <= dl._retry_sleep(Exception("reset"), attempts) <= base * 1.1
|
||||
assert dl._retry_sleep(Exception("reset"), 30) == 60.0 # capped
|
||||
|
||||
|
||||
def test_retry_sleep_invalid_retry_after_falls_back():
|
||||
v = dl._retry_sleep(_exc_with_headers({"Retry-After": "soon"}), 1)
|
||||
assert dl.RETRY_WAIT <= v <= dl.RETRY_WAIT * 1.1
|
||||
|
||||
|
||||
def test_retry_sleep_naive_http_date_treated_as_utc():
|
||||
# "-0000" makes parsedate_to_datetime return a naive datetime; must clamp, not raise
|
||||
when = format_datetime(datetime.now(timezone.utc).replace(tzinfo=None) + timedelta(seconds=10))
|
||||
assert when.endswith("-0000")
|
||||
v = dl._retry_sleep(_exc_with_headers({"Retry-After": when}), 1)
|
||||
assert 5 <= v <= 10
|
||||
|
||||
|
||||
def test_retry_sleep_nan_retry_after_falls_back():
|
||||
# float("nan") parses; a non-finite wait must fall back to the exponential path
|
||||
v = dl._retry_sleep(_exc_with_headers({"Retry-After": "nan"}), 1)
|
||||
assert dl.RETRY_WAIT <= v <= dl.RETRY_WAIT * 1.1
|
||||
|
||||
|
||||
def test_retry_sleep_honors_retry_after_on_wrapped_cause():
|
||||
# RnetSession raises MaxRetriesError whose __cause__ HTTPError carries the response
|
||||
exc = Exception("max retries exceeded")
|
||||
exc.__cause__ = _exc_with_headers({"Retry-After": "3"})
|
||||
assert dl._retry_sleep(exc, 1) == 3.0
|
||||
|
||||
|
||||
def test_transient_429_with_retry_after_does_not_abort_track(server, tmp_path, monkeypatch):
|
||||
# a short rate-limit window must retry the segment, not kill the whole batch
|
||||
dl.DOWNLOAD_CANCELLED.clear()
|
||||
monkeypatch.setattr(dl, "RETRY_WAIT", 0.01)
|
||||
server.fail_429_first = 2
|
||||
files = _run(server, tmp_path, [{"url": _url(server, f"/seg{i}.bin")} for i in range(3)], max_workers=1)
|
||||
assert len(files) == 3
|
||||
assert all(f.read_bytes() == PARENT for f in files)
|
||||
assert not dl.DOWNLOAD_CANCELLED.is_set()
|
||||
assert len(server.requests) == 5 # 3 successes + 2 rate-limited attempts
|
||||
|
||||
@@ -0,0 +1,156 @@
|
||||
"""Multiprocess segment-download invariants.
|
||||
|
||||
Pins the spawned-child fan-out path in ``requests()``/``_download_multiprocess``:
|
||||
|
||||
- Strided assignment plus the ``index_offset``/``index_stride`` filename mapping puts every
|
||||
segment's bytes in the correctly named file (a broken stride mapping misfiles content).
|
||||
- The parent normalizes child progress to segment granularity: it drops child ``advance``
|
||||
and ``total``/speed events and emits exactly one ``advance=1`` per ``file_downloaded``, so
|
||||
the advances summed by a caller equal the segment count.
|
||||
- A child that dies (its worker sends ``__mp_error__``) surfaces as a ``RuntimeError`` from
|
||||
the generator without hanging.
|
||||
- A process-global ``DOWNLOAD_CANCELLED`` set mid-batch terminates the spawned children (which
|
||||
carry their own fresh flag and can't see it) and the teardown join reaps them promptly.
|
||||
|
||||
Children are spawned and re-import the module, so monkeypatched module attributes do NOT reach
|
||||
them; these tests drive real child processes against a localhost server on 127.0.0.1 (reachable
|
||||
from the children) and avoid any child-side patching. Segments are a few hundred bytes so the
|
||||
real spawns stay fast.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import multiprocessing
|
||||
import threading
|
||||
import time
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
import pytest
|
||||
|
||||
dl = importlib.import_module("unshackle.core.downloaders.requests")
|
||||
|
||||
SEG_REPEAT = 32 # 8-digit index * 32 = 256-byte body, distinct per index
|
||||
|
||||
|
||||
def _seg_body(index: int) -> bytes:
|
||||
"""Segment body that encodes its own index, so misfiled striding is detectable."""
|
||||
return (f"{index:08d}" * SEG_REPEAT).encode()
|
||||
|
||||
|
||||
class _Handler(BaseHTTPRequestHandler):
|
||||
def do_GET(self) -> None:
|
||||
server = self.server
|
||||
delay = getattr(server, "delay", 0.0)
|
||||
if delay:
|
||||
time.sleep(delay)
|
||||
try:
|
||||
index = int(self.path.rsplit("/", 1)[1].split(".")[0])
|
||||
except ValueError:
|
||||
self.send_response(404)
|
||||
self.send_header("Content-Length", "0")
|
||||
self.end_headers()
|
||||
return
|
||||
body = _seg_body(index)
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def log_message(self, *args): # keep the test output quiet
|
||||
pass
|
||||
|
||||
|
||||
class _Server(ThreadingHTTPServer):
|
||||
daemon_threads = True
|
||||
|
||||
def handle_error(self, request, client_address): # a cancelled child closes its socket mid-write
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server():
|
||||
srv = _Server(("127.0.0.1", 0), _Handler)
|
||||
srv.delay = 0.0
|
||||
thread = threading.Thread(target=srv.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
yield srv
|
||||
srv.shutdown()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_cancel():
|
||||
# spawned children carry their own flag, but the parent loop reads this process's flag;
|
||||
# keep it clean so cancel state never leaks between tests
|
||||
dl.DOWNLOAD_CANCELLED.clear()
|
||||
yield
|
||||
dl.DOWNLOAD_CANCELLED.clear()
|
||||
|
||||
|
||||
def _url(srv, path):
|
||||
return f"http://127.0.0.1:{srv.server_address[1]}{path}"
|
||||
|
||||
|
||||
def test_multiprocess_stride_downloads_all_segments_correctly(server, tmp_path):
|
||||
# >= MP_MIN_SEGMENTS with processes=2 engages the fan-out via the public entry point
|
||||
n = dl.MP_MIN_SEGMENTS
|
||||
urls = [{"url": _url(server, f"/seg/{i}.bin")} for i in range(n)]
|
||||
|
||||
total_advance = 0
|
||||
for event in dl.requests(urls, output_dir=tmp_path, filename="seg_{i:04}.bin", max_workers=4, processes=2):
|
||||
total_advance += event.get("advance", 0)
|
||||
|
||||
files = sorted(tmp_path.glob("seg_*.bin"))
|
||||
assert len(files) == n
|
||||
for f in files:
|
||||
index = int(f.stem.split("_")[1])
|
||||
assert f.read_bytes() == _seg_body(index), f"segment {index} holds another index's bytes (stride mapping bug)"
|
||||
# parent drops child advances and emits one advance=1 per file_downloaded
|
||||
assert total_advance == n
|
||||
|
||||
|
||||
def test_multiprocess_child_crash_raises(server, tmp_path):
|
||||
# _download_multiprocess has no MP_MIN_SEGMENTS gate, so drive it directly with 2 urls.
|
||||
# a dict url without "url" raises KeyError while requests() builds the save path (before
|
||||
# any retry loop), so the child sends __mp_error__ fast; the sibling downloads normally.
|
||||
good = {"url": _url(server, "/seg/0.bin")}
|
||||
bad = {"nourl": True}
|
||||
with pytest.raises(RuntimeError):
|
||||
for _ in dl._download_multiprocess(
|
||||
urls=[good, bad],
|
||||
output_dir=tmp_path,
|
||||
filename="seg_{i:04}.bin",
|
||||
headers=None,
|
||||
cookies=None,
|
||||
proxy=None,
|
||||
max_workers=2,
|
||||
adaptive=False,
|
||||
spec={"kind": "none"},
|
||||
processes=2,
|
||||
debug_logger=None,
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
def test_multiprocess_cancel_terminates_children(server, tmp_path):
|
||||
server.delay = 0.5 # slow segments so cancel lands well before the batch could finish
|
||||
n = dl.MP_MIN_SEGMENTS
|
||||
urls = [{"url": _url(server, f"/seg/{i}.bin")} for i in range(n)]
|
||||
|
||||
baseline = len(multiprocessing.active_children())
|
||||
start = time.time()
|
||||
progressed = 0
|
||||
for event in dl.requests(urls, output_dir=tmp_path, filename="seg_{i:04}.bin", max_workers=4, processes=2):
|
||||
if event.get("advance") or event.get("written"):
|
||||
progressed += 1
|
||||
if progressed >= 2:
|
||||
dl.DOWNLOAD_CANCELLED.set() # sibling-track cancel: parent must tear the children down
|
||||
elapsed = time.time() - start
|
||||
|
||||
# cancel stopped the work: nowhere near all segments finished, and no hang
|
||||
assert len(list(tmp_path.glob("seg_*.bin"))) < n
|
||||
assert elapsed < 15
|
||||
# the finally-block join reaps the spawned children promptly
|
||||
deadline = time.time() + 5
|
||||
while time.time() < deadline and len(multiprocessing.active_children()) > baseline:
|
||||
time.sleep(0.05)
|
||||
assert len(multiprocessing.active_children()) <= baseline
|
||||
@@ -0,0 +1,164 @@
|
||||
"""Ranged-parallel single-URL download invariants.
|
||||
|
||||
Pins the byte-range fan-out that splits one large file across parts:
|
||||
|
||||
- a range-capable server above the size threshold is split into several byte
|
||||
windows, the parts merge byte-identical, and the ``.!dev`` control marker is
|
||||
removed on success;
|
||||
- a ranged part that receives a 200 (server ignored the Range) fails after
|
||||
retries rather than writing a wrong-length body into the pre-truncated file;
|
||||
- that failure sets only the ranged download's LOCAL abort event, never the
|
||||
process-global DOWNLOAD_CANCELLED (which would poison sibling tracks);
|
||||
- ``_probe_ranged`` parses the Content-Range total on a 206 and declines a 200
|
||||
or a content-encoded 206.
|
||||
|
||||
Tests 2/3 exercise ``_dispatch_parts`` directly: the public single-URL path in
|
||||
``requests()`` deliberately swallows a ranged failure and falls back to a plain
|
||||
sequential download, so the raise the invariant guarantees is only observable at
|
||||
the dispatch layer.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import threading
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
import pytest
|
||||
from requests import Session
|
||||
|
||||
dl = importlib.import_module("unshackle.core.downloaders.requests")
|
||||
|
||||
PAYLOAD = bytes(range(256)) * 128 # 32 KiB, distinctive repeating byte pattern
|
||||
|
||||
|
||||
class _RangeHandler(BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
srv = self.server
|
||||
rng = self.headers.get("Range")
|
||||
with srv.lock:
|
||||
srv.requests.append(rng)
|
||||
mode = srv.mode
|
||||
|
||||
if mode == "gzip_probe":
|
||||
# a 206 that (wrongly) carries a content-encoding: the probe must decline it
|
||||
self.send_response(206)
|
||||
self.send_header("Content-Range", f"bytes 0-0/{len(PAYLOAD)}")
|
||||
self.send_header("Content-Encoding", "gzip")
|
||||
self.send_header("Content-Length", "1")
|
||||
self.end_headers()
|
||||
self.wfile.write(b"\x00")
|
||||
return
|
||||
|
||||
# always_200 ignores Range entirely; probe_only_206 answers the 0-0 probe
|
||||
# with a real 206 but hands every wider part window a 200 full body
|
||||
serve_full = mode == "always_200" or not rng or (mode == "probe_only_206" and rng != "bytes=0-0")
|
||||
if serve_full:
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Length", str(len(PAYLOAD)))
|
||||
self.end_headers()
|
||||
self.wfile.write(PAYLOAD)
|
||||
return
|
||||
|
||||
start, end = (int(x) for x in rng.removeprefix("bytes=").split("-"))
|
||||
body = PAYLOAD[start : end + 1]
|
||||
self.send_response(206)
|
||||
self.send_header("Content-Range", f"bytes {start}-{end}/{len(PAYLOAD)}")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_cancelled():
|
||||
# ranged failures must never touch the process-global flag; isolate every test from it
|
||||
dl.DOWNLOAD_CANCELLED.clear()
|
||||
yield
|
||||
dl.DOWNLOAD_CANCELLED.clear()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server():
|
||||
srv = ThreadingHTTPServer(("127.0.0.1", 0), _RangeHandler)
|
||||
srv.lock = threading.Lock()
|
||||
srv.requests = []
|
||||
srv.mode = "range" # honor every Range with a proper 206
|
||||
thread = threading.Thread(target=srv.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
yield srv
|
||||
srv.shutdown()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
def _url(srv, path="/payload.bin"):
|
||||
return f"http://127.0.0.1:{srv.server_address[1]}{path}"
|
||||
|
||||
|
||||
def _run(srv, tmp_path, urls, **kwargs):
|
||||
for _ in dl.requests(urls, output_dir=tmp_path, filename="seg_{i:04}.bin", **kwargs):
|
||||
pass
|
||||
return sorted(tmp_path.glob("seg_*.bin"))
|
||||
|
||||
|
||||
def test_ranged_parallel_merges_byte_identical(server, tmp_path, monkeypatch):
|
||||
# shrink the thresholds so a 32 KiB payload splits into several parts
|
||||
monkeypatch.setattr(dl, "RANGE_PARALLEL_MIN_SIZE", 8 * 1024)
|
||||
monkeypatch.setattr(dl, "RANGE_PARALLEL_PART_SIZE", 2 * 1024)
|
||||
files = _run(server, tmp_path, [{"url": _url(server)}], max_workers=4)
|
||||
assert len(files) == 1
|
||||
assert files[0].read_bytes() == PAYLOAD
|
||||
# the split must actually have happened: >1 distinct window beyond the 0-0 probe
|
||||
windows = {r for r in server.requests if r and r != "bytes=0-0"}
|
||||
assert len(windows) >= 2
|
||||
# control marker removed once the parts finalize
|
||||
assert not files[0].with_name(files[0].name + ".!dev").exists()
|
||||
|
||||
|
||||
def test_ranged_part_requires_206(server, tmp_path, monkeypatch):
|
||||
server.mode = "probe_only_206"
|
||||
monkeypatch.setattr(dl, "RETRY_WAIT", 0.01)
|
||||
save_path = tmp_path / "seg_0000.bin"
|
||||
gen = dl._dispatch_parts(
|
||||
url=_url(server),
|
||||
save_path=save_path,
|
||||
session=Session(),
|
||||
total_size=len(PAYLOAD),
|
||||
max_workers=4,
|
||||
)
|
||||
with pytest.raises(IOError, match="expected 206 for ranged part"):
|
||||
for _ in gen:
|
||||
pass
|
||||
# not finalized: the control marker survives (completed=False) and the target
|
||||
# is still the pre-truncated zero-fill, never the served bytes
|
||||
assert save_path.with_name(save_path.name + ".!dev").exists()
|
||||
assert save_path.read_bytes() != PAYLOAD
|
||||
|
||||
|
||||
def test_ranged_part_failure_stays_local(server, tmp_path, monkeypatch):
|
||||
server.mode = "probe_only_206"
|
||||
monkeypatch.setattr(dl, "RETRY_WAIT", 0.01)
|
||||
dl.DOWNLOAD_CANCELLED.clear()
|
||||
save_path = tmp_path / "seg_0000.bin"
|
||||
gen = dl._dispatch_parts(
|
||||
url=_url(server),
|
||||
save_path=save_path,
|
||||
session=Session(),
|
||||
total_size=len(PAYLOAD),
|
||||
max_workers=4,
|
||||
)
|
||||
with pytest.raises(IOError, match="expected 206 for ranged part"):
|
||||
for _ in gen:
|
||||
pass
|
||||
# a failed part rides the local abort event; the global cancel must stay clear
|
||||
assert not dl.DOWNLOAD_CANCELLED.is_set()
|
||||
|
||||
|
||||
def test_probe_ranged_parses_content_range(server):
|
||||
session = Session()
|
||||
server.mode = "range"
|
||||
assert dl._probe_ranged(_url(server), session) == (len(PAYLOAD), True)
|
||||
server.mode = "always_200"
|
||||
assert dl._probe_ranged(_url(server), session) == (0, False)
|
||||
server.mode = "gzip_probe"
|
||||
assert dl._probe_ranged(_url(server), session) == (0, False)
|
||||
@@ -0,0 +1,175 @@
|
||||
"""Tail-boost success-path invariants.
|
||||
|
||||
Pins the tail-boost engage path in ``requests()``: when only a few segments remain
|
||||
and workers would otherwise idle, each remaining segment is probed for range support
|
||||
and split into intra-segment byte-range parts written into a pre-truncated ``.tp.!dev``
|
||||
target, finalized by the last part worker via os.replace. The invariants:
|
||||
|
||||
- a boosted segment's final file is byte-identical to the served content, every
|
||||
segment reports exactly one advance=1 (part-mode byte advances are swallowed), and
|
||||
a clean run leaves no ``.!dev`` / ``.tp.!dev`` strays;
|
||||
- finalize is last-part-wins: on success save_path exists and the ``.tp.!dev`` target
|
||||
is gone;
|
||||
- a permanently failing part fails the batch and the boosted segment is never
|
||||
os.replace'd into place (no corrupt finalized segment).
|
||||
|
||||
Determinism: with ``max_workers`` < segment count <= 2*max_workers and adaptive=True,
|
||||
the leading ``max_workers`` segments submit upfront and the trailing few stay in
|
||||
``remaining`` (the tail top-up only releases boost-declined indices). The controller
|
||||
starts at the cap and never grows, so as soon as two leaders finish the spare workers
|
||||
outnumber the tail and ``_tail_boost_engages`` fires. Module timing/size constants are
|
||||
lowered via monkeypatch; the whole batch runs in-process (threads), so the patched
|
||||
values are seen by the workers.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import threading
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
import pytest
|
||||
|
||||
dl = importlib.import_module("unshackle.core.downloaders.requests")
|
||||
|
||||
BODY = bytes((i * 7 + 11) % 256 for i in range(16384)) # 16 KiB, distinctive, range-sliceable
|
||||
|
||||
# lowered so a 16 KiB segment clears the min and splits into several 4 KiB parts
|
||||
BOOST_MIN = 4096
|
||||
BOOST_PART = 4096
|
||||
|
||||
MAX_WORKERS = 6
|
||||
SEG_COUNT = 8 # > MAX_WORKERS and <= 2*MAX_WORKERS: leaders 0-5 submit, tail 6,7 boost
|
||||
|
||||
|
||||
class _Handler(BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
server = self.server
|
||||
rng = self.headers.get("Range")
|
||||
with server.lock:
|
||||
server.requests.append((self.path, rng))
|
||||
if rng and rng.startswith("bytes="):
|
||||
start_s, _, end_s = rng[len("bytes=") :].partition("-")
|
||||
start = int(start_s)
|
||||
end = int(end_s) if end_s else len(BODY) - 1
|
||||
if server.break_offset_ranges and start > 0:
|
||||
# misbehave for every non-leading window: send the whole body as 200 so
|
||||
# part-mode's 206 check fails permanently for that part
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Length", str(len(BODY)))
|
||||
self.end_headers()
|
||||
self.wfile.write(BODY)
|
||||
return
|
||||
body = BODY[start : end + 1]
|
||||
self.send_response(206)
|
||||
self.send_header("Content-Range", f"bytes {start}-{end}/{len(BODY)}")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
return
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Length", str(len(BODY)))
|
||||
self.end_headers()
|
||||
self.wfile.write(BODY)
|
||||
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server():
|
||||
srv = ThreadingHTTPServer(("127.0.0.1", 0), _Handler)
|
||||
srv.lock = threading.Lock()
|
||||
srv.requests = []
|
||||
srv.break_offset_ranges = False
|
||||
thread = threading.Thread(target=srv.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
yield srv
|
||||
srv.shutdown()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clean_cancel():
|
||||
# a failing batch sets the process-global cancel; keep tests independent
|
||||
dl.DOWNLOAD_CANCELLED.clear()
|
||||
yield
|
||||
dl.DOWNLOAD_CANCELLED.clear()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _boost_constants(monkeypatch):
|
||||
monkeypatch.setattr(dl, "TAIL_BOOST_MIN_SEGMENT_SIZE", BOOST_MIN)
|
||||
monkeypatch.setattr(dl, "TAIL_BOOST_PART_SIZE", BOOST_PART)
|
||||
|
||||
|
||||
def _url(srv, path):
|
||||
return f"http://127.0.0.1:{srv.server_address[1]}{path}"
|
||||
|
||||
|
||||
def _urls(srv):
|
||||
return [{"url": _url(srv, f"/seg{i}.bin")} for i in range(SEG_COUNT)]
|
||||
|
||||
|
||||
def _run(srv, tmp_path, urls):
|
||||
advances = 0
|
||||
for ev in dl.requests(urls, output_dir=tmp_path, filename="seg_{i:04}.bin", max_workers=MAX_WORKERS, adaptive=True):
|
||||
a = ev.get("advance")
|
||||
if a:
|
||||
advances += a
|
||||
return sorted(tmp_path.glob("seg_*.bin")), advances
|
||||
|
||||
|
||||
def _ranged_counts(srv):
|
||||
with srv.lock:
|
||||
reqs = list(srv.requests)
|
||||
counts: dict[str, int] = {}
|
||||
for path, rng in reqs:
|
||||
if rng:
|
||||
counts[path] = counts.get(path, 0) + 1
|
||||
return counts
|
||||
|
||||
|
||||
def _boost_engaged(srv):
|
||||
# a boosted path sees a 0-0 probe plus >=2 part windows; a normal segment sees no Range
|
||||
return any(n > 1 for n in _ranged_counts(srv).values())
|
||||
|
||||
|
||||
def test_tail_boost_split_produces_byte_identical_segment(server, tmp_path):
|
||||
files, advances = _run(server, tmp_path, _urls(server))
|
||||
|
||||
assert _boost_engaged(server), "tail boost never engaged; test would silently cover the normal path"
|
||||
assert len(files) == SEG_COUNT
|
||||
assert all(f.read_bytes() == BODY for f in files)
|
||||
# part-mode byte advances are swallowed; each segment reports exactly one advance=1
|
||||
assert advances == SEG_COUNT
|
||||
# a clean run renames every target into place, so no partial/part strays survive
|
||||
assert not list(tmp_path.glob("*.!dev"))
|
||||
assert not list(tmp_path.glob("*.tp.!dev"))
|
||||
|
||||
|
||||
def test_tail_boost_finalize_last_part_wins(server, tmp_path):
|
||||
files, _ = _run(server, tmp_path, _urls(server))
|
||||
|
||||
assert _boost_engaged(server), "tail boost never engaged; finalize path untested"
|
||||
assert len(files) == SEG_COUNT
|
||||
# last part wins: the pre-truncated .tp.!dev is os.replace'd into save_path, leaving no target
|
||||
assert not list(tmp_path.glob("*.tp.!dev"))
|
||||
for path in _ranged_counts(server):
|
||||
i = int(path.removeprefix("/seg").removesuffix(".bin"))
|
||||
save_path = tmp_path / f"seg_{i:04}.bin"
|
||||
assert save_path.exists()
|
||||
assert not save_path.with_name(f"{save_path.name}.tp.!dev").exists()
|
||||
|
||||
|
||||
def test_tail_boost_part_failure_fails_batch_but_keeps_siblings_files_clean(server, tmp_path, monkeypatch):
|
||||
server.break_offset_ranges = True # every non-leading part window 200s -> permanent 206 failure
|
||||
monkeypatch.setattr(dl, "RETRY_WAIT", 0.01)
|
||||
|
||||
with pytest.raises(Exception):
|
||||
_run(server, tmp_path, _urls(server))
|
||||
|
||||
# the boosted tail segments (6, 7) hit the failing windows: no part run reaches parts_left==0
|
||||
# cleanly, so neither is finalized. sibling leaders that did finalize must be uncorrupted.
|
||||
for i in (6, 7):
|
||||
assert not (tmp_path / f"seg_{i:04}.bin").exists()
|
||||
for f in tmp_path.glob("seg_*.bin"):
|
||||
assert f.read_bytes() == BODY
|
||||
@@ -1,6 +1,7 @@
|
||||
import math
|
||||
import multiprocessing
|
||||
import os
|
||||
import random
|
||||
import socket
|
||||
import statistics
|
||||
import sys
|
||||
@@ -10,6 +11,8 @@ import traceback
|
||||
from collections import deque
|
||||
from concurrent.futures import FIRST_COMPLETED, wait
|
||||
from concurrent.futures.thread import ThreadPoolExecutor
|
||||
from datetime import datetime, timezone
|
||||
from email.utils import parsedate_to_datetime
|
||||
from http.cookiejar import CookieJar
|
||||
from pathlib import Path
|
||||
from queue import Empty, Queue
|
||||
@@ -34,9 +37,6 @@ READ_TIMEOUT = 30
|
||||
HEDGE_FACTOR = 3
|
||||
HEDGE_MIN_WAIT = 5.0
|
||||
|
||||
# racers read per network arrival (read1) so superseded hedge losers exit promptly
|
||||
RACER_READ1 = True
|
||||
|
||||
# Adaptive chunk sizing — benchmarked optimal range
|
||||
MIN_CHUNK = 524_288 # 512KB
|
||||
MAX_CHUNK = 4_194_304 # 4MB
|
||||
@@ -207,6 +207,43 @@ class AdaptiveWorkerController:
|
||||
return self.target
|
||||
|
||||
|
||||
def _retry_sleep(exc: Exception, attempts: int) -> float:
|
||||
"""Seconds to wait before retry number ``attempts``.
|
||||
|
||||
download() owns segment retries (session-level retries are suppressed), so backoff
|
||||
lives here: honor the server's Retry-After when the failure carries a response,
|
||||
else exponential growth from RETRY_WAIT with jitter so concurrently rate-limited
|
||||
segments don't retry in lockstep. Capped at session.MAX_BACKOFF either way.
|
||||
"""
|
||||
from unshackle.core.session import MAX_BACKOFF
|
||||
|
||||
response = getattr(exc, "response", None)
|
||||
if response is None:
|
||||
# RnetSession wraps forcelist failures in MaxRetriesError with the HTTPError as __cause__
|
||||
response = getattr(getattr(exc, "__cause__", None), "response", None)
|
||||
if response is not None:
|
||||
try:
|
||||
retry_after = response.headers.get("Retry-After")
|
||||
except Exception:
|
||||
retry_after = None
|
||||
if retry_after:
|
||||
wait: Optional[float] = None
|
||||
try:
|
||||
wait = float(retry_after)
|
||||
except ValueError:
|
||||
try:
|
||||
retry_date = parsedate_to_datetime(retry_after)
|
||||
if retry_date.tzinfo is None:
|
||||
retry_date = retry_date.replace(tzinfo=timezone.utc)
|
||||
wait = (retry_date - datetime.now(timezone.utc)).total_seconds()
|
||||
except Exception:
|
||||
wait = None
|
||||
if wait is not None and math.isfinite(wait):
|
||||
return min(max(wait, 0.0), MAX_BACKOFF)
|
||||
backoff = RETRY_WAIT * (2 ** (attempts - 1))
|
||||
return min(backoff + random.uniform(0, backoff * 0.1), MAX_BACKOFF)
|
||||
|
||||
|
||||
def _is_requests_session(session: Any) -> bool:
|
||||
"""Check if the session is a standard requests.Session (supports resp.raw)."""
|
||||
return isinstance(session, Session)
|
||||
@@ -396,6 +433,11 @@ def _dispatch_parts(
|
||||
total_bytes += advance
|
||||
yield {"advance": advance}
|
||||
|
||||
if DOWNLOAD_CANCELLED.is_set() or abort.is_set():
|
||||
# a cross-track cancel makes part workers return silently with partials, so
|
||||
# every future completes without error; finalizing here would strip the .!dev
|
||||
# control file and pass off a hole-filled pre-truncated file as complete
|
||||
return
|
||||
yield {"file_downloaded": save_path, "written": total_size}
|
||||
completed = True
|
||||
except KeyboardInterrupt:
|
||||
@@ -450,12 +492,11 @@ 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).
|
||||
racing: Optional predicate selecting read granularity per iteration. When None,
|
||||
read1 (per-network-arrival reads) is used whenever `claimed` is set and
|
||||
RACER_READ1 is on. When provided, read1 is used only
|
||||
on iterations where it returns True (this segment is being hedge-raced, so a
|
||||
superseded loser must notice `claimed()` mid-stream); otherwise a blocking
|
||||
full-chunk read is used, avoiding the per-record read tax on non-racing segments.
|
||||
racing: Optional predicate selecting read granularity per iteration. Requires
|
||||
`claimed`. read1 (per-network-arrival reads) is used only on iterations where
|
||||
it returns True (this segment is being hedge-raced, so a superseded loser must
|
||||
notice `claimed()` mid-stream); otherwise a blocking full-chunk read is used,
|
||||
avoiding the per-record read tax on non-racing segments.
|
||||
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.
|
||||
@@ -504,6 +545,10 @@ def download(
|
||||
while True:
|
||||
if claimed is not None and claimed():
|
||||
return
|
||||
if DOWNLOAD_CANCELLED.is_set() or (abort is not None and abort.is_set()):
|
||||
# a worker waking from a retry nap after cancel/batch-abort must exit here,
|
||||
# before opening a new request or the .!dev handle mid-teardown
|
||||
return
|
||||
if not part_mode:
|
||||
written = 0
|
||||
last_speed_refresh = _time()
|
||||
@@ -598,14 +643,12 @@ def download(
|
||||
if use_rnet:
|
||||
chunks = stream.stream()
|
||||
elif use_raw:
|
||||
_read1 = getattr(stream.raw, "read1", None) if claimed is not None and RACER_READ1 else None
|
||||
_read1 = getattr(stream.raw, "read1", None) if claimed is not None else None
|
||||
if _read1 is not None and racing is not None:
|
||||
# only pay the per-arrival read1 tax while this segment is actually
|
||||
# racing; a blocking full-chunk read otherwise (the common case)
|
||||
_raw_read = stream.raw.read
|
||||
chunks = iter(lambda: _read1(chunk_size) if racing() else _raw_read(chunk_size), b"")
|
||||
elif _read1 is not None:
|
||||
chunks = iter(lambda: _read1(chunk_size), b"")
|
||||
else:
|
||||
chunks = iter(lambda: stream.raw.read(chunk_size), b"")
|
||||
else:
|
||||
@@ -701,7 +744,15 @@ def download(
|
||||
pass
|
||||
if not part_mode:
|
||||
resume_offset = tmp_file.stat().st_size if tmp_file.exists() else 0
|
||||
time.sleep(RETRY_WAIT)
|
||||
delay = _retry_sleep(exc, attempts)
|
||||
if abort is not None:
|
||||
# interruptible nap: backoff can reach MAX_BACKOFF, and teardown always
|
||||
# sets the batch abort, so a parked worker must wake and exit at once
|
||||
# instead of stalling shutdown (or the merge) for the full delay
|
||||
if abort.wait(delay):
|
||||
return
|
||||
else:
|
||||
time.sleep(delay)
|
||||
attempts += 1
|
||||
|
||||
|
||||
@@ -774,7 +825,10 @@ def _mp_worker(queue: Any, kwargs: dict[str, Any]) -> None:
|
||||
if overrides:
|
||||
for pair in overrides.split(","):
|
||||
name, _, value = pair.partition("=")
|
||||
setattr(sys.modules[__name__], name, float(value))
|
||||
# only the timing constants the bench relays; a stale/foreign env var
|
||||
# must not be able to rewrite arbitrary module globals in every child
|
||||
if name in ("READ_TIMEOUT", "RETRY_WAIT", "HEDGE_MIN_WAIT"):
|
||||
setattr(sys.modules[__name__], name, float(value))
|
||||
spec = kwargs.pop("_session_spec")
|
||||
kwargs["session"] = _rebuild_session(spec)
|
||||
for event in requests(**kwargs):
|
||||
@@ -866,6 +920,13 @@ def _download_multiprocess(
|
||||
dead_ticks = 0
|
||||
try:
|
||||
while done_count < len(procs):
|
||||
if DOWNLOAD_CANCELLED.is_set():
|
||||
# cross-track cancel (e.g. a sibling track failed): spawned children carry
|
||||
# their own fresh flag and can't see it, so terminate them like KeyboardInterrupt
|
||||
for p in procs:
|
||||
p.terminate()
|
||||
yield dict(downloaded="[yellow]CANCELLED")
|
||||
return
|
||||
try:
|
||||
event = queue.get(timeout=0.1)
|
||||
except Empty:
|
||||
@@ -893,9 +954,16 @@ def _download_multiprocess(
|
||||
continue # parent owns the aggregate total
|
||||
if "downloaded" in event:
|
||||
continue # child speed/status string; parent computes its own aggregate speed
|
||||
if "advance" in event:
|
||||
# children report mixed granularity (segment counts, or raw bytes when a
|
||||
# length-1 strided chunk routes through the single-file path); the parent
|
||||
# bar counts segments, so advance=1 is emitted per file_downloaded instead
|
||||
continue
|
||||
written = event.get("written")
|
||||
if written:
|
||||
total_bytes += written
|
||||
if "file_downloaded" in event:
|
||||
yield dict(advance=1)
|
||||
yield event
|
||||
|
||||
now = time.time()
|
||||
|
||||
@@ -701,7 +701,7 @@ class RnetSession:
|
||||
|
||||
if response.status_code not in self.status_forcelist:
|
||||
return response
|
||||
last_exception = HTTPError(f"Received status code: {response.status_code}")
|
||||
last_exception = HTTPError(f"Received status code: {response.status_code}", response=response)
|
||||
self.log.warning(
|
||||
f"{response.status_code} {response.reason}({urlparse(url).path}). Retrying... "
|
||||
f"({attempt + 1}/{max_retries})"
|
||||
|
||||
Reference in New Issue
Block a user