diff --git a/scripts/bench_downloader.py b/scripts/bench_downloader.py index c44ed231..1a4e28de 100644 --- a/scripts/bench_downloader.py +++ b/scripts/bench_downloader.py @@ -29,6 +29,7 @@ from __future__ import annotations import http.server import importlib import math +import multiprocessing import os import socket import socketserver @@ -223,7 +224,18 @@ def run_batch( except BaseException as exc: # noqa: BLE001 - bench records failure, keeps sweeping err = exc wall = time.perf_counter() - start - # measured immediately at completion: this is the exact moment merge would sweep + if err is not None: + # a failed batch shuts down without waiting, so a leaked racer may sit in one last + # blocked read (<= READ_TIMEOUT) holding its .!dev handle. merge never sweeps after + # a failure; what races the leak is the cleanup that runs once the error propagates, + # so drain before measuring. a handle still held past the deadline is a real leak. + deadline = time.monotonic() + dl_mod.READ_TIMEOUT + dl_mod.RETRY_WAIT + 2 + while time.monotonic() < deadline and ( + any(t.name.startswith("ThreadPoolExecutor") for t in threading.enumerate()) + or multiprocessing.active_children() + ): + time.sleep(0.1) + # on success this is the exact moment merge would sweep strays, locked = sweep_strays(out_dir) pool_threads = sum(1 for t in threading.enumerate() if t.name.startswith("ThreadPoolExecutor")) return wall, completions, err, strays, locked, pool_threads @@ -319,7 +331,9 @@ def sweep( 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) bad = incomplete_segments(out_dir, len(urls), expected_bytes) - samples.append((secs, comps, downloaded_mib(out_dir), bad, err is not None, strays, locked, pool_threads)) + samples.append( + (secs, comps, downloaded_mib(out_dir), bad, err is not None, 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 @@ -327,8 +341,19 @@ def sweep( mbps = best[2] / best[0] if best[0] else 0.0 # sweep safety reports the WORST run: one locked handle in any run is the bug rows.append( - (w, median_s, mbps, bucket, tail_s, tail_pct, best[3], best[4], - max(s[5] for s in samples), max(s[6] for s in samples), max(s[7] for s in samples)) + ( + w, + median_s, + mbps, + bucket, + tail_s, + tail_pct, + best[3], + best[4], + max(s[5] for s in samples), + max(s[6] for s in samples), + max(s[7] for s in samples), + ) ) return rows @@ -370,11 +395,17 @@ def print_report(rows: list[Row], runs: int, total_segments: int) -> None: @click.option("--seg-size", default=2 * MIB, show_default=True, help="Bytes per synthetic segment.") @click.option("--workers", default="1,2,4,8,16", show_default=True, help="Comma list of max_workers to sweep.") @click.option("--runs", default=1, show_default=True, help="Runs per worker count; median reported.") -@click.option("--latency", "latency_ms", default=0, show_default=True, help="Per-response delay before first byte (ms).") -@click.option("--rate", "rate_kib", default=0, show_default=True, help="Per-response throttle in KiB/s (0 = unthrottled).") +@click.option( + "--latency", "latency_ms", default=0, show_default=True, help="Per-response delay before first byte (ms)." +) +@click.option( + "--rate", "rate_kib", default=0, show_default=True, help="Per-response throttle in KiB/s (0 = unthrottled)." +) @click.option("--fault-stall", default=0, show_default=True, help="First N segments stall forever mid-body.") @click.option("--fault-reset", default=0, show_default=True, help="First N segments RST mid-body on attempt 1.") -@click.option("--fault-503", "fault_503", default=0, show_default=True, help="First N segments return 503 on attempt 1.") +@click.option( + "--fault-503", "fault_503", default=0, show_default=True, help="First N segments return 503 on attempt 1." +) @click.option( "--fault-tail-slow", default=0, show_default=True, help=f"Last N segments throttled to {TAIL_SLOW_KIB} KiB/s." ) @@ -391,7 +422,9 @@ def print_report(rows: list[Row], runs: int, total_segments: int) -> None: @click.option("--impersonate", help="rnet browser preset (e.g. Chrome131); uses RnetSession in either mode.") @click.option("--adaptive", is_flag=True, help="Use the adaptive worker controller (max_workers becomes the cap).") @click.option("--procs", default=1, show_default=True, help="Fan segments across N processes (multiprocess path).") -@click.option("--no-read1", is_flag=True, help="A/B: force racers back to blocking full-chunk reads (pre-fix behavior).") +@click.option( + "--no-read1", is_flag=True, help="A/B: force racers back to blocking full-chunk reads (pre-fix behavior)." +) def main( segments: int, seg_size: int, @@ -424,6 +457,8 @@ def main( saved = {k: getattr(dl_mod, k) for k in fast} for k, v in fast.items(): setattr(dl_mod, k, v) + # --procs children are spawned fresh and would otherwise keep the real timeouts + os.environ["UNSHACKLE_DL_TIMING_OVERRIDES"] = ",".join(f"{k}={v}" for k, v in fast.items()) if no_read1: if hasattr(dl_mod, "RACER_READ1"): @@ -478,6 +513,7 @@ def main( server.shutdown() for k, v in saved.items(): setattr(dl_mod, k, v) + os.environ.pop("UNSHACKLE_DL_TIMING_OVERRIDES", None) if __name__ == "__main__": diff --git a/tests/test_downloader_leaks.py b/tests/test_downloader_leaks.py new file mode 100644 index 00000000..d397e5a6 --- /dev/null +++ b/tests/test_downloader_leaks.py @@ -0,0 +1,135 @@ +"""Regression tests: no locked file handles or leaked worker threads after a batch ends. + +Covers the WinError 32 family: a racer leaked past a no-wait shutdown must never hold +(or reopen) its .!dev handle once the batch is over, even after the process-global +cancel flag is cleared for the next track. On Windows the unlink assertions fail +outright if a handle is still open; on POSIX the thread accounting catches the +same leak. +""" + +from __future__ import annotations + +import http.server +import importlib +import socketserver +import threading +import time +from pathlib import Path +from typing import Any, Iterator + +import pytest + +from unshackle.core.constants import DOWNLOAD_CANCELLED +from unshackle.core.downloaders.requests import requests as download_batch + +dl_mod = importlib.import_module("unshackle.core.downloaders.requests") + +SEG_SIZE = 256 * 1024 +FAST = {"READ_TIMEOUT": 1, "RETRY_WAIT": 0.2, "MAX_ATTEMPTS": 2, "HEDGE_MIN_WAIT": 0.3} + + +class _Server(socketserver.ThreadingTCPServer): + allow_reuse_address = True + daemon_threads = True + + def __init__(self) -> None: + self.body = (b"unshackle-test-" * (SEG_SIZE // 15 + 1))[:SEG_SIZE] + self.stall_release = threading.Event() + super().__init__(("127.0.0.1", 0), _Handler) + + def handle_error(self, request: Any, client_address: Any) -> None: # expected fault noise + pass + + +class _Handler(http.server.BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, *_: Any) -> None: + pass + + def do_GET(self) -> None: + server: _Server = self.server # type: ignore[assignment] + body = server.body + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + if self.path.startswith("/stall/"): # half the body, then hold: never completes + self.wfile.write(body[: len(body) // 2]) + self.wfile.flush() + server.stall_release.wait(30) + elif self.path.startswith("/slow/"): # trickle so a cancel can land mid-body + for off in range(0, len(body), 16 * 1024): + self.wfile.write(body[off : off + 16 * 1024]) + self.wfile.flush() + time.sleep(0.05) + else: + self.wfile.write(body) + + +@pytest.fixture() +def server() -> Iterator[_Server]: + srv = _Server() + threading.Thread(target=srv.serve_forever, daemon=True).start() + yield srv + srv.stall_release.set() + srv.shutdown() + + +@pytest.fixture(autouse=True) +def fast_timeouts() -> Iterator[None]: + saved = {k: getattr(dl_mod, k) for k in FAST} + for k, v in FAST.items(): + setattr(dl_mod, k, v) + DOWNLOAD_CANCELLED.clear() + yield + for k, v in saved.items(): + setattr(dl_mod, k, v) + DOWNLOAD_CANCELLED.clear() + + +def _pool_threads() -> int: + return sum(1 for t in threading.enumerate() if t.name.startswith("ThreadPoolExecutor")) + + +def _assert_drained_and_unlocked(out_dir: Path, baseline: int) -> None: + """Workers must exit within one read timeout; then every stray must be deletable.""" + deadline = time.monotonic() + FAST["READ_TIMEOUT"] + FAST["RETRY_WAIT"] + 3 + while time.monotonic() < deadline and _pool_threads() > baseline: + time.sleep(0.05) + assert _pool_threads() <= baseline, "worker threads leaked past the drain deadline" + for stray in out_dir.glob("*.!dev"): + stray.unlink() # PermissionError here on Windows = a handle is still held + # a resurrected worker would reopen/recreate its .!dev within one retry cycle + time.sleep(FAST["RETRY_WAIT"] + 0.3) + assert not list(out_dir.glob("*.!dev")), "a leaked worker recreated its temp file after cleanup" + + +def test_failed_batch_frees_handles_even_after_flag_clear(server: _Server, tmp_path: Path) -> None: + """Stall fault fails the batch; clearing DOWNLOAD_CANCELLED (as dl.py does for the + next track) must not resurrect the leaked racer's retry loop.""" + baseline = _pool_threads() + host, port = str(server.server_address[0]), int(server.server_address[1]) + urls = [f"http://{host}:{port}/stall/0"] + [f"http://{host}:{port}/ok/{i}" for i in range(1, 4)] + + with pytest.raises(Exception): + for _ in download_batch(urls=urls, output_dir=tmp_path, filename="seg_{i}.bin", max_workers=4): + pass + + assert DOWNLOAD_CANCELLED.is_set() + DOWNLOAD_CANCELLED.clear() # what dl.py does after a failed track + _assert_drained_and_unlocked(tmp_path, baseline) + + +def test_external_cancel_frees_handles(server: _Server, tmp_path: Path) -> None: + """DOWNLOAD_CANCELLED set mid-stream (sibling-track failure / SIGINT path): the batch + must wind down without leaving locked handles or live workers.""" + baseline = _pool_threads() + host, port = str(server.server_address[0]), int(server.server_address[1]) + urls = [f"http://{host}:{port}/slow/{i}" for i in range(4)] + + for event in download_batch(urls=urls, output_dir=tmp_path, filename="seg_{i}.bin", max_workers=4): + if "advance" in event or "written" in event: + DOWNLOAD_CANCELLED.set() + + DOWNLOAD_CANCELLED.clear() + _assert_drained_and_unlocked(tmp_path, baseline)