From 22cc739e8382c94e449bd9eba0942872f1668fd0 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Sun, 6 Sep 2026 20:03:08 -0600 Subject: [PATCH] fix(downloader): correct truncation, ranged sizing, and session carry Check a sequential download against the complete length whenever one is known, and settle the unknown-total case on the 416. An empty body of unknown length, a 206 with no complete length, and a 206 starting past the resume offset all passed the size checks before and were renamed into place. Four restart paths spent no attempt, so a host that always answered the same way cycled forever. --- tests/commands/test_cdm_cache.py | 145 ++++++++ tests/core/test_session_headers.py | 335 ++++++++++++++++++ tests/core/test_session_pool.py | 143 +++++++- tests/downloaders/test_probe_ranged.py | 214 +++++++++++ tests/downloaders/test_ranged_parallel.py | 114 +++++- .../test_short_range_completion.py | 273 ++++++++++++++ tests/downloaders/test_tail_boost.py | 61 +++- .../unit/test_cdm_patch_leaves_no_shadow.py | 40 +++ .../unit/test_licence_declared_kwargs.py | 25 +- tests/remote/unit/test_server_cdm_multikey.py | 2 +- .../remote/unit/test_session_service_lock.py | 252 +++++++++++++ unshackle/core/downloaders/requests.py | 247 +++++++++---- unshackle/core/session.py | 165 ++++++++- 13 files changed, 1895 insertions(+), 121 deletions(-) create mode 100644 tests/commands/test_cdm_cache.py create mode 100644 tests/core/test_session_headers.py create mode 100644 tests/downloaders/test_probe_ranged.py create mode 100644 tests/downloaders/test_short_range_completion.py create mode 100644 tests/remote/unit/test_cdm_patch_leaves_no_shadow.py create mode 100644 tests/remote/unit/test_session_service_lock.py diff --git a/tests/commands/test_cdm_cache.py b/tests/commands/test_cdm_cache.py new file mode 100644 index 00000000..fb957d98 --- /dev/null +++ b/tests/commands/test_cdm_cache.py @@ -0,0 +1,145 @@ +"""``dl.get_cdm`` builds each CDM once per run. + +Regression: the callers ask for a CDM once per title, and the dedupe guard at the +call site compares object identity. No CDM class defines ``__eq__``, so the guard +was always true and every title paid a fresh build. A remote device costs an RSA +key pair and a blocking request to its host. +""" + +from __future__ import annotations + +import logging +from typing import Any, Optional + +import pytest + +import unshackle.core.cdm as cdm_mod +from unshackle.commands.dl import dl +from unshackle.core.cdm import loader +from unshackle.core.config import config + + +class StubCdm: + """A distinct object per build, so a reused CDM is visible by identity.""" + + def __init__(self, cdm_name: str, service_name: str) -> None: + self.cdm_name = cdm_name + self.service_name = service_name + + +@pytest.fixture +def builds(monkeypatch) -> list[tuple[str, str]]: + """Record every ``load_cdm`` call and hand back a fresh stub for each one.""" + calls: list[tuple[str, str]] = [] + + def fake_load_cdm(cdm_name: str, *, service_name: str = "", vaults: Optional[Any] = None) -> StubCdm: + calls.append((cdm_name, service_name)) + return StubCdm(cdm_name, service_name) + + # The package resolves ``load_cdm`` lazily, so a ``setattr`` here would write the real + # function into the package dict on undo and shadow that lookup for the whole process. + monkeypatch.setitem(cdm_mod.__dict__, "load_cdm", fake_load_cdm) + return calls + + +def make_dl() -> dl: + # __new__ skips the CLI-driven __init__; get_cdm only needs these three. + instance = dl.__new__(dl) + instance.log = logging.getLogger("test_cdm_cache") + instance.cdm_override = None + instance.vaults = None + return instance + + +def test_repeated_calls_build_the_cdm_once(builds, monkeypatch) -> None: + monkeypatch.setattr(config, "cdm", {"EXAMPLE": "device_a"}) + runner = make_dl() + + first = runner.get_cdm("EXAMPLE", "default") + for _ in range(4): + assert runner.get_cdm("EXAMPLE", "default") is first + + assert builds == [("device_a", "EXAMPLE")] + + +def test_the_title_loop_guard_sees_the_same_object(builds, monkeypatch) -> None: + """The call site keeps its CDM when ``get_cdm`` returns what it already holds.""" + monkeypatch.setattr(config, "cdm", {"EXAMPLE": "device_a"}) + runner = make_dl() + + held = runner.get_cdm("EXAMPLE", "default") + for _ in range(3): + quality_based_cdm = runner.get_cdm("EXAMPLE", "default", drm="widevine", quality=2160) + assert not (quality_based_cdm and quality_based_cdm != held) + + assert len(builds) == 1 + + +def test_a_different_service_gets_its_own_cdm(builds, monkeypatch) -> None: + """A CDM is built with its ``service_name``, so two services must not share one.""" + monkeypatch.setattr(config, "cdm", {"EXAMPLE": "device_a", "OTHER": "device_a"}) + runner = make_dl() + + example = runner.get_cdm("EXAMPLE", "default") + other = runner.get_cdm("OTHER", "default") + + assert example is not other + assert builds == [("device_a", "EXAMPLE"), ("device_a", "OTHER")] + + +def test_a_different_profile_gets_its_own_cdm(builds, monkeypatch) -> None: + monkeypatch.setattr(config, "cdm", {"EXAMPLE": {"one": "device_one", "two": "device_two"}}) + runner = make_dl() + + one = runner.get_cdm("EXAMPLE", "one") + two = runner.get_cdm("EXAMPLE", "two") + + assert one is not two + assert runner.get_cdm("EXAMPLE", "one") is one + assert builds == [("device_one", "EXAMPLE"), ("device_two", "EXAMPLE")] + + +def test_quality_selection_reuses_the_device_it_resolves_to(builds, monkeypatch) -> None: + """Two spellings of the same device share one CDM; a different device does not.""" + monkeypatch.setattr(config, "cdm", {"EXAMPLE": {">=1080": "device_hd", "<1080": "device_sd"}}) + runner = make_dl() + + hd = runner.get_cdm("EXAMPLE", "default", quality=2160) + assert runner.get_cdm("EXAMPLE", "default", quality=1080) is hd + sd = runner.get_cdm("EXAMPLE", "default", quality=720) + + assert sd is not hd + assert builds == [("device_hd", "EXAMPLE"), ("device_sd", "EXAMPLE")] + + +def test_an_override_is_cached_under_the_name_it_resolves_to(builds, monkeypatch) -> None: + monkeypatch.setattr(config, "cdm", {"EXAMPLE": "device_a"}) + runner = make_dl() + runner.cdm_override = "device_b" + + override = runner.get_cdm("EXAMPLE", "default") + + assert runner.get_cdm("EXAMPLE", "default") is override + assert builds == [("device_b", "EXAMPLE")] + + +@pytest.fixture +def leaked_shadow(monkeypatch) -> None: + """Reproduce the residue that a ``setattr`` on the CDM package leaves behind. + + ``load_cdm`` reaches the package only through its lazy ``__getattr__``. A ``setattr`` on + the package writes the real function into the package dict when it undoes, which shadows + ``__getattr__`` for the rest of the process. Two remote tests do exactly that. + """ + monkeypatch.setitem(cdm_mod.__dict__, "load_cdm", loader.load_cdm) + + +def test_a_shadowed_package_attribute_still_reaches_the_stub(leaked_shadow, builds, monkeypatch) -> None: + """The stub intercepts ``get_cdm`` whatever an earlier test left on the CDM package.""" + monkeypatch.setattr(config, "cdm", {"EXAMPLE": "device_a"}) + runner = make_dl() + + cdm = runner.get_cdm("EXAMPLE", "default") + + assert isinstance(cdm, StubCdm) + assert builds == [("device_a", "EXAMPLE")] diff --git a/tests/core/test_session_headers.py b/tests/core/test_session_headers.py new file mode 100644 index 00000000..7d4ce7aa --- /dev/null +++ b/tests/core/test_session_headers.py @@ -0,0 +1,335 @@ +"""Removing a header or a cookie from an RnetSession must stop it reaching the wire. + +``client.update(headers=...)`` merges into the rnet client default headers and has no +removal, and the client keeps its own cookie store. These tests read what a real local +server received, because the session dicts pass even with the leak in place. +""" + +from __future__ import annotations + +import threading +from http.cookiejar import Cookie, CookieJar +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Iterator + +import pytest +import rnet + +from unshackle.core.session import RnetCookieAdapter, RnetSession + +pytestmark = pytest.mark.unit + + +class Recorder(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + received: list[dict[str, str]] = [] + cookies_sent: list[str] = [] + + def do_GET(self) -> None: + type(self).received.append({k.lower(): v for k, v in self.headers.items()}) + # rnet sends one Cookie header per cookie on HTTP/1.1, so a dict of the headers + # keeps only the last one and hides a cookie that leaked. + type(self).cookies_sent.append("; ".join(self.headers.get_all("Cookie") or [])) + body = b"ok" + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + if self.path == "/login": + self.send_header("Set-Cookie", "sid=abc123; Path=/") + elif self.path.startswith("/set/"): + self.send_header("Set-Cookie", f"{self.path[len('/set/') :]}; Path=/") + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args: object) -> None: + pass + + +@pytest.fixture +def server() -> Iterator[str]: + Recorder.received = [] + Recorder.cookies_sent = [] + httpd = ThreadingHTTPServer(("127.0.0.1", 0), Recorder) + threading.Thread(target=httpd.serve_forever, daemon=True).start() + try: + yield f"http://127.0.0.1:{httpd.server_port}" + finally: + httpd.shutdown() + httpd.server_close() + + +def make_cookie(name: str, value: str, domain: str) -> Cookie: + return Cookie( + version=0, + name=name, + value=value, + port=None, + port_specified=False, + domain=domain, + domain_specified=True, + domain_initial_dot=domain.startswith("."), + path="/", + path_specified=True, + secure=False, + expires=None, + discard=False, + comment=None, + comment_url=None, + rest={}, + ) + + +def test_popped_header_stops_being_sent(server: str) -> None: + session = RnetSession() + session.headers.update({"X-Auth": "secret", "X-Keep": "yes"}) + + session.get(f"{server}/one") + assert Recorder.received[-1]["x-auth"] == "secret" + + session.headers.pop("X-Auth", None) + session.get(f"{server}/two") + assert "x-auth" not in Recorder.received[-1] + assert Recorder.received[-1]["x-keep"] == "yes" + + +def test_deleted_header_stops_being_sent(server: str) -> None: + session = RnetSession() + session.headers["X-Auth"] = "secret" + + session.get(f"{server}/one") + assert "x-auth" in Recorder.received[-1] + + del session.headers["X-Auth"] + session.get(f"{server}/two") + assert "x-auth" not in Recorder.received[-1] + + +def test_clear_removes_every_header(server: str) -> None: + session = RnetSession() + session.headers.update({"X-Auth": "secret", "X-Trace": "on"}) + + session.get(f"{server}/one") + assert "x-auth" in Recorder.received[-1] + + session.headers.clear() + session.get(f"{server}/two") + assert "x-auth" not in Recorder.received[-1] + assert "x-trace" not in Recorder.received[-1] + + +def test_header_removal_keeps_cookies(server: str) -> None: + """A removed header must not cost the session the cookies it has already earned.""" + session = RnetSession() + session.headers["X-Auth"] = "secret" + + session.get(f"{server}/login") + session.get(f"{server}/before") + assert Recorder.received[-1].get("cookie") == "sid=abc123" + + session.headers.pop("X-Auth", None) + session.get(f"{server}/after") + assert Recorder.received[-1].get("cookie") == "sid=abc123" + + +def test_surviving_headers_keep_the_impersonate_order(server: str) -> None: + """A rebuilt client must overlay the preset headers, not replace them.""" + session = RnetSession(impersonate=rnet.Impersonate.Chrome131) + session.headers.update({"X-Auth": "secret", "X-Keep": "yes"}) + + session.get(f"{server}/one") + before = [k for k in Recorder.received[-1] if k != "x-auth"] + + session.headers.pop("X-Auth", None) + session.get(f"{server}/two") + assert [k for k in Recorder.received[-1]] == before + + +def test_header_removal_keeps_cookies_for_a_host_not_yet_visited(server: str) -> None: + """A cookie unshackle set itself must survive the rebuild. + + ``rebuild_client`` can only read the rnet cookie jar per origin, and it knows only the + origins the session has requested. Cookies loaded from a cookie file, or set through + ``session.cookies``, are for a host the session often reaches later. + """ + session = RnetSession() + session.headers["X-Auth"] = "secret" + session.get(f"{server}/one") + + session.cookies.set("sid", "abc123", domain="example.test") + session.headers.pop("X-Auth", None) + + raw = session._client.get_cookies("https://example.test") + assert raw is not None and b"sid=abc123" in raw + + +def test_rebuild_survives_an_origin_arriving_mid_rebuild(server: str, monkeypatch) -> None: + """``request`` records origins from the download threads while the rebuild reads them. + + Iterating the live set raises ``RuntimeError: Set changed size during iteration`` out of the + header removal. Adding an origin from inside the copy stands in for that race. + """ + session = RnetSession() + session.headers["X-Auth"] = "secret" + session.get(f"{server}/one") + + original = RnetSession.__dict__["copy_cookies"].__func__ + + def racing_copy(source, target, origin: str) -> None: + session._origins.add(f"https://late-{len(session._origins)}.test") + original(source, target, origin) + + monkeypatch.setattr(RnetSession, "copy_cookies", staticmethod(racing_copy)) + + session.headers.pop("X-Auth", None) + + session.get(f"{server}/two") + assert "x-auth" not in Recorder.received[-1] + + +def test_clear_drops_a_server_set_cookie(server: str) -> None: + """``cookies.clear()`` must reach the rnet cookie store, not only the local dicts.""" + session = RnetSession() + + session.get(f"{server}/set/sid=abc123") + session.get(f"{server}/before") + assert Recorder.cookies_sent[-1] == "sid=abc123" + + session.cookies.clear() + session.get(f"{server}/after") + assert Recorder.cookies_sent[-1] == "" + + +def test_clear_leaves_the_session_able_to_take_new_cookies(server: str) -> None: + """A cleared session must still hold the next cookie, and only that one.""" + session = RnetSession() + + session.get(f"{server}/set/sid=abc123") + session.cookies.clear() + + session.get(f"{server}/set/fresh=1") + session.get(f"{server}/after") + assert Recorder.cookies_sent[-1] == "fresh=1" + + +def test_delete_drops_only_the_named_server_set_cookie(server: str) -> None: + """``del cookies[name]`` must drop that name from the store and leave the rest.""" + session = RnetSession() + + session.get(f"{server}/set/sid=abc123") + session.get(f"{server}/set/other=keep") + + del session.cookies["sid"] + + session.get(f"{server}/after") + assert Recorder.cookies_sent[-1] == "other=keep" + + +def test_clear_by_name_drops_a_server_set_cookie(server: str) -> None: + """``clear(name=...)`` is the requests spelling a service reaches for.""" + session = RnetSession() + + session.get(f"{server}/set/sid=abc123") + session.get(f"{server}/before") + assert Recorder.cookies_sent[-1] == "sid=abc123" + + session.cookies.clear(name="sid") + session.get(f"{server}/after") + assert Recorder.cookies_sent[-1] == "" + + +def test_clear_before_the_client_exists_empties_the_jar_and_sends_nothing(server: str) -> None: + """``clear()`` on a session that has made no request must empty every cookie view. + + ``dl.save_cookies`` writes ``cookies.jar`` back to the cookie file and only ever merges, + so a cleared cookie left in the jar returns to disk. + """ + session = RnetSession() + jar = CookieJar() + jar.set_cookie(make_cookie("sid", "abc123", "127.0.0.1")) + session.cookies.update(jar) + session.cookies.set("other", "keep", domain="127.0.0.1") + + session.cookies.clear() + + assert dict(session.cookies) == {} + assert session.cookies.get_dict() == {} + assert list(session.cookies.jar) == [] + + session.get(f"{server}/after") + assert Recorder.cookies_sent[-1] == "" + + +def test_a_later_header_rebuild_does_not_resurrect_cleared_cookies(server: str) -> None: + """``rebuild_client`` replays the buffered cookies and copies the old jar per origin. + + A ``clear()`` that left either populated hands the cookie back on the next header removal. + """ + session = RnetSession() + session.headers["X-Auth"] = "secret" + + session.get(f"{server}/set/sid=abc123") + session.cookies.clear() + + session.headers.pop("X-Auth", None) + + session.get(f"{server}/after") + assert Recorder.cookies_sent[-1] == "" + + +def test_a_request_during_the_rebuild_still_carries_headers_and_cookies(server: str, monkeypatch) -> None: + """A request from another thread must never see a client with neither headers nor cookies. + + ``rebuild_client`` publishes the new client before it carries the headers and the cookies, + so a request in that window goes out unauthenticated. + """ + session = RnetSession() + session.headers.update({"X-Auth": "secret", "X-Trace": "on"}) + + session.get(f"{server}/set/sid=abc123") + session.get(f"{server}/before") + assert Recorder.cookies_sent[-1] == "sid=abc123" + + original = RnetCookieAdapter.flush_to_client + rebuilding: list[bool] = [] + + def flush_with_a_request(self) -> None: + if rebuilding: + session.get(f"{server}/mid") + original(self) + + monkeypatch.setattr(RnetCookieAdapter, "flush_to_client", flush_with_a_request) + rebuilding.append(True) + session.headers.pop("X-Trace", None) + + assert Recorder.received[-1].get("x-auth") == "secret" + assert Recorder.cookies_sent[-1] == "sid=abc123" + + +def test_the_rebuilt_client_is_published_only_when_it_is_ready(server: str, monkeypatch) -> None: + """The new client must go on ``_client`` last, after the headers, the flush and the copy.""" + session = RnetSession() + session.headers["X-Auth"] = "secret" + session.get(f"{server}/set/sid=abc123") + + old = session._client + seen: list[object] = [] + rebuilding: list[bool] = [] + + original_flush = RnetCookieAdapter.flush_to_client + original_copy = RnetSession.__dict__["copy_cookies"].__func__ + + def recording_flush(self) -> None: + if rebuilding: + seen.append(session._client) + original_flush(self) + + def recording_copy(source, target, origin: str) -> None: + seen.append(session._client) + original_copy(source, target, origin) + + monkeypatch.setattr(RnetCookieAdapter, "flush_to_client", recording_flush) + monkeypatch.setattr(RnetSession, "copy_cookies", staticmethod(recording_copy)) + rebuilding.append(True) + session.headers.pop("X-Auth", None) + + assert seen and all(client is old for client in seen) + assert session._client is not old diff --git a/tests/core/test_session_pool.py b/tests/core/test_session_pool.py index 06d3a139..2d977351 100644 --- a/tests/core/test_session_pool.py +++ b/tests/core/test_session_pool.py @@ -4,11 +4,17 @@ only the adapter mounts.""" from __future__ import annotations +from copy import copy +from typing import Any +from unittest.mock import patch + import pytest import requests from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry -from unshackle.core.service import TimeoutHTTPAdapter, grow_session_pool +from unshackle.core.service import DEFAULT_TIMEOUT, TimeoutHTTPAdapter, grow_session_pool +from unshackle.core.utils.sslciphers import SSLCiphers pytestmark = pytest.mark.unit @@ -20,43 +26,166 @@ class FakeRnetSession: raise AssertionError("grow_session_pool must not touch a non-requests session") -def test_small_pool_is_remounted_on_both_schemes() -> None: +def sent_timeout(adapter: HTTPAdapter) -> Any: + """The timeout an adapter hands to HTTPAdapter.send when the caller passes none.""" + seen: dict[str, Any] = {} + + def record(self: object, request: object, **kwargs: Any) -> str: + seen.update(kwargs) + return "sent" + + with patch.object(HTTPAdapter, "send", record): + assert adapter.send(object(), timeout=None) == "sent" + + return seen["timeout"] + + +def test_subclass_adapter_keeps_its_class_and_tls_context() -> None: + """A service-mounted SSLCiphers must survive the growth with its TLS context. + + CBS, friDay, iP, MY5 and PMTP all mount SSLCiphers on the service session and none + override get_session(). Replacing it with a plain adapter drops DEFAULT:@SECLEVEL=2 and + check_hostname=False, so metadata succeeds and every segment then fails at TLS or 403. + """ session = requests.Session() + session.mount("https://", SSLCiphers(security_level=2)) old = session.get_adapter("https://x") + context = old._ssl_context assert old._pool_maxsize < 64 grow_session_pool(session, 64) + grown = session.get_adapter("https://x") + assert grown is old + assert isinstance(grown, SSLCiphers) + assert grown._ssl_context is context + assert grown.poolmanager.connection_pool_kw["ssl_context"] is context + assert grown._pool_maxsize == 64 + assert grown._pool_block is True + + +def test_a_distinct_http_adapter_is_grown_not_overwritten() -> None: + """Every mounted adapter grows, and no mount overwrites another. + + The https:// adapter must not land on the http:// prefix, and a service that mounts on a + narrow prefix must be grown as well. + """ + session = requests.Session() + secure = SSLCiphers(security_level=2) + narrow = SSLCiphers(security_level=2) + plain = TimeoutHTTPAdapter(pool_connections=10, pool_maxsize=10) + session.mount("https://", secure) + session.mount("https://cbs.example/", narrow) + session.mount("http://", plain) + + grow_session_pool(session, 64) + + assert session.get_adapter("https://x") is secure + assert session.get_adapter("https://cbs.example/a") is narrow + assert session.get_adapter("http://x") is plain + assert session.get_adapter("http://x") is not session.get_adapter("https://x") + for adapter in (secure, narrow, plain): + assert adapter._pool_maxsize == 64 + assert adapter._pool_connections == 64 + + +def test_shared_adapter_is_grown_in_place_on_both_schemes() -> None: + session = requests.Session() + adapter = TimeoutHTTPAdapter(pool_connections=10, pool_maxsize=10, max_retries=Retry(total=5)) + session.mount("https://", adapter) + session.mount("http://", adapter) + retries = adapter.max_retries + timeout = adapter.default_timeout + + grow_session_pool(session, 64) + https = session.get_adapter("https://x") - http = session.get_adapter("http://x") - assert isinstance(https, TimeoutHTTPAdapter) - assert https is http - assert https is not old + assert https is adapter + assert session.get_adapter("http://x") is adapter + assert type(https) is TimeoutHTTPAdapter + assert https.max_retries is retries + assert https.default_timeout == timeout assert https._pool_maxsize == 64 + assert https._pool_connections == 64 assert https._pool_block is True - assert https.max_retries is old.max_retries + assert https.poolmanager.connection_pool_kw["maxsize"] == 64 + + +def test_cached_proxy_manager_is_rebuilt_at_the_new_size() -> None: + """proxy_manager_for caches one manager per proxy URL, sized when it was first built. + + authenticate() and get_tracks() run before the growth, so a service behind --proxy has + already filled that cache. Leaving it there keeps every proxied segment on the small pool. + """ + session = requests.Session() + adapter = TimeoutHTTPAdapter(pool_connections=4, pool_maxsize=4) + session.mount("https://", adapter) + proxy = "http://proxy.example:8080" + stale = adapter.proxy_manager_for(proxy) + assert stale.connection_pool_kw["maxsize"] == 4 + + grow_session_pool(session, 64) + + fresh = adapter.proxy_manager_for(proxy) + assert fresh.connection_pool_kw["maxsize"] == 64 + assert fresh is not stale def test_large_enough_pool_is_left_alone() -> None: session = requests.Session() session.mount("https://", HTTPAdapter(pool_connections=64, pool_maxsize=64)) old = session.get_adapter("https://x") + pool = old.poolmanager grow_session_pool(session, 64) assert session.get_adapter("https://x") is old + assert old.poolmanager is pool, "an adapter that needs no growth must keep its warm connections" def test_equal_size_pool_is_left_alone() -> None: session = requests.Session() session.mount("https://", HTTPAdapter(pool_maxsize=32)) old = session.get_adapter("https://x") + pool = old.poolmanager grow_session_pool(session, 32) assert session.get_adapter("https://x") is old + assert old.poolmanager is pool @pytest.mark.parametrize("session", [object(), FakeRnetSession()]) def test_non_requests_session_is_ignored(session: object) -> None: grow_session_pool(session, 64) + + +def test_copied_adapter_keeps_the_default_timeout() -> None: + """default_timeout must survive copy(), which no_retry_session runs per segment. + + copy() runs HTTPAdapter.__getstate__/__setstate__, which carry only the names in + __attrs__. Without default_timeout there, the copy quietly falls back to DEFAULT_TIMEOUT + on the segment path and drops the service timeout. + + The last two lines pin the getattr fallback in send(). A subclass that redefines + __attrs__ from HTTPAdapter.__attrs__, as SSLCiphers does, drops default_timeout again. + """ + twin = copy(TimeoutHTTPAdapter(timeout=(3, 300), pool_connections=4, pool_maxsize=4)) + assert twin.default_timeout == (3, 300) + assert sent_timeout(twin) == (3, 300) + + del twin.default_timeout + assert sent_timeout(twin) == DEFAULT_TIMEOUT + + +def test_growing_does_not_revert_a_custom_default_timeout() -> None: + session = requests.Session() + session.mount("https://", TimeoutHTTPAdapter(timeout=(1, 5), pool_connections=4, pool_maxsize=4)) + old = session.get_adapter("https://x") + + grow_session_pool(session, 64) + + grown = session.get_adapter("https://x") + assert grown.default_timeout == (1, 5) + assert grown is old + assert copy(grown).default_timeout == (1, 5), "the per-segment copy must keep the service timeout" diff --git a/tests/downloaders/test_probe_ranged.py b/tests/downloaders/test_probe_ranged.py new file mode 100644 index 00000000..e1ef8714 --- /dev/null +++ b/tests/downloaders/test_probe_ranged.py @@ -0,0 +1,214 @@ +"""``probe_ranged`` cost and fidelity. + +The probe asks for one byte to learn the complete length and whether the host honours +ranges. It runs once for every single-URL download and once per tail-boost candidate, +so what it costs and what it sends both matter: + +- it must leave the connection reusable. A streamed response whose body is never read + cannot go back in the pool. Every probe would then tear down a socket, and the next + request would pay a fresh handshake. +- it must ask for an unencoded body like ``download()`` does. A CDN that answers a + probe with a Content-Encoding makes the probe decline, and the ranged-parallel path + is silently lost for a file that supports it. +- it must carry the URL item's own request options. Without them a host that needs an + auth header answers the probe with a 4xx, and the download falls back to one + connection even though the ranges were there. + +The drain is bounded by the declared Content-Length, not by the status code. A host +that ignores Range may answer 200 with the whole resource, and one that answers 206 +with the whole resource is just as possible. Reading either would download the file +only to decide not to use it. +""" + +import importlib +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest +from requests import Session + +from unshackle.core.constants import DOWNLOAD_CANCELLED + +# the package re-exports the `requests` function, shadowing the module of the same name +dl = importlib.import_module("unshackle.core.downloaders.requests") + +PAYLOAD = bytes((i * 11 + 3) % 256 for i in range(32 * 1024)) +BIG = bytes(8 * 1024 * 1024) # larger than any socket buffer, so a full drain is visible +PROBES = 6 +AUTH = "let-me-in" + + +class _ProbeServer(ThreadingHTTPServer): + allow_reuse_address = True + daemon_threads = True + + def __init__(self): + self.lock = threading.Lock() + self.requests: list = [] # (Range, Accept-Encoding, X-Auth) + self.connections = 0 + self.mode = "range" + self.written = 0 + self.broke = False + super().__init__(("127.0.0.1", 0), _ProbeHandler) + + def handle_error(self, request, client_address): # the client hanging up is expected here + pass + + +class _ProbeHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, *_): + pass + + def setup(self): + super().setup() + with self.server.lock: + self.server.connections += 1 + + def do_GET(self): + srv = self.server + rng = self.headers.get("Range") + with srv.lock: + srv.requests.append((rng, self.headers.get("Accept-Encoding"), self.headers.get("X-Auth"))) + + if srv.mode == "needs_auth" and self.headers.get("X-Auth") != AUTH: + self.send_response(403) + self.send_header("Content-Length", "0") + self.end_headers() + return + + if srv.mode == "ignores_range": + self.send_response(200) + self.send_header("Content-Length", str(len(BIG))) + self.end_headers() + self._write_slowly(BIG) + return + + if srv.mode == "fat_206": + # a 206 status on a whole-resource body: the status alone cannot bound the drain + self.send_response(206) + self.send_header("Content-Range", f"bytes 0-{len(BIG) - 1}/{len(BIG)}") + self.send_header("Content-Length", str(len(BIG))) + self.end_headers() + self._write_slowly(BIG) + return + + if not rng: + self.send_response(200) + self.send_header("Content-Length", str(len(PAYLOAD))) + self.end_headers() + self.wfile.write(PAYLOAD) + return + + start_s, _, end_s = rng.removeprefix("bytes=").partition("-") + start = int(start_s) + end = int(end_s) if end_s else len(PAYLOAD) - 1 + 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))) + if srv.mode == "encode_unless_identity" and self.headers.get("Accept-Encoding") != "identity": + self.send_header("Content-Encoding", "gzip") + self.end_headers() + self.wfile.write(body) + + def _write_slowly(self, body): + # chunked writes so the count records how far the client actually read + srv = self.server + try: + for at in range(0, len(body), 64 * 1024): + self.wfile.write(body[at : at + 64 * 1024]) + with srv.lock: + srv.written += len(body[at : at + 64 * 1024]) + except OSError: + with srv.lock: + srv.broke = True + self.close_connection = True + + +@pytest.fixture +def server(monkeypatch): + DOWNLOAD_CANCELLED.clear() + monkeypatch.setattr(dl, "RETRY_WAIT", 0) + srv = _ProbeServer() + thread = threading.Thread(target=srv.serve_forever, daemon=True) + thread.start() + yield srv + srv.shutdown() + thread.join(timeout=5) + srv.server_close() + DOWNLOAD_CANCELLED.clear() + + +def url(srv, path="/payload.bin"): + return f"http://127.0.0.1:{srv.server_address[1]}{path}" + + +def test_probe_reuses_the_pooled_connection(server): + session = Session() + + for _ in range(PROBES): + assert dl.probe_ranged(url(server), session) == (len(PAYLOAD), True) + + assert server.connections == 1 + + +def test_probe_does_not_drain_a_200_body(server): + server.mode = "ignores_range" + session = Session() + + assert dl.probe_ranged(url(server), session) == (0, False) + + for _ in range(50): + with server.lock: + if server.broke or server.written >= len(BIG): + break + time.sleep(0.1) + with server.lock: + assert server.written < len(BIG) + + +def test_probe_asks_for_identity_encoding(server): + server.mode = "encode_unless_identity" + session = Session() + + assert dl.probe_ranged(url(server), session) == (len(PAYLOAD), True) + assert server.requests[-1][1] == "identity" + + # a caller that names its own encoding still wins + assert dl.probe_ranged(url(server), session, headers={"Accept-Encoding": "gzip"}) == (0, False) + assert server.requests[-1][1] == "gzip" + + +def test_probe_uses_the_url_items_headers(server, tmp_path, monkeypatch): + server.mode = "needs_auth" + monkeypatch.setattr(dl, "RANGE_PARALLEL_MIN_SIZE", 8 * 1024) + item = {"url": url(server), "headers": {"X-Auth": AUTH}} + + for _ in dl.requests([item], output_dir=tmp_path, filename="seg_{i:04}.bin", max_workers=4): + pass + + save_path = tmp_path / "seg_0000.bin" + assert save_path.read_bytes() == PAYLOAD + probes = [r for r in server.requests if r[0] == "bytes=0-0"] + assert probes and all(r[2] == AUTH for r in probes) + # the probe was answered, so the ranged path ran instead of one sequential connection + assert [r for r in server.requests if r[0] and r[0] != "bytes=0-0"] + + +def test_probe_does_not_drain_a_fat_206_body(server): + """A 206 carrying the whole resource must not be read to decide whether to use ranges.""" + server.mode = "fat_206" + session = Session() + + assert dl.probe_ranged(url(server), session) == (len(BIG), True) + + for _ in range(50): + with server.lock: + if server.broke or server.written >= len(BIG): + break + time.sleep(0.1) + with server.lock: + assert server.written < len(BIG) diff --git a/tests/downloaders/test_ranged_parallel.py b/tests/downloaders/test_ranged_parallel.py index 1f7f0ad5..b31d928b 100644 --- a/tests/downloaders/test_ranged_parallel.py +++ b/tests/downloaders/test_ranged_parallel.py @@ -13,7 +13,9 @@ Pins the byte-range fan-out that splits one large file across parts: or a content-encoded 206; - the public single-URL path falls back to a sequential download only when the server does not really support ranges, and re-raises a 4xx instead of paying - MAX_ATTEMPTS a second time from byte 0. + MAX_ATTEMPTS a second time from byte 0; +- the split plan is worker-driven: every worker gets a part, down to a 4 MiB floor + that stops a large worker count cutting a file into slivers. Tests 2/3 exercise ``dispatch_parts`` directly: the public path swallows a range-support failure and falls back, so the raise that the invariant guarantees @@ -118,7 +120,7 @@ def run(srv, tmp_path, urls, **kwargs): 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) + monkeypatch.setattr(dl, "RANGE_PARALLEL_MIN_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 @@ -182,7 +184,7 @@ def test_permanent_part_failure_is_not_retried_sequentially(server, tmp_path, mo server.mode = "deny_parts" monkeypatch.setattr(dl, "RETRY_WAIT", 0.01) monkeypatch.setattr(dl, "RANGE_PARALLEL_MIN_SIZE", 8 * 1024) - monkeypatch.setattr(dl, "RANGE_PARALLEL_PART_SIZE", 8 * 1024) + monkeypatch.setattr(dl, "RANGE_PARALLEL_MIN_PART_SIZE", 8 * 1024) with pytest.raises(HTTPError): run(server, tmp_path, [{"url": url(server)}], max_workers=4) # a sequential fallback would issue a plain unranged GET; only the probe and the @@ -198,7 +200,7 @@ def test_range_unsupported_falls_back_to_sequential(server, tmp_path, monkeypatc server.mode = "probe_only_206" monkeypatch.setattr(dl, "RETRY_WAIT", 0.01) monkeypatch.setattr(dl, "RANGE_PARALLEL_MIN_SIZE", 8 * 1024) - monkeypatch.setattr(dl, "RANGE_PARALLEL_PART_SIZE", 8 * 1024) + monkeypatch.setattr(dl, "RANGE_PARALLEL_MIN_PART_SIZE", 8 * 1024) files = run(server, tmp_path, [{"url": url(server)}], max_workers=4) assert len(files) == 1 assert files[0].read_bytes() == PAYLOAD @@ -214,3 +216,107 @@ def test_permanent_status_unwraps_causes(): wrapped = RuntimeError("max retries") wrapped.__cause__ = HTTPError(response=_Resp(404)) assert dl.permanent_status(wrapped) == 404 + + +BIG = bytes((i * 13 + 7) % 256 for i in range(8 * 1024 * 1024)) + + +class _BigHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self): + srv = self.server + rng = self.headers.get("Range") + with srv.lock: + srv.requests.append(rng) + if not rng: + self.send_response(200) + self.send_header("Content-Length", str(len(BIG))) + self.end_headers() + self.wfile.write(BIG) + return + start_s, _, end_s = rng.removeprefix("bytes=").partition("-") + start = int(start_s) + end = int(end_s) if end_s else len(BIG) - 1 + body = BIG[start : end + 1] + self.send_response(206) + self.send_header("Content-Range", f"bytes {start}-{end}/{len(BIG)}") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): + pass + + +@pytest.fixture +def big_server(): + srv = ThreadingHTTPServer(("127.0.0.1", 0), _BigHandler) + srv.lock = threading.Lock() + srv.requests = [] + thread = threading.Thread(target=srv.serve_forever, daemon=True) + thread.start() + yield srv + srv.shutdown() + thread.join(timeout=5) + + +_REAL_SPLIT = dl.split_ranges + + +class _SplitStop(Exception): + """Ends dispatch_parts once the split plan is known, before any request is made.""" + + +def plan_split(monkeypatch, tmp_path, total_size, max_workers): + """The (size, max_parts, part_target, parts) that dispatch_parts asks split_ranges for.""" + record: list = [] + + def spy(size, max_parts, part_target): + # _REAL_SPLIT, not dl.split_ranges: a second call in one test would otherwise + # chain through the previous spy and raise before recording + record.append((size, max_parts, part_target, _REAL_SPLIT(size, max_parts, part_target))) + raise _SplitStop + + monkeypatch.setattr(dl, "split_ranges", spy) + gen = dl.dispatch_parts( + url="http://127.0.0.1:1/unreachable", + save_path=tmp_path / "planned.bin", + session=Session(), + total_size=total_size, + max_workers=max_workers, + ) + with pytest.raises(_SplitStop): + next(gen) + return record[0] + + +def test_part_count_matches_worker_count(tmp_path, monkeypatch): + _, _, _, parts = plan_split(monkeypatch, tmp_path, 64 * 1024 * 1024, 16) + assert len(parts) == 16 + + _, _, _, parts = plan_split(monkeypatch, tmp_path, 16 * 1024 * 1024, 4) + assert len(parts) == 4 + + +def test_part_floor_caps_sliver_count(tmp_path, monkeypatch): + # below the 4 MiB floor a part's handshake costs more than its transfer + _, _, target, parts = plan_split(monkeypatch, tmp_path, 16 * 1024 * 1024, 64) + assert target == 4 * 1024 * 1024 + assert len(parts) == 4 + + +def test_worker_driven_parts_merge_byte_identical(big_server, tmp_path): + save_path = tmp_path / "big.bin" + for _ in dl.dispatch_parts( + url=url(big_server, "/big.bin"), + save_path=save_path, + session=Session(), + total_size=len(BIG), + max_workers=2, + ): + pass + + assert save_path.read_bytes() == BIG + assert len({r for r in big_server.requests if r}) == 2 + assert not save_path.with_name(save_path.name + ".!dev").exists() diff --git a/tests/downloaders/test_short_range_completion.py b/tests/downloaders/test_short_range_completion.py new file mode 100644 index 00000000..14628f95 --- /dev/null +++ b/tests/downloaders/test_short_range_completion.py @@ -0,0 +1,273 @@ +"""Completeness of a sequential download that the server serves in short pieces. + +``download()`` renames its ``.!dev`` partial into place and reports ``file_downloaded`` +only when the resource is whole. A host that caps every range answers an open-ended +resume (``Range: bytes=N-``) with an honest short 206: the body matches the response's +own ``Content-Length``, so the short-read check is satisfied while the file is still +truncated. These tests pin the completeness decision on the three paths that reach it: + +- a resume with a known complete length must continue from the served end until the + complete length is on disk; +- a resume with an unknown complete length (``Content-Range: bytes N-M/*``) has no + number to compare against, so the only proof of the end is the server refusing the + next range with a 416; +- a first request that carries no Range at all, which a host may still answer with a + short 206. Its Content-Range names the complete length, so the decision has a number + to compare against even though nothing was resumed. + +The guards beside them pin what the decision must not change: an unknown-length body +with no partial on disk still finalizes, a stale oversized partial is still discarded +on a 416, a resume that arrives whole costs no extra request, and an empty resume body +still raises instead of repeating the same request forever. +""" + +import importlib +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest +from requests import Session + +from unshackle.core.constants import DOWNLOAD_CANCELLED + +# the package re-exports the `requests` function, shadowing the module of the same name +dl = importlib.import_module("unshackle.core.downloaders.requests") + +TOTAL = 256 * 1024 +PAYLOAD = bytes((i * 37 + 5) % 256 for i in range(TOTAL)) + +SEED = 64 * 1024 # bytes a prior run left in the .!dev partial +CAP = 32 * 1024 # most the capped server serves for one range +OPEN_CAP = 16 * 1024 # most the halving server serves for an open-ended range +CUT = 96 * 1024 # where the halving server cuts the unranged body + + +class _ShortServer(ThreadingHTTPServer): + allow_reuse_address = True + daemon_threads = True + + def __init__(self): + self.lock = threading.Lock() + self.requests: list = [] + self.mode = "cap" + super().__init__(("127.0.0.1", 0), _ShortHandler) + + def handle_error(self, request, client_address): # the client hanging up is expected here + pass + + +class _ShortHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, *_): + pass + + def do_GET(self): + srv = self.server + header = self.headers.get("Range") + with srv.lock: + srv.requests.append(header) + mode = srv.mode + + if header is None: + if mode == "fresh206": + # a 206 to an un-Ranged request: the complete length is only in Content-Range + body = PAYLOAD[:CAP] + self.send_response(206) + self.send_header("Content-Range", f"bytes 0-{CAP - 1}/{TOTAL}") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + return + if mode == "chunked200": + # no Content-Length: an unknown-length body the downloader cannot size + self.send_response(200) + self.send_header("Transfer-Encoding", "chunked") + self.end_headers() + self.wfile.write(b"%x\r\n" % len(PAYLOAD) + PAYLOAD + b"\r\n0\r\n\r\n") + return + self.send_response(200) + self.send_header("Content-Length", str(TOTAL)) + self.end_headers() + if mode == "halving": + # honest length, then cut: download() credits a resume and bounds the next request + self.wfile.write(PAYLOAD[:CUT]) + self.close_connection = True + return + self.wfile.write(PAYLOAD) + return + + start_s, _, end_s = header.removeprefix("bytes=").partition("-") + start = int(start_s) + want_end = int(end_s) if end_s else None + + if start >= TOTAL: + self.send_response(416) + self.send_header("Content-Range", f"bytes */{TOTAL}") + self.send_header("Content-Length", "0") + self.end_headers() + return + + if mode == "empty206": + self.send_response(206) + self.send_header("Content-Range", f"bytes {start}-{TOTAL - 1}/{TOTAL}") + self.send_header("Transfer-Encoding", "chunked") + self.end_headers() + self.wfile.write(b"0\r\n\r\n") + return + + if want_end is not None: + end = min(TOTAL - 1, want_end) # a bounded sub-range is always served whole + elif mode == "whole_remainder": + end = TOTAL - 1 + elif mode == "halving": + end = min(TOTAL - 1, start + OPEN_CAP - 1) + else: + end = min(TOTAL - 1, start + CAP - 1) + + body = PAYLOAD[start : end + 1] + total = "*" if mode == "cap_star" else str(TOTAL) + self.send_response(206) + self.send_header("Content-Range", f"bytes {start}-{end}/{total}") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + +@pytest.fixture +def server(monkeypatch): + DOWNLOAD_CANCELLED.clear() + monkeypatch.setattr(dl, "RETRY_WAIT", 0) + srv = _ShortServer() + thread = threading.Thread(target=srv.serve_forever, daemon=True) + thread.start() + yield srv + srv.shutdown() + thread.join(timeout=5) + srv.server_close() + DOWNLOAD_CANCELLED.clear() + + +def url(srv): + return f"http://127.0.0.1:{srv.server_address[1]}/seg.mp4" + + +def run(srv, save_path): + return list(dl.download(url=url(srv), save_path=save_path, session=Session(), segmented=True)) + + +def assert_complete(save_path): + # size first: a truncation failure then names the byte counts instead of diffing blobs + assert save_path.stat().st_size == TOTAL + assert save_path.read_bytes() == PAYLOAD + + +def finalized(events): + return [event["written"] for event in events if "file_downloaded" in event] + + +def seed_partial(save_path, size): + tmp_file = save_path.with_name(f"{save_path.name}.!dev") + tmp_file.write_bytes(PAYLOAD[:size]) + return tmp_file + + +def test_capped_open_ended_resume_finishes_the_file(server, tmp_path): + server.mode = "cap" + save_path = tmp_path / "seg.mp4" + tmp_file = seed_partial(save_path, SEED) + + events = run(server, save_path) + + assert_complete(save_path) + assert finalized(events) == [TOTAL] + assert not tmp_file.exists() + + +def test_short_tail_after_chunk_halving_finishes_the_file(server, tmp_path, monkeypatch): + server.mode = "halving" + monkeypatch.setattr(dl, "REQUEST_CHUNK_SIZE", 64 * 1024) + monkeypatch.setattr(dl, "MIN_REQUEST_CHUNK", 8 * 1024) + monkeypatch.setattr(dl, "MIN_RESUME_PROGRESS", 1024) + save_path = tmp_path / "seg.mp4" + + events = run(server, save_path) + + assert_complete(save_path) + assert finalized(events) == [TOTAL] + # the bounded sub-ranges in the middle were each served whole and continued + bounded = [r for r in server.requests if r and r.split("-")[-1]] + assert len(bounded) >= 2 + + +def test_zero_byte_resume_body_raises_instead_of_spinning(server, tmp_path): + server.mode = "empty206" + save_path = tmp_path / "seg.mp4" + seed_partial(save_path, SEED) + + with pytest.raises(IOError): + run(server, save_path) + + assert not save_path.exists() + # an empty body must spend an attempt, not repeat the same request with no bound + assert len(server.requests) <= dl.MAX_ATTEMPTS + + +def test_unknown_total_resume_finalizes_on_416(server, tmp_path): + server.mode = "cap_star" + save_path = tmp_path / "seg.mp4" + tmp_file = seed_partial(save_path, SEED) + + events = run(server, save_path) + + assert_complete(save_path) + assert finalized(events) == [TOTAL] + assert not tmp_file.exists() + + +def test_stale_partial_416_still_restarts_clean(server, tmp_path): + server.mode = "cap" + save_path = tmp_path / "seg.mp4" + tmp_file = save_path.with_name(f"{save_path.name}.!dev") + tmp_file.write_bytes(b"\x00" * (TOTAL + 4096)) + + events = run(server, save_path) + + assert_complete(save_path) + assert finalized(events) == [TOTAL] + assert None in server.requests # restarted from byte 0 with an unranged GET + assert not tmp_file.exists() + + +def test_unknown_length_body_still_finalizes(server, tmp_path): + server.mode = "chunked200" + save_path = tmp_path / "seg.mp4" + + events = run(server, save_path) + + assert_complete(save_path) + assert finalized(events) == [TOTAL] + + +def test_complete_resume_finalizes_without_an_extra_request(server, tmp_path): + server.mode = "whole_remainder" + save_path = tmp_path / "seg.mp4" + seed_partial(save_path, SEED) + + events = run(server, save_path) + + assert_complete(save_path) + assert finalized(events) == [TOTAL] + # a known complete length that the first response satisfies costs no confirming request + assert len(server.requests) == 1 + + +def test_short_206_to_an_unranged_first_request_finishes_the_file(server, tmp_path): + """No resume, so the only complete length is the one the 206's Content-Range names.""" + server.mode = "fresh206" + save_path = tmp_path / "seg.mp4" + + events = run(server, save_path) + + assert_complete(save_path) + assert finalized(events) == [TOTAL] diff --git a/tests/downloaders/test_tail_boost.py b/tests/downloaders/test_tail_boost.py index f7de6789..83fae148 100644 --- a/tests/downloaders/test_tail_boost.py +++ b/tests/downloaders/test_tail_boost.py @@ -11,7 +11,9 @@ target, finalized by the last part worker via os.replace. The invariants: - 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). + os.replace'd into place (no corrupt finalized segment); +- a part that fails for a reason a retry could fix puts its segment back on the normal + single-worker path, without aborting the sibling segments or the track. 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 @@ -50,9 +52,16 @@ class _Handler(BaseHTTPRequestHandler): 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.deny_offset_ranges and start > 0: + # every non-leading window is refused for good (expired token, geo block) + self.send_response(403) + self.send_header("Content-Length", "0") + self.end_headers() + return 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 + # every non-leading window answers 200 with the whole body: part-mode's 206 + # check fails, but 200 is not permanent, so the segment goes back on the + # normal single-worker path self.send_response(200) self.send_header("Content-Length", str(len(BODY))) self.end_headers() @@ -80,6 +89,7 @@ def server(): srv.lock = threading.Lock() srv.requests = [] srv.break_offset_ranges = False + srv.deny_offset_ranges = False thread = threading.Thread(target=srv.serve_forever, daemon=True) thread.start() yield srv @@ -160,15 +170,52 @@ def test_tail_boost_finalize_last_part_wins(server, tmp_path): 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 +def test_tail_boost_part_failure_falls_back_to_normal_path(server, tmp_path, monkeypatch): + server.break_offset_ranges = True # every part window past the first fails, but not for good + monkeypatch.setattr(dl, "RETRY_WAIT", 0.01) + + files, advances = run(server, tmp_path, seg_urls(server)) + + assert boost_engaged(server), "tail boost never engaged; the fallback path is untested" + # a transient part failure returns the segment to the normal path instead of killing the track + assert len(files) == SEG_COUNT + assert all(f.read_bytes() == BODY for f in files) + # one advance per segment: a requeued segment must not be counted twice + assert advances == SEG_COUNT + assert not dl.DOWNLOAD_CANCELLED.is_set() + assert not list(tmp_path.glob("*.!dev")) + assert not list(tmp_path.glob("*.tp.!dev")) + + +def test_tail_boost_failure_does_not_abort_sibling_segments(server, tmp_path, monkeypatch): + server.break_offset_ranges = True + monkeypatch.setattr(dl, "RETRY_WAIT", 0.01) + + files, _ = run(server, tmp_path, seg_urls(server)) + + assert boost_engaged(server), "tail boost never engaged; the fallback path is untested" + # the parts of one segment share an abort of their own: setting the batch-wide one would + # strand every other in-flight worker and end the batch with segments missing + assert len(files) == SEG_COUNT + with server.lock: + reqs = list(server.requests) + boosted_paths = {path for path, rng in reqs if rng and rng != "bytes=0-0"} + assert boosted_paths + # each boosted segment went back on the normal path: a plain unranged GET for it + plain = {path for path, rng in reqs if rng is None} + assert boosted_paths <= plain + + +def test_tail_boost_permanent_part_failure_still_fails_the_batch(server, tmp_path, monkeypatch): + server.deny_offset_ranges = True # every part window past the first is a 403 monkeypatch.setattr(dl, "RETRY_WAIT", 0.01) with pytest.raises(Exception): run(server, tmp_path, seg_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. + assert boost_engaged(server), "tail boost never engaged; the permanent-failure rule is untested" + # a 4xx will not get better on the normal path either, so the boosted tail segments + # (6, 7) are never finalized and the batch fails instead of paying a second full attempt for i in (6, 7): assert not (tmp_path / f"seg_{i:04}.bin").exists() for f in tmp_path.glob("seg_*.bin"): diff --git a/tests/remote/unit/test_cdm_patch_leaves_no_shadow.py b/tests/remote/unit/test_cdm_patch_leaves_no_shadow.py new file mode 100644 index 00000000..841dcd73 --- /dev/null +++ b/tests/remote/unit/test_cdm_patch_leaves_no_shadow.py @@ -0,0 +1,40 @@ +"""No test may patch ``load_cdm`` on the CDM package with ``setattr``. + +``unshackle.core.cdm`` resolves ``load_cdm`` through a lazy ``__getattr__`` (PEP 562), so the +name never sits in the package dict. ``monkeypatch.setattr`` on the package reads the old value +through ``__getattr__`` and, on undo, writes the real function into the package dict. That key +shadows ``__getattr__`` for the rest of the process, and every later test that patches +``unshackle.core.cdm.loader`` loses its stub. ``tests/commands/test_cdm_cache.py`` ran the real +loader and raised ``ValueError: device_b does not exist``. + +``monkeypatch.setitem(cdm_mod.__dict__, "load_cdm", ...)`` leaves no residue: it deletes the key +it added. This guard names the file at fault instead of leaving the damage to surface in an +unrelated test. +""" + +from __future__ import annotations + +import re +from pathlib import Path +from typing import List + +import pytest + +pytestmark = pytest.mark.unit + +TESTS_ROOT = Path(__file__).resolve().parents[2] + +# A setattr of load_cdm on anything but the loader module, which holds the name for real. +PACKAGE_SETATTR = re.compile(r"""setattr\(\s*(?!loader\b)[\w.]+\s*,\s*["']load_cdm["']""") + + +def test_no_test_patches_load_cdm_on_the_package() -> None: + offenders: List[str] = [] + for path in sorted(TESTS_ROOT.rglob("test_*.py")): + if path == Path(__file__).resolve(): + continue + for number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1): + if PACKAGE_SETATTR.search(line): + offenders.append(f"{path.relative_to(TESTS_ROOT)}:{number}: {line.strip()}") + + assert not offenders, "patch the package dict with monkeypatch.setitem instead:\n" + "\n".join(offenders) diff --git a/tests/remote/unit/test_licence_declared_kwargs.py b/tests/remote/unit/test_licence_declared_kwargs.py index 5601514d..b70fc5c2 100644 --- a/tests/remote/unit/test_licence_declared_kwargs.py +++ b/tests/remote/unit/test_licence_declared_kwargs.py @@ -33,9 +33,6 @@ PSSH_B64 = base64.b64encode(b"pssh").decode() SESSION_ID = b"sid-bytes" -# --- fake services, one per licence signature shape ------------------------------------- - - class _Recorder: """Base for the fake services; every licence call lands in `calls`.""" @@ -86,9 +83,6 @@ class VarKwService(_Recorder): return "ok" -# --- stub DRM objects that keep the real get_content_keys call shape --------------------- - - class StubWidevine(Widevine): """Real `Widevine.get_content_keys`, with the PSSH parsing skipped.""" @@ -153,14 +147,11 @@ class FakePlayReadyCdm: return None -# --- fixtures --------------------------------------------------------------------------- - - @pytest.fixture def widevine_env(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(wv_pssh_mod, "PSSH", lambda b64: SimpleNamespace(b64=b64)) monkeypatch.setattr(drm_mod, "Widevine", StubWidevine) - monkeypatch.setattr(cdm_mod, "load_cdm", lambda *a, **k: FakeWidevineCdm()) + monkeypatch.setitem(cdm_mod.__dict__, "load_cdm", lambda *a, **k: FakeWidevineCdm()) monkeypatch.setattr(detect_mod, "is_widevine_cdm", lambda cdm: True) monkeypatch.setattr(handlers, "ensure_track_drm", lambda track, session=None, init_data=None: None) monkeypatch.setattr(handlers, "resolve_device_name", lambda *a, **k: "dev") @@ -173,7 +164,7 @@ def widevine_env(monkeypatch: pytest.MonkeyPatch) -> None: def playready_env(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(pr_pssh_mod, "PSSH", lambda data: SimpleNamespace(data=data)) monkeypatch.setattr(drm_mod, "PlayReady", StubPlayReady) - monkeypatch.setattr(cdm_mod, "load_cdm", lambda *a, **k: FakePlayReadyCdm()) + monkeypatch.setitem(cdm_mod.__dict__, "load_cdm", lambda *a, **k: FakePlayReadyCdm()) monkeypatch.setattr(detect_mod, "is_playready_cdm", lambda cdm: True) monkeypatch.setattr(handlers, "ensure_track_drm", lambda track, session=None, init_data=None: None) monkeypatch.setattr(handlers, "resolve_device_name", lambda *a, **k: "dev") @@ -192,9 +183,6 @@ def _server_cdm(service: Any, drm_type: str) -> dict[str, str]: ) -# --- server_cdm, Widevine --------------------------------------------------------------- - - def test_server_cdm_widevine_classic_signature_gets_three_arguments(widevine_env: None) -> None: service = ClassicService() keys = _server_cdm(service, "widevine") @@ -252,9 +240,6 @@ def test_server_cdm_widevine_lambda_is_never_itself_filtered(widevine_env: None) assert "session_id" not in service.calls[0] -# --- server_cdm, PlayReady -------------------------------------------------------------- - - def test_server_cdm_playready_classic_signature_gets_three_arguments(playready_env: None) -> None: service = ClassicService() keys = _server_cdm(service, "playready") @@ -281,9 +266,6 @@ def test_server_cdm_playready_session_service_is_not_passed_session_id(playready assert service.calls[0]["session_id"] is None # the signature default, never an offered value -# --- proxy license ---------------------------------------------------------------------- - - def _proxy(service: Any, drm_type: str) -> Any: return handlers.handle_proxy_license( service=service, @@ -333,9 +315,6 @@ def test_proxy_license_rejects_an_unknown_drm_type() -> None: _proxy(ClassicService(), "monalisa") -# --- the base-class PlayReady delegate -------------------------------------------------- - - class _DelegatingService(Service): """A service with no PlayReady implementation: the base class delegates to Widevine.""" diff --git a/tests/remote/unit/test_server_cdm_multikey.py b/tests/remote/unit/test_server_cdm_multikey.py index 6c072bad..3ad2de73 100644 --- a/tests/remote/unit/test_server_cdm_multikey.py +++ b/tests/remote/unit/test_server_cdm_multikey.py @@ -50,7 +50,7 @@ class _FakePlayReady: def playready_env(monkeypatch): monkeypatch.setattr(pr_pssh_mod, "PSSH", _FakePRPSSH) monkeypatch.setattr(drm_mod, "PlayReady", _FakePlayReady) - monkeypatch.setattr(cdm_mod, "load_cdm", lambda *a, **k: object()) + monkeypatch.setitem(cdm_mod.__dict__, "load_cdm", lambda *a, **k: object()) monkeypatch.setattr(detect_mod, "is_playready_cdm", lambda cdm: True) monkeypatch.setattr(handlers, "ensure_track_drm", lambda track, session=None, init_data=None: None) monkeypatch.setattr(handlers, "resolve_device_name", lambda *a, **k: "dev") diff --git a/tests/remote/unit/test_session_service_lock.py b/tests/remote/unit/test_session_service_lock.py new file mode 100644 index 00000000..04d4112a --- /dev/null +++ b/tests/remote/unit/test_session_service_lock.py @@ -0,0 +1,252 @@ +"""One remote session owns one service instance, so its handlers must take turns. + +``session_titles_handler`` and ``session_tracks_handler`` call the service from a worker +thread. Both calls reach the same long-lived ``SessionEntry.service_instance``, and services +keep per-title state on it: HTTP session headers, per-title tokens, cached manifest +attributes. Two requests on one remote session that run in that instance at the same time +corrupt each other's state, and the track response then hands the client the wrong headers. + +The probe here reports how many service calls are inside the instance at once. Its first +entrant waits a short window for company, so the test sees a genuinely concurrent run instead +of racing past it. A slow machine can only make a concurrent run look serialised, so these +tests cannot flake red. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import threading +from types import SimpleNamespace +from typing import Any, Iterator, List + +import pytest + +from unshackle.core.api import handlers +from unshackle.core.api.session_store import SessionEntry, get_session_store +from unshackle.core.titles.movie import Movie + +pytestmark = pytest.mark.unit + +# How long the first service call waits for a second one to join it. +OVERLAP_WINDOW = 0.5 + +# Bounds the wait when the loop is blocked, so a regression fails the test instead of hanging. +HANDSHAKE_TIMEOUT = 5.0 + + +class Overlap: + """Records how many service calls are inside the service instance at once.""" + + def __init__(self, window: float = OVERLAP_WINDOW) -> None: + self.window = window + self.max_active = 0 + self.trace: List[str] = [] + self._lock = threading.Lock() + self._active = 0 + self._joined = threading.Event() + self._waited = False + + @contextlib.contextmanager + def call(self, label: str) -> Iterator[None]: + with self._lock: + self._active += 1 + self.max_active = max(self.max_active, self._active) + self.trace.append(f"+{label}") + if self._active > 1: + self._joined.set() + first = not self._waited + self._waited = True + # Only the first call pays the window. Serialised handlers then cost one wait, not one + # per call, and a concurrent second call releases it as soon as it arrives. + if first: + self._joined.wait(timeout=self.window) + try: + yield + finally: + with self._lock: + self._active -= 1 + self.trace.append(f"-{label}") + + +def a_movie(title_id: str) -> Movie: + return Movie(id_=title_id, service=object, name=f"Film {title_id}", year=2024, language=None) + + +def empty_tracks() -> Any: + return SimpleNamespace(videos=[], audio=[], subtitles=[], attachments=[]) + + +class ServiceProbe: + """A service whose calls report their overlap into a shared recorder.""" + + def __init__(self, overlap: Overlap, name: str = "probe") -> None: + self.overlap = overlap + self.name = name + self.session = SimpleNamespace(headers={}, cookies=[]) + + def get_titles(self) -> List[Movie]: + with self.overlap.call(f"{self.name}:titles"): + return [a_movie("movie-0001"), a_movie("movie-0002")] + + def get_tracks(self, title: Any) -> Any: + with self.overlap.call(f"{self.name}:tracks:{title.id}"): + return empty_tracks() + + def get_chapters(self, title: Any) -> List[Any]: + with self.overlap.call(f"{self.name}:chapters:{title.id}"): + return [] + + +def seed_titles(session: SessionEntry) -> None: + """Put both titles in the remote session, as a prior titles call would.""" + for title_id in ("movie-0001", "movie-0002"): + session.title_map[title_id] = a_movie(title_id) + + +def bodies(responses: List[Any]) -> List[dict]: + parsed = [json.loads(response.body) for response in responses] + for body in parsed: + assert "error" not in body, body + return parsed + + +async def two_track_calls(session_id: str) -> List[Any]: + return list( + await asyncio.gather( + handlers.session_tracks_handler({"title_id": "movie-0001"}, session_id), + handlers.session_tracks_handler({"title_id": "movie-0002"}, session_id), + ) + ) + + +@pytest.mark.asyncio +async def test_two_track_requests_never_share_the_service() -> None: + overlap = Overlap() + store = get_session_store() + session = await store.create("EXAMPLE", ServiceProbe(overlap), session_id="lock-test-tracks") + seed_titles(session) + try: + responses = await two_track_calls(session.session_id) + finally: + await store.delete(session.session_id) + + bodies(responses) + assert overlap.max_active == 1, overlap.trace + + +@pytest.mark.asyncio +async def test_the_track_lock_covers_the_chapters_call() -> None: + overlap = Overlap() + store = get_session_store() + session = await store.create("EXAMPLE", ServiceProbe(overlap), session_id="lock-test-chapters") + seed_titles(session) + try: + responses = await two_track_calls(session.session_id) + finally: + await store.delete(session.session_id) + + bodies(responses) + # Every call closes before the next one opens, so one title's chapters are read before the + # other title's tracks are fetched. + assert len(overlap.trace) == 8, overlap.trace + for opened, closed in zip(overlap.trace[0::2], overlap.trace[1::2]): + assert closed == f"-{opened[1:]}", overlap.trace + + +@pytest.mark.asyncio +async def test_titles_and_tracks_never_overlap() -> None: + overlap = Overlap() + store = get_session_store() + session = await store.create("EXAMPLE", ServiceProbe(overlap), session_id="lock-test-mixed") + seed_titles(session) + try: + responses = list( + await asyncio.gather( + handlers.session_titles_handler(session.session_id), + handlers.session_tracks_handler({"title_id": "movie-0001"}, session.session_id), + ) + ) + finally: + await store.delete(session.session_id) + + bodies(responses) + assert overlap.max_active == 1, overlap.trace + + +@pytest.mark.asyncio +async def test_a_second_session_still_runs_in_parallel() -> None: + overlap = Overlap() + store = get_session_store() + first = await store.create("EXAMPLE", ServiceProbe(overlap, "a"), session_id="lock-test-parallel-a") + second = await store.create("EXAMPLE", ServiceProbe(overlap, "b"), session_id="lock-test-parallel-b") + try: + responses = list( + await asyncio.gather( + handlers.session_titles_handler(first.session_id), + handlers.session_titles_handler(second.session_id), + ) + ) + finally: + await store.delete(first.session_id) + await store.delete(second.session_id) + + bodies(responses) + assert first.lock is not second.lock + assert overlap.max_active == 2, overlap.trace + + +class BlockingTracks: + """Holds ``get_tracks`` open until the event loop releases it.""" + + def __init__(self) -> None: + self.entered = threading.Event() + self.resume = threading.Event() + self.loop_was_blocked = False + self.session = SimpleNamespace(headers={}, cookies=[]) + + def get_tracks(self, title: Any) -> Any: + self.entered.set() + self.loop_was_blocked = not self.resume.wait(timeout=HANDSHAKE_TIMEOUT) + return empty_tracks() + + def get_chapters(self, title: Any) -> List[Any]: + return [] + + +@pytest.mark.asyncio +async def test_the_loop_stays_free_while_the_lock_is_held() -> None: + service = BlockingTracks() + store = get_session_store() + session = await store.create("EXAMPLE", service, session_id="lock-test-loop") + seed_titles(session) + try: + task = asyncio.create_task(handlers.session_tracks_handler({"title_id": "movie-0001"}, session.session_id)) + try: + # Reached only while the loop is free, which is what the lock placement buys. + await asyncio.to_thread(service.entered.wait, HANDSHAKE_TIMEOUT) + held = session.lock.locked() + finally: + service.resume.set() + response = await asyncio.wait_for(task, HANDSHAKE_TIMEOUT) + finally: + await store.delete(session.session_id) + + bodies([response]) + assert not service.loop_was_blocked, "get_tracks blocked the event loop" + assert held, "the session lock was not held around the threaded service call" + + +@pytest.mark.asyncio +async def test_each_session_entry_gets_its_own_lock() -> None: + store = get_session_store() + first = await store.create("EXAMPLE", object(), session_id="lock-test-field-a") + second = await store.create("EXAMPLE", object(), session_id="lock-test-field-b") + try: + assert isinstance(first.lock, asyncio.Lock) + assert isinstance(second.lock, asyncio.Lock) + assert first.lock is not second.lock + finally: + await store.delete(first.session_id) + await store.delete(second.session_id) diff --git a/unshackle/core/downloaders/requests.py b/unshackle/core/downloaders/requests.py index 97f5be9f..7b36da7e 100644 --- a/unshackle/core/downloaders/requests.py +++ b/unshackle/core/downloaders/requests.py @@ -1,6 +1,7 @@ import math import multiprocessing import os +import pickle import random import re import socket @@ -48,6 +49,9 @@ PROGRESS_WINDOW = 2 CONNECT_TIMEOUT = 10 READ_TIMEOUT = 30 +# a probe asks for one byte, so a larger body is a host that ignored the range +PROBE_DRAIN_LIMIT = 4096 + # re-request a tail segment stuck past max(HEDGE_FACTOR * median segment time, HEDGE_MIN_WAIT) HEDGE_FACTOR = 3 HEDGE_MIN_WAIT = 5.0 @@ -58,7 +62,7 @@ DEFAULT_CHUNK = 524_288 SPEED_ROLLING_WINDOW = 10 # seconds of history to keep for speed calculation RANGE_PARALLEL_MIN_SIZE = 64 * 1024 * 1024 -RANGE_PARALLEL_PART_SIZE = 16 * 1024 * 1024 +RANGE_PARALLEL_MIN_PART_SIZE = 4 * 1024 * 1024 # One CPython interpreter caps sustained segment throughput (GIL in the ssl read path); # fan a big segment batch across spawned children to reach line rate. Below this count the @@ -193,7 +197,8 @@ class AdaptiveWorkerController: of throughput samples exists. The target starts at ``start`` (the cap by default) and follows AIMD from there: an error burst halves it and takes a one-tick cooldown, the controller reverts a probe upward that plateaus, and otherwise it climbs by - ``ADAPTIVE_STEP``. Target is + ``ADAPTIVE_STEP``. It judges a probe on the trailing tick before and after the change, so + both sides of the ``ADAPTIVE_PLATEAU_GAIN`` comparison measure the same span. Target is always clamped to ``[ADAPTIVE_MIN, cap]``. """ @@ -228,15 +233,18 @@ class AdaptiveWorkerController: def record_error(self, now: float) -> None: self._errors += 1 - def speed(self, now: float) -> float: - """Rolling bytes/sec over the window.""" + def speed(self, now: float, span: Optional[float] = None) -> float: + """Rolling bytes/sec over the window, or over the trailing ``span`` seconds when given.""" self.prune(now) if not self._samples: return 0.0 - span = now - self._samples[0][0] - if span <= 0: + if span is not None and span > 0: + cutoff = now - span + return sum(n for t, n in self._samples if t > cutoff) / span + elapsed = now - self._samples[0][0] + if elapsed <= 0: return 0.0 - return sum(n for _, n in self._samples) / span + return sum(n for _, n in self._samples) / elapsed def warmed_up(self, now: float) -> bool: # half a tick of samples gives enough of a baseline to probe against; waiting for @@ -267,6 +275,7 @@ class AdaptiveWorkerController: errors = self._errors self._errors = 0 speed_now = self.speed(now) + probe_speed = self.speed(now, min(self.tick, self.window)) old = self.target if self._cooldown: @@ -284,13 +293,13 @@ class AdaptiveWorkerController: # tail guard: too little work to saturate the target, so a low measured speed # here reflects starvation rather than a plateau, so hold and skip the probe/revert this tick pass - elif self._last_action == "increase" and speed_now < ADAPTIVE_PLATEAU_GAIN * self._speed_before_increase: + elif self._last_action == "increase" and probe_speed < ADAPTIVE_PLATEAU_GAIN * self._speed_before_increase: # last probe upward did not pay off: revert it and hold self.target = max(ADAPTIVE_MIN, self._target_before_increase) self._last_action = None reason = "plateau" elif self.target < self.cap: - self._speed_before_increase = speed_now + self._speed_before_increase = probe_speed self._target_before_increase = self.target self.target = min(self.cap, self.target + ADAPTIVE_STEP) self._last_action = "increase" @@ -328,6 +337,18 @@ def permanent_status(exc: BaseException) -> Optional[int]: return None +class TailPartFailed(Exception): + """A tail-boost range part failed. ``index`` names the segment it was split from. + + The batch loop needs the index to put that segment back on the normal single-worker + path; the failure itself stays reachable as ``__cause__`` for ``permanent_status``. + """ + + def __init__(self, index: int) -> None: + super().__init__(f"tail-boost part failed for segment {index}") + self.index = index + + def retry_sleep(exc: Exception, attempts: int) -> float: """Seconds to wait before retry number ``attempts``. @@ -376,14 +397,16 @@ def no_retry_session(session: Session) -> Session: download() owns segment retries, so a mounted urllib3 ``Retry`` nests inside a single attempt: one 429 carrying ``Retry-After: 60`` costs five one-minute sleeps before download() even sees the failure. requests takes no per-call retry override, and track - threads share the caller's session, so remounting it is not an option. A shallow adapter - copy keeps the poolmanager object, so TLS and pooling behaviour does not change. + threads share the caller's session, so remounting it is not an option. """ view = copy(session) view.adapters = OrderedDict() for prefix, adapter in session.adapters.items(): twin = copy(adapter) twin.max_retries = Retry(0, read=False) + if isinstance(adapter, HTTPAdapter) and isinstance(twin, HTTPAdapter): + twin.poolmanager = adapter.poolmanager + twin.proxy_manager = adapter.proxy_manager view.adapters[prefix] = twin return view @@ -409,6 +432,8 @@ def is_content_encoded(value: Optional[str]) -> bool: def probe_ranged(url: str, session: Any, **kwargs: Any) -> tuple[int, bool]: headers = {**(kwargs.get("headers") or {}), "Range": "bytes=0-0"} + if not any(str(k).lower() == "accept-encoding" for k in headers): + headers["Accept-Encoding"] = "identity" rest = {k: v for k, v in kwargs.items() if k != "headers"} if is_rnet_session(session): rest.setdefault("read_timeout", READ_TIMEOUT) @@ -430,6 +455,12 @@ def probe_ranged(url: str, session: Any, **kwargs: Any) -> tuple[int, bool]: total = content_range.rsplit("/", 1)[-1].strip() return (int(total), True) if total.isdigit() else (0, False) finally: + try: + declared = int(resp.headers.get("Content-Length") or resp.headers.get("content-length") or 0) + if resp.status_code == 206 and 0 < declared <= PROBE_DRAIN_LIMIT: + _ = resp.content + except Exception: + pass try: resp.close() except Exception: @@ -476,6 +507,17 @@ def parse_content_range(value: Any) -> Optional[tuple[int, int, Optional[int]]]: return int(match.group(1)), int(match.group(2)), None if match.group(3) == "*" else int(match.group(3)) +def parse_content_range_unsatisfied(value: Any) -> Optional[int]: + """Parse a 416's ``Content-Range`` into the complete length. Gives None when absent or unknown. + + A 416 answers with ``bytes */N`` (or ``bytes */*``), which the 206 form does not cover. + """ + if isinstance(value, bytes): + value = value.decode("latin1", "replace") + match = re.fullmatch(r"bytes\s+\*/(\d+)", str(value or "").strip(), re.IGNORECASE) + return int(match.group(1)) if match else None + + def range_clamped_to_end(content_range: tuple[int, int, Optional[int]], slice_end: int) -> bool: """Whether a 206 stops early only because the requested range ends past the resource. @@ -541,6 +583,18 @@ def plan_tail_parts(size: int, spare: int) -> list[tuple[int, int]]: return split_ranges(size, spare, TAIL_BOOST_PART_SIZE) +def submission_window(target: int) -> int: + """Futures to keep queued for ``target`` download workers. + + Submitting every segment up front convoys on the pool's global lock and starves the event + drain, so the drain loop meters submission through a window it tops up. The window is + twice the worker target: a worker that finishes a segment finds its next one already staged + instead of idling until the next drain cycle stages it. A pure function, so the ratio is + unit-testable and fixed mode and adaptive mode cannot drift apart. + """ + return target * 2 + + def tail_boost_engages(remaining: int, pending: int, target: int) -> bool: """True when idle workers outnumber the remaining whole segments. @@ -564,7 +618,8 @@ def dispatch_parts( save_path.parent.mkdir(parents=True, exist_ok=True) control_file = save_path.with_name(f"{save_path.name}.!dev") - parts = split_ranges(total_size, max_workers, RANGE_PARALLEL_PART_SIZE) + part_target = max(RANGE_PARALLEL_MIN_PART_SIZE, math.ceil(total_size / max_workers)) + parts = split_ranges(total_size, max_workers, part_target) control_file.write_bytes(b"") with open(save_path, "wb") as f: @@ -755,29 +810,24 @@ def download( _time = time.time use_raw = is_requests_session(session) - # item carries its own Range (DASH SegmentBase / HLS EXT-X-BYTERANGE slice): never - # overwrite it with a resume Range (that would fetch the parent's tail, not the slice); - # retries rewrite the whole slice in "wb" mode instead item_range = has_range_header(kwargs) attempts = 1 written = 0 resumes = 0 - # None until a cut proves the host bounds one request; the sequential path also needs a - # known complete length, or a short body could not be told apart from a finished download request_limit: Optional[int] = None known_total: Optional[int] = None + high_water = 0 progress_seeded = False 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 secured = written if part_mode else resume_offset + resume_start = resume_offset last_speed_refresh = _time() try: @@ -818,16 +868,26 @@ def download( stream = session.get(url, stream=True, **request_kwargs) if (not part_mode) and (not item_range) and resume_offset > 0 and stream.status_code == 416: - # our Range started past the end (a stale or oversized .!dev from a prior run). - # discard it and restart clean rather than raise_for_status → burn retries. + complete = parse_content_range_unsatisfied( + stream.headers.get("Content-Range") or stream.headers.get("content-range") + ) try: stream.close() except Exception: pass + if complete == resume_offset and known_total in (None, resume_offset): + if not segmented and not progress_seeded: + # nothing is read on this path, so the bar has to be seeded here + progress_seeded = True + yield dict(total=resume_offset) + yield dict(advance=resume_offset) + os.replace(tmp_file, save_path) + yield dict(file_downloaded=save_path, written=resume_offset) + if segmented: + yield dict(advance=1) + break tmp_file.unlink(missing_ok=True) - resume_offset = 0 - written = 0 - continue + raise IOError(f"resume from byte {resume_offset} got a 416, complete length {complete}") stream.raise_for_status() @@ -848,14 +908,8 @@ def download( resumed = False resume_offset = 0 else: - try: - stream.close() - except Exception: - pass tmp_file.unlink(missing_ok=True) - resume_offset = 0 - written = 0 - continue + raise IOError(f"resume from byte {resume_offset} got Content-Range {content_range!r}") elif content_range[2] is not None: known_total = content_range[2] if (not part_mode) and resume_offset > 0 and not resumed: @@ -895,6 +949,15 @@ def download( # a 200 carries the whole resource, so its length is the complete length known_total = content_length or None + if known_total is None and not part_mode and not item_range and stream.status_code == 206: + fresh_range = parse_content_range( + stream.headers.get("Content-Range") or stream.headers.get("content-range") + ) + if fresh_range is not None: + if fresh_range[0] != resume_offset: + raise IOError(f"206 body starts at byte {fresh_range[0]}, expected {resume_offset}") + known_total = fresh_range[2] + if item_range and not part_mode and stream.status_code != 206: # server ignored the slice's Range and sent the whole parent (RFC 9110 allows # a 200 here); writing that as the segment would silently corrupt the merge, so @@ -926,15 +989,12 @@ def download( f"byte-range segment got Content-Range {content_range!r}, " f"expected {slice_start}-{slice_end if slice_end is not None else ''}" ) + if not content_encoded: + content_length = content_range[1] - content_range[0] + 1 if resumed and content_encoded: - try: - stream.close() - except Exception: - pass tmp_file.unlink(missing_ok=True) - resume_offset = 0 - continue + raise IOError(f"resume from byte {resume_offset} came back content-encoded") limiter = speed_limiter chunk_size = adaptive_chunk_size(content_length) @@ -1055,9 +1115,17 @@ def download( raise IOError(f"Failed to read {content_length} bytes from the track URI.") if not part_mode: - if chunk_end is not None and known_total is not None and resume_offset + written < known_total: + if known_total is not None: + incomplete = resume_offset + written < known_total + else: + incomplete = written == 0 or (stream.status_code == 206 and not item_range) + if incomplete: + if resume_offset + written <= resume_start: + raise IOError(f"range from byte {resume_start} returned no bytes past it") resume_offset += written - attempts = 1 + if resume_offset > high_water: + high_water = resume_offset + attempts = 1 continue os.replace(tmp_file, save_path) yield dict(file_downloaded=save_path, written=resume_offset + written) @@ -1114,22 +1182,49 @@ def download( attempts += 1 +def rebuildable_adapter(adapter: Any) -> bool: + """Whether a child can remount ``adapter`` with the same transport behaviour. + + A child mounts a plain HTTPAdapter. TimeoutHTTPAdapter counts as the same transport here: + its one override supplies a default timeout when the caller passes none, and every request + this module sends carries an explicit timeout. Any other subclass changes the transport + itself, which a spec cannot carry: SSLCiphers swaps the SSL context, so a child would + negotiate a different cipher list and the host would reject it. + """ + from unshackle.core.service import TimeoutHTTPAdapter + + return type(adapter) in (HTTPAdapter, TimeoutHTTPAdapter) + + def build_session_spec(session: Optional[Any]) -> Optional[dict[str, Any]]: """Picklable spec to rebuild ``session`` in a child process. Gives None when it is not cheaply rebuildable. Live sessions (sockets, TLS state, threads) cannot cross a process boundary, so each child - reconstructs its own from cheap state. RnetSession needs a resolvable impersonate preset. - Without one it is not cheaply rebuildable, and the caller falls back to a single process. + reconstructs its own from cheap state. State a child cannot reproduce gives None, and the + caller falls back to a single process: a slower correct download beats a broken one. + RnetSession needs a resolvable impersonate preset. A requests session needs every mounted + adapter to be one a child can remount, and a spec spawn can pickle. """ if session is None: return {"kind": "none"} if is_requests_session(session): - return { + if not all(rebuildable_adapter(adapter) for adapter in session.adapters.values()): + return None + spec: dict[str, Any] = { "kind": "requests", "headers": dict(session.headers), "cookies": session.cookies, # RequestsCookieJar pickles cleanly "proxies": dict(session.proxies), + "params": copy(session.params), + "verify": session.verify, + "cert": session.cert, + "auth": session.auth, } + try: + pickle.dumps(spec) + except Exception: + return None + return spec if is_rnet_session(session): name = session.impersonate_name if not name: @@ -1139,6 +1234,7 @@ def build_session_spec(session: Optional[Any]) -> Optional[dict[str, Any]]: "impersonate": name, "headers": dict(session.headers), "cookies": session.cookies.get_dict_by_domain(), + "origin_cookies": session.export_origin_cookies(), "proxy": session.proxies.get("all") or session.proxies.get("https") or session.proxies.get("http"), } return None @@ -1164,6 +1260,11 @@ def rebuild_session(spec: dict[str, Any], max_workers: int) -> Optional[Any]: rs.cookies.update(spec["cookies"]) if spec.get("proxies"): rs.proxies.update(spec["proxies"]) + if spec.get("params"): + rs.params = spec["params"] + rs.verify = spec.get("verify", True) + rs.cert = spec.get("cert") + rs.auth = spec.get("auth") adapter = HTTPAdapter(pool_connections=max_workers, pool_maxsize=max_workers, pool_block=True) rs.mount("https://", adapter) rs.mount("http://", adapter) @@ -1183,6 +1284,7 @@ def rebuild_session(spec: dict[str, Any], max_workers: int) -> Optional[Any]: ns.cookies.update(jar) if spec.get("proxy"): ns.proxies.update({"all": spec["proxy"]}) + ns.import_origin_cookies(spec.get("origin_cookies") or {}) return ns return None # "none": child builds its own Session from the passed headers/cookies/proxy @@ -1209,8 +1311,9 @@ def mp_worker(queue: Any, kwargs: dict[str, Any]) -> None: kwargs["session"] = rebuild_session(spec, int(kwargs.get("max_workers") or 1)) for event in requests(**kwargs): queue.put(event) - except Exception: - queue.put({"__mp_error__": traceback.format_exc()}) + except Exception as exc: + summary = " ".join(f"{type(exc).__name__}: {exc}".split()) + queue.put({"__mp_error__": f"{summary}\n{traceback.format_exc()}"}) finally: queue.put({"__mp_done__": True}) @@ -1326,7 +1429,7 @@ def download_multiprocess( if "__mp_error__" in event: for p in procs: p.terminate() - raise RuntimeError(f"segment download child failed:\n{event['__mp_error__']}") + raise RuntimeError(f"segment download child failed: {event['__mp_error__']}") if "total" in event: continue if "downloaded" in event: @@ -1581,7 +1684,8 @@ def requests( and not has_range_header(url_item) and not url_item["save_path"].exists() ): - total_size, supports_ranges = probe_ranged(url_item["url"], session) + probe_kwargs = {k: v for k, v in url_item.items() if k not in ("url", "save_path")} + total_size, supports_ranges = probe_ranged(url_item["url"], session, **probe_kwargs) if supports_ranges and total_size >= RANGE_PARALLEL_MIN_SIZE: try: yield from dispatch_parts( @@ -1730,6 +1834,7 @@ def requests( # indices probed and left for the normal single-worker path, each decided once tail_boosted: set[int] = set() tail_skipped: set[int] = set() + part_aborts: list[threading.Event] = [] def submit(count: int, only: Optional[set[int]] = None) -> None: # only= submits just those indices (used in the tail window to release @@ -1762,6 +1867,7 @@ def requests( part_lock: threading.Lock, req_kwargs: dict[str, Any], ) -> None: + failure: Optional[Exception] = None try: # part_mode download writes [start, end] into the pre-truncated target and # emits only byte-`advance`; swallow those, since tail segments report at segment @@ -1777,16 +1883,24 @@ def requests( **req_kwargs, ): pass - except BaseException: - abort.set() # stop sibling parts of this segment; a real failure fails the batch - raise + except BaseException as exc: + abort.set() # stop the sibling parts of this segment, and only those + if not isinstance(exc, Exception): + raise # KeyboardInterrupt and friends still end the batch + failure = exc with part_lock: parts_left[0] -= 1 - finalize = parts_left[0] == 0 - if not finalize: + last = parts_left[0] == 0 + if DOWNLOAD_CANCELLED.is_set(): + return # cancelled: keep the partial target for the stray sweep, don't finalize + if failure is not None or abort.is_set(): + if last: + part_target.unlink(missing_ok=True) + if failure is not None: + raise TailPartFailed(index) from failure + return # a sibling part failed; that one reports the segment + if not last: return - if DOWNLOAD_CANCELLED.is_set() or abort.is_set(): - return # cancelled/aborted: keep the partial target for the stray sweep, don't finalize with seg_lock: if index in seg_done: part_target.unlink(missing_ok=True) @@ -1849,8 +1963,8 @@ def requests( f.truncate(size) seg_start[index] = time.time() tail_boosted.add(index) - # a failed part fails the whole batch anyway, so parts share batch_abort - abort = batch_abort + abort = threading.Event() + part_aborts.append(abort) parts_left = [len(parts)] part_lock = threading.Lock() for start, end in parts: @@ -1872,12 +1986,10 @@ def requests( ) boosted += 1 - # submitting every segment up front convoys on the pool's global lock and starves the - # event drain, so fixed mode meters submission through a window the loop tops up - queue_depth = max_workers * 2 + queue_depth = submission_window(max_workers) if controller: - submit(controller.update(time.time())) + submit(controller.update(time.time(), len(remaining))) else: submit(queue_depth) @@ -1914,15 +2026,12 @@ def requests( if controller: target = controller.update(now, len(pending) + len(remaining)) maybe_tail_boost(target) - if len(pending) < target: - # In the final stride (unstarted segments <= one worker target), hold back - # boost candidates so idle workers accumulate for a range-split; only - # release segments the boost already declined (too small / no range) so - # nothing stalls. Otherwise top up in FIFO order as usual. + queue_target = submission_window(target) + if len(pending) < queue_target: if remaining and len(remaining) <= target: submit(target - len(pending), only=tail_skipped) else: - submit(target - len(pending)) + submit(queue_target - len(pending)) elif len(pending) < queue_depth: submit(queue_depth - len(pending)) @@ -1975,6 +2084,12 @@ def requests( exc = future.exception() if isinstance(exc, KeyboardInterrupt): raise KeyboardInterrupt() + elif isinstance(exc, TailPartFailed) and permanent_status(exc) is None: + if exc.index not in tail_skipped: + tail_boosted.discard(exc.index) + tail_skipped.add(exc.index) + remaining.appendleft((exc.index, urls[exc.index])) + continue elif exc: DOWNLOAD_CANCELLED.set() yield dict(downloaded="[red]FAILING") @@ -2006,6 +2121,8 @@ def requests( # batch_abort outlives the shutdown so a leaked worker exits at its next # arrival/timeout instead of retrying after the global flag is cleared batch_abort.set() + for part_abort in part_aborts: + part_abort.set() # batch_abort is set, so a loser whose read we unblock here sees it and exits without # retrying; unblock only on the success path (the wait below still guarantees handles # are released before merge sweeps). rnet streams have no reachable socket -> no-op. diff --git a/unshackle/core/session.py b/unshackle/core/session.py index 1d8ae78b..f1c8fff3 100644 --- a/unshackle/core/session.py +++ b/unshackle/core/session.py @@ -230,13 +230,14 @@ class RnetResponse: class RnetSessionHeaders(CaseInsensitiveDict): """Dict-like headers that write to the rnet client through update().""" - def __init__(self, client: Any) -> None: - self._client = client + def __init__(self, session: Optional[RnetSession] = None) -> None: + self._session = session + self._client: Any = None super().__init__() def sync(self) -> None: """Push current headers to the rnet client.""" - if self._client is not None and hasattr(self, "_store") and self._store: + if self._client is not None and hasattr(self, "_store"): self._client.update(headers={k: v for k, v in self.items()}) def __setitem__(self, key: str, value: str) -> None: @@ -255,21 +256,31 @@ class RnetSessionHeaders(CaseInsensitiveDict): super().__setitem__(k, v) self.sync() - def pop(self, key: str, *args: Any) -> Any: - result = super().pop(key, *args) - # rnet doesn't support removing individual headers, but we track locally - # and always send the full set on next update - return result + def rebuild(self) -> None: + """Remake the rnet client so removed headers stop going out. + + rnet's update() merges into the client default headers and offers no removal, + so a header dropped from this dict alone keeps reaching every later host. + """ + if self._session is not None and self._client is not None: + self._session.rebuild_client() def __delitem__(self, key: str) -> None: super().__delitem__(key) + self.rebuild() + + def clear(self) -> None: + for key in list(self.keys()): + super().__delitem__(key) + self.rebuild() class RnetCookieAdapter(MutableMapping): """Cookie adapter that bridges requests-style cookie access to rnet.""" - def __init__(self, client: Any) -> None: + def __init__(self, client: Any, session: Optional[RnetSession] = None) -> None: self._client = client + self._session = session self._cookies: dict[str, dict[str, str]] = {} self._flat: dict[str, str] = {} self._original_cookies: list[Any] = [] @@ -294,6 +305,36 @@ class RnetCookieAdapter(MutableMapping): except Exception: pass + def client_urls(self, domain: Optional[str] = None) -> list[str]: + """URLs the rnet cookie jar could hold a cookie under. + + rnet keys its jar by URL and cannot list the jar whole, so a removal has to name every + URL this class has set a cookie on and every origin the HTTP session has requested. + """ + urls = {f"https://{d.lstrip('.')}" for d in self._cookies if d} + urls.add("https://localhost") + if self._session is not None: + urls |= set(self._session._origins) + if domain is not None: + host = domain.lstrip(".") + urls = {url for url in urls if (urlparse(url).hostname or "") == host} + urls.add(f"https://{host}") + return sorted(urls) + + def remove_cookie_on_client(self, name: str, domain: Optional[str] = None) -> None: + """Drop a cookie from the rnet cookie jar. + + The client is built with cookie_store=True, so a cookie the server set lives in that + jar alone. Dropping it from the local dicts leaves it going out on the wire. + """ + if self._client is None: + return + for url in self.client_urls(domain): + try: + self._client.remove_cookie(url, name) + except Exception: + pass + @property def jar(self) -> CookieJar: """Return a CookieJar with original Cookie objects (requests compat). @@ -355,6 +396,8 @@ class RnetCookieAdapter(MutableMapping): self._flat.pop(name, None) for domain_cookies in self._cookies.values(): domain_cookies.pop(name, None) + self._original_cookies = [cookie for cookie in self._original_cookies if cookie.name != name] + self.remove_cookie_on_client(name) def __contains__(self, name: object) -> bool: return name in self._flat @@ -406,6 +449,12 @@ class RnetCookieAdapter(MutableMapping): else: for domain_cookies in self._cookies.values(): domain_cookies.pop(name, None) + self._original_cookies = [ + cookie + for cookie in self._original_cookies + if cookie.name != name or (domain is not None and cookie.domain != domain) + ] + self.remove_cookie_on_client(name, domain) elif domain is not None: removed = self._cookies.pop(domain, {}) for k in removed: @@ -413,9 +462,15 @@ class RnetCookieAdapter(MutableMapping): still_exists = any(k in dc for dc in self._cookies.values()) if not still_exists: self._flat.pop(k, None) + self._original_cookies = [cookie for cookie in self._original_cookies if cookie.domain != domain] + for k in removed: + self.remove_cookie_on_client(k, domain) else: self._flat.clear() self._cookies.clear() + self._original_cookies.clear() + if self._client is not None: + self._client.clear_cookies() def items(self) -> list[tuple[str, str]]: return list(self._flat.items()) @@ -530,9 +585,10 @@ class RnetSession: self._client_kwargs = dict(client_kwargs) self._client: Optional[rnet.BlockingClient] = None + self._origins: set[str] = set() - self.headers = RnetSessionHeaders(None) - self.cookies = RnetCookieAdapter(None) + self.headers = RnetSessionHeaders(self) + self.cookies = RnetCookieAdapter(None, self) self.proxies = RnetProxyDict(self) if "headers" in session_kwargs: @@ -559,13 +615,90 @@ class RnetSession: def ensure_client(self) -> rnet.BlockingClient: """Lazily make the rnet client on first use, flushing any buffered state.""" if self._client is None: - self._client = rnet.BlockingClient(**self._client_kwargs) - self.headers._client = self._client + client = rnet.BlockingClient(**self._client_kwargs) + self.headers._client = client self.headers.sync() - self.cookies._client = self._client + self.cookies._client = client self.cookies.flush_to_client() + self._client = client return self._client + def rebuild_client(self) -> None: + """Replace the rnet client with one that carries only the current header set. + + rnet merges header updates and has no header removal, so a new client is the only + way to drop a header the caller deleted. Do not pass the headers as default_headers: + that replaces the impersonate header set instead of overlaying it, which breaks the + fingerprint. + + A new client also starts with an empty cookie jar. This method replays the cookies + unshackle set itself, then the cookies of every origin this HTTP session has used. + That order lets a value the server rotated win over the buffered one. The new client + goes on the HTTP session last, so a request from another thread keeps the old client + until the new one holds the headers and the cookies. + """ + old_client = self._client + if old_client is None: + return + client = rnet.BlockingClient(**self._client_kwargs) + self.headers._client = client + self.headers.sync() + self.cookies._client = client + self.cookies.flush_to_client() + for origin in list(self._origins): + self.copy_cookies(old_client, client, origin) + self._client = client + + @staticmethod + def set_cookie_header(client: rnet.BlockingClient, origin: str, header: str) -> None: + """Set every cookie in a ``name=value; ...`` header string on one origin of *client*.""" + for pair in header.split(";"): + name, separator, value = pair.partition("=") + if not separator: + continue + try: + client.set_cookie(origin, rnet.Cookie(name.strip(), value.strip())) + except Exception: + pass + + @staticmethod + def copy_cookies(source: rnet.BlockingClient, target: rnet.BlockingClient, origin: str) -> None: + """Copy one origin's cookies between clients. rnet cannot list the whole cookie jar.""" + try: + raw = source.get_cookies(origin) + except Exception: + return + if not raw: + return + RnetSession.set_cookie_header(target, origin, raw.decode("utf-8", errors="replace")) + + def export_origin_cookies(self) -> dict[str, str]: + """Cookie headers the live client holds, one per origin this HTTP session has requested. + + The client owns its own cookie jar, so a cookie a server set lives there and in no + part of :class:`RnetCookieAdapter`, which records only what was set through it. rnet + cannot list the jar, so this covers the origins already requested and no others. + """ + if self._client is None: + return {} + exported: dict[str, str] = {} + for origin in list(self._origins): + try: + raw = self._client.get_cookies(origin) + except Exception: + continue + if raw: + exported[origin] = raw.decode("utf-8", errors="replace") + return exported + + def import_origin_cookies(self, exported: dict[str, str]) -> None: + """Load cookie headers from :meth:`export_origin_cookies` into this HTTP session's client.""" + if not exported: + return + client = self.ensure_client() + for origin, header in exported.items(): + self.set_cookie_header(client, origin, header) + def build_url(self, url: str, params: Optional[Any] = None) -> str: """Encode params into the URL (rnet ignores the params kwarg). @@ -632,6 +765,10 @@ class RnetSession: url = self.build_url(url, kwargs.pop("params", None)) + parsed_origin = urlparse(url) + if parsed_origin.scheme and parsed_origin.netloc: + self._origins.add(f"{parsed_origin.scheme}://{parsed_origin.netloc}") + kwargs.setdefault("allow_redirects", True) if not self.verify: