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: