diff --git a/tests/core/test_adaptive_workers.py b/tests/core/test_adaptive_workers.py index fd3fdd2c..2d9ecc81 100644 --- a/tests/core/test_adaptive_workers.py +++ b/tests/core/test_adaptive_workers.py @@ -1,7 +1,6 @@ -"""Unit tests for the opt-in adaptive worker controller (slow-start then AIMD hill-climb -over segment concurrency) and the plumbing that carries the flag into the downloader. Pure -controller tests only. No threads, sockets, or real downloads; the clock is supplied -explicitly.""" +"""Unit tests for the opt-in adaptive worker controller (AIMD hill-climb over segment +concurrency) and the plumbing that carries the flag into the downloader. Pure controller +tests only. No threads, sockets, or real downloads; the clock is supplied explicitly.""" from __future__ import annotations @@ -55,31 +54,31 @@ def run_tick(c: AdaptiveWorkerController, rate: float, now: float, errors: int = return c.update(now), now -def test_starts_at_min_six_cap() -> None: - assert AdaptiveWorkerController(cap=16).target == 6 - assert AdaptiveWorkerController(cap=8).target == 6 - assert AdaptiveWorkerController(cap=4).target == 4 # min(6, cap) +def test_starts_at_cap() -> None: + assert AdaptiveWorkerController(cap=16).target == 16 + assert AdaptiveWorkerController(cap=4).target == 4 assert AdaptiveWorkerController(cap=1).target == 2 # clamped up to ADAPTIVE_MIN + assert AdaptiveWorkerController(cap=16, start=6).target == 6 # explicit start wins def test_no_op_before_warmup() -> None: - c = AdaptiveWorkerController(cap=16, tick=TICK, window=WINDOW) + c = AdaptiveWorkerController(cap=16, start=6, tick=TICK, window=WINDOW) assert c.update(0.0) == 6 # arms last_tick before any samples are recorded c.record_bytes(500_000, 3.0) # a full tick has elapsed since arming, but under half a tick of samples exists -> hold assert c.update(4.5) == 6 c.record_bytes(500_000, 6.5) - # warmed up (first sample 3.0, now 8.5) -> slow-start probe fires - assert c.update(8.5) == 12 + # warmed up (first sample 3.0, now 8.5) -> the probe fires + assert c.update(8.5) == 8 -def test_ramps_up_while_speed_improves() -> None: - c = AdaptiveWorkerController(cap=16, tick=TICK, window=WINDOW) +def test_climbs_while_speed_improves() -> None: + c = AdaptiveWorkerController(cap=16, start=6, tick=TICK, window=WINDOW) now = 0.0 c.update(now) # arm rate = 1_000_000.0 now = fill_window(c, rate, now) - targets = [c.update(now)] # first evaluation: slow-start doubles 6 -> 12 + targets = [c.update(now)] # first evaluation: 6 -> 8 for _ in range(6): rate *= 2 # sustained >10% windowed gain each tick target, now = run_tick(c, rate, now) @@ -89,23 +88,23 @@ def test_ramps_up_while_speed_improves() -> None: assert max(targets) <= 16 -def test_reverts_on_plateau_and_ends_slow_start() -> None: - c = AdaptiveWorkerController(cap=16, tick=TICK, window=WINDOW) +def test_reverts_on_plateau() -> None: + c = AdaptiveWorkerController(cap=16, start=6, tick=TICK, window=WINDOW) now = 0.0 c.update(now) rate = 1_000_000.0 now = fill_window(c, rate, now) - assert c.update(now) == 12 # first increase slow-starts 6 -> 12 + assert c.update(now) == 8 # first increase 6 -> 8 # hold the rate flat: <10% gain -> the tick reverts the increase to its restore point target, now = run_tick(c, rate, now) assert target == 6 - # slow-start is over: the next successful probe is additive (+2), not a double + # a paying probe climbs again from the restore point target, now = run_tick(c, rate * 4, now) assert target == 8 def test_halves_and_cooldown_on_error_burst() -> None: - c = AdaptiveWorkerController(cap=16, tick=TICK, window=WINDOW) + c = AdaptiveWorkerController(cap=16, start=6, tick=TICK, window=WINDOW) now = 0.0 c.update(now) now = fill_window(c, 1_000_000.0, now) @@ -150,7 +149,7 @@ def test_error_burst_floors_at_min() -> None: def test_tail_guard_suppresses_increase() -> None: # with fewer units of work than the target, the tick must hold instead of probing up - c = AdaptiveWorkerController(cap=16, tick=TICK, window=WINDOW) + c = AdaptiveWorkerController(cap=16, start=6, tick=TICK, window=WINDOW) now = 0.0 c.update(now) # arm at target 6 now = fill_window(c, 1_000_000.0, now) @@ -159,16 +158,16 @@ def test_tail_guard_suppresses_increase() -> None: def test_tail_guard_none_path_unchanged() -> None: # without a tail count, the update still probes upward on the first evaluation - c = AdaptiveWorkerController(cap=16, tick=TICK, window=WINDOW) + c = AdaptiveWorkerController(cap=16, start=6, tick=TICK, window=WINDOW) now = 0.0 c.update(now) now = fill_window(c, 1_000_000.0, now) - assert c.update(now) == 12 # first evaluation slow-starts 6 -> 12 + assert c.update(now) == 8 # first evaluation 6 -> 8 def test_tail_guard_still_halves_on_error_burst() -> None: # the guard suppresses growth but never blocks an error-burst decrease - c = AdaptiveWorkerController(cap=16, tick=TICK, window=WINDOW) + c = AdaptiveWorkerController(cap=16, start=6, tick=TICK, window=WINDOW) now = 0.0 c.update(now) now = fill_window(c, 1_000_000.0, now) diff --git a/tests/core/test_dash_collapse.py b/tests/core/test_dash_collapse.py index 82155ab0..0305778b 100644 --- a/tests/core/test_dash_collapse.py +++ b/tests/core/test_dash_collapse.py @@ -48,6 +48,12 @@ def test_first_segment_gap_after_init_ineligible(): assert collapsible([(URL, "900-1999")], init_len=863) is False +def test_gap_between_segments_ineligible(): + # the merge path would skip 2000-2999; collapsing the parent would include it + segments = [(URL, "863-1999"), (URL, "3000-3999")] + assert collapsible(segments, init_len=863) is False + + def test_overlapping_ranges_ineligible(): segments = [(URL, "863-2500"), (URL, "2000-3999")] assert collapsible(segments, init_len=863) is False diff --git a/tests/core/test_dash_collapse_integration.py b/tests/core/test_dash_collapse_integration.py index 07e2308d..96c6b9c4 100644 --- a/tests/core/test_dash_collapse_integration.py +++ b/tests/core/test_dash_collapse_integration.py @@ -27,6 +27,8 @@ from unshackle.core.tracks.track import DownloadContext # distinctive parent resource for the collapse case; init is its [0, INIT_LEN) prefix PARENT = bytes(i % 256 for i in range(10240)) INIT_LEN = 512 +# last indexed byte for the clipped case; the parent carries 1024 bytes past it +CLIP_END = len(PARENT) - 1025 # separate resources for the non-collapse (mixed-URL) case INIT2 = bytes(range(100)) @@ -111,6 +113,27 @@ def collapse_mpd(srv) -> str: """ +def clipped_mpd(srv) -> str: + # same parent, but the ranges stop short of its end -> the collapsed file must be trimmed + # to the last indexed byte, not left holding the parent's trailing bytes. + return f""" + + {base(srv)} + + + + media.mp4 + + + + + + + + +""" + + def mixed_mpd(srv) -> str: # Two SegmentURLs on distinct paths -> mixed URLs -> predicate False -> segmented path. return f""" @@ -162,6 +185,17 @@ def test_collapse_downloads_parent_whole_in_one_direct_request(server, tmp_path) assert not list(tmp_path.glob("**/*.!dev")) +def test_collapse_trims_parent_bytes_past_the_last_segment(server, tmp_path): + server.routes["/media.mp4"] = PARENT + track = make_track(clipped_mpd(server), server) + ctx = make_ctx(tmp_path, "clipped.mp4") + + DASH.download_track(track, ctx) + + # the parent's trailing 1024 bytes are not part of any segment, so they must not survive + assert ctx.save_path.read_bytes() == PARENT[: CLIP_END + 1] + + def test_non_collapse_takes_segmented_merge_path(server, tmp_path): server.routes["/init.bin"] = INIT2 server.routes["/part0.bin"] = PART0 diff --git a/tests/core/test_retry_after.py b/tests/core/test_retry_after.py new file mode 100644 index 00000000..0f4ba101 --- /dev/null +++ b/tests/core/test_retry_after.py @@ -0,0 +1,69 @@ +"""RnetSession.get_sleep_time must survive hostile or malformed Retry-After headers.""" + +from datetime import datetime, timedelta, timezone +from email.utils import format_datetime +from types import SimpleNamespace + +import pytest + +from unshackle.core.session import RnetSession + + +@pytest.fixture +def session(): + return RnetSession() + + +def response(retry_after=None): + headers = {} if retry_after is None else {"Retry-After": retry_after} + return SimpleNamespace(headers=headers) + + +def test_numeric_retry_after_passes_through(session): + assert session.get_sleep_time(response("3"), 1) == 3.0 + + +def test_huge_numeric_retry_after_is_clamped(session): + assert session.get_sleep_time(response("86400"), 1) == session.max_backoff + + +def test_huge_numeric_retry_after_respects_session_cap(): + session = RnetSession(max_backoff=5.0) + assert session.get_sleep_time(response("86400"), 1) == 5.0 + + +def test_http_date_retry_after_honoured(session): + when = format_datetime(datetime.now(timezone.utc) + timedelta(seconds=10), usegmt=True) + assert 5 <= session.get_sleep_time(response(when), 1) <= 10 + + +def test_naive_http_date_does_not_raise(session): + # "-0000" makes parsedate_to_datetime return a naive datetime + when = format_datetime(datetime.now(timezone.utc) + timedelta(seconds=10)).replace("+0000", "-0000") + assert 5 <= session.get_sleep_time(response(when), 1) <= 10 + + +def test_past_http_date_yields_no_sleep(session): + when = format_datetime(datetime.now(timezone.utc) - timedelta(seconds=30), usegmt=True) + assert session.get_sleep_time(response(when), 1) <= 0 + + +def test_malformed_retry_after_falls_back_to_backoff(session): + value = session.get_sleep_time(response("not-a-date"), 1) + assert 0 < value <= session.max_backoff + + +def test_empty_retry_after_falls_back_to_backoff(session): + value = session.get_sleep_time(response(""), 1) + assert 0 < value <= session.max_backoff + + +def test_exponential_backoff_without_retry_after(session): + assert session.get_sleep_time(response(), 0) == 0.0 + base = session.backoff_factor * 4 + assert base * 0.9 <= session.get_sleep_time(response(), 3) <= base * 1.1 + assert session.get_sleep_time(response(), 30) == session.max_backoff + + +def test_no_response_still_backs_off(session): + assert session.get_sleep_time(None, 30) == session.max_backoff diff --git a/tests/downloaders/test_failure_surfacing.py b/tests/downloaders/test_failure_surfacing.py index ea8f5861..176b9d4b 100644 --- a/tests/downloaders/test_failure_surfacing.py +++ b/tests/downloaders/test_failure_surfacing.py @@ -6,7 +6,7 @@ import pytest import requests as rq from unshackle.core.constants import DOWNLOAD_CANCELLED -from unshackle.core.downloaders.requests import download, has_range_header, parse_content_range, request_range_start +from unshackle.core.downloaders.requests import download, has_range_header, parse_content_range, request_range from unshackle.core.downloaders.requests import requests as requests_downloader # the package re-exports the `requests` function, shadowing the module of the same name @@ -76,6 +76,12 @@ class _Server: self.send_206(conn, BODY[5:], f"bytes 5-{len(BODY) - 1}/{len(BODY)}") elif self.mode == "resume_missing_content_range": self.send_206(conn, BODY) + elif self.mode == "range_capped_206": + # honors the requested start but caps the span, as size-capping CDNs do + start, _, end = headers["range"].removeprefix("bytes=").partition("-") + start, end = int(start), int(end) + capped = start + (end - start) // 2 + self.send_206(conn, BODY[start : capped + 1], f"bytes {start}-{capped}/{len(BODY)}") else: conn.sendall( b"HTTP/1.1 416 Range Not Satisfiable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" @@ -197,6 +203,23 @@ def test_slice_segment_rejects_a_206_for_the_wrong_range(tmp_path, server): assert len(server.requests) == downloader.MAX_ATTEMPTS +@pytest.mark.parametrize("server", ["range_capped_206"], indirect=True) +def test_slice_segment_rejects_a_206_capped_short_of_the_requested_end(tmp_path, server): + # the body is a valid prefix of the slice, so only the Content-Range end catches it; + # writing it would finalize a truncated segment and corrupt the merge + save_path = tmp_path / "0.mp4" + with pytest.raises(OSError): + list( + download( + url=server.url, + save_path=save_path, + session=rq.Session(), + headers={"Range": f"bytes=0-{len(BODY) - 1}"}, + ) + ) + assert not save_path.exists() + + @pytest.mark.parametrize("server", ["range_unsatisfiable"], indirect=True) def test_slice_segment_never_gets_a_resume_range(tmp_path, server): save_path = tmp_path / "0.mp4" @@ -295,9 +318,9 @@ def test_parse_content_range(value, expected): ({}, None), ({"Accept": "*/*"}, None), ({"Range": "bytes=-500"}, None), - ({"Range": "bytes=100-"}, 100), - ({"range": "bytes=100-227"}, 100), + ({"Range": "bytes=100-"}, (100, None)), + ({"range": "bytes=100-227"}, (100, 227)), ], ) -def test_request_range_start(headers, expected): - assert request_range_start(headers) == expected +def test_request_range(headers, expected): + assert request_range(headers) == expected diff --git a/tests/downloaders/test_session_spec.py b/tests/downloaders/test_session_spec.py new file mode 100644 index 00000000..cf41156c --- /dev/null +++ b/tests/downloaders/test_session_spec.py @@ -0,0 +1,125 @@ +"""Session spec fidelity across the multiprocess spawn boundary. + +``build_session_spec``/``rebuild_session`` are how a spawned child gets the parent's session +state. The cross-process half needs a live CDN to observe, so these tests pin the two halves +that can be checked deterministically: the spec is picklable (spawn requires it), and a +rebuilt session carries the same cookie domains as the parent. A spec that flattens cookies +to name/value rebuilds every cookie scoped to localhost, so a cookie-authed CDN sees none of +them and 403s only when ``--download-processes`` is above 1. +""" + +import pickle + +from requests import Session + +from unshackle.core.downloaders.requests import build_session_spec, rebuild_session +from unshackle.core.session import RnetSession +from unshackle.core.session import session as make_session + + +def rnet_session_with_cookies() -> RnetSession: + """RnetSession carrying two domain-scoped cookies plus one set from a plain dict. + + The dict-set cookie lands in the adapter's flat map only, never in its domain map, so it + is the case a domain-only serialization would silently drop. + """ + ns = make_session("Chrome131") + ns.cookies.set("cdn_token", "abc", domain="cdn.example.com") + ns.cookies.set("api_token", "def", domain="api.example.com") + ns.cookies.update({"loose": "ghi"}) + return ns + + +def test_rnet_spec_round_trips_cookie_domains() -> None: + ns = rnet_session_with_cookies() + + spec = build_session_spec(ns) + assert spec is not None + rebuilt = rebuild_session(pickle.loads(pickle.dumps(spec)), 1) + assert rebuilt is not None + + assert rebuilt.cookies.get_dict_by_domain() == ns.cookies.get_dict_by_domain() + assert rebuilt.cookies.get_dict() == ns.cookies.get_dict() + assert rebuilt.cookies.get("cdn_token", domain="cdn.example.com") == "abc" + assert rebuilt.cookies.get("api_token", domain="api.example.com") == "def" + assert rebuilt.cookies.get_dict(domain="localhost") == {} + + +def test_rnet_spec_keeps_dict_set_cookies() -> None: + ns = rnet_session_with_cookies() + + spec = build_session_spec(ns) + assert spec is not None + rebuilt = rebuild_session(spec, 1) + assert rebuilt is not None + + assert rebuilt.cookies["loose"] == "ghi" + + +def test_rnet_spec_cookies_carry_domain_information() -> None: + ns = rnet_session_with_cookies() + + spec = build_session_spec(ns) + assert spec is not None + + cookies = spec["cookies"] + assert set(cookies) == {"cdn.example.com", "api.example.com", ""} + assert all(isinstance(jar, dict) for jar in cookies.values()) + assert cookies["cdn.example.com"] == {"cdn_token": "abc"} + + +def test_rnet_spec_is_picklable() -> None: + spec = build_session_spec(rnet_session_with_cookies()) + assert spec is not None + + assert pickle.loads(pickle.dumps(spec)) == spec + + +def test_rnet_spec_without_impersonate_is_not_rebuildable() -> None: + assert build_session_spec(RnetSession()) is None + + +def test_rebuilt_requests_session_pool_is_sized_to_worker_count() -> None: + """A child rebuilding a requests session must size its pool, not inherit urllib3's default 10.""" + spec = build_session_spec(Session()) + assert spec is not None + + rebuilt = rebuild_session(spec, 24) + assert rebuilt is not None + + for scheme in ("https://example.com", "http://example.com"): + adapter = rebuilt.get_adapter(scheme) + assert adapter._pool_maxsize == 24 + assert adapter._pool_connections == 24 + assert adapter._pool_block is True + assert adapter.poolmanager.connection_pool_kw["maxsize"] == 24 + + +def test_rebuilt_rnet_session_has_no_adapters() -> None: + """rnet pools internally; the sizing fix must not reach into that path.""" + spec = build_session_spec(rnet_session_with_cookies()) + assert spec is not None + + rebuilt = rebuild_session(spec, 24) + assert rebuilt is not None + + assert not isinstance(rebuilt, Session) + assert not hasattr(rebuilt, "adapters") + + +def test_unrebuildable_spec_stays_none() -> None: + assert rebuild_session({"kind": "none"}, 24) is None + + +def test_requests_spec_round_trips_cookie_domains() -> None: + rs = Session() + rs.cookies.set("cdn_token", "abc", domain="cdn.example.com") + rs.cookies.set("api_token", "def", domain="api.example.com") + + spec = build_session_spec(rs) + assert spec is not None + rebuilt = rebuild_session(pickle.loads(pickle.dumps(spec)), 1) + assert rebuilt is not None + + assert rebuilt.cookies.get_dict(domain="cdn.example.com") == {"cdn_token": "abc"} + assert rebuilt.cookies.get_dict(domain="api.example.com") == {"api_token": "def"} diff --git a/unshackle/core/downloaders/requests.py b/unshackle/core/downloaders/requests.py index 50ecd700..3e351404 100644 --- a/unshackle/core/downloaders/requests.py +++ b/unshackle/core/downloaders/requests.py @@ -166,10 +166,9 @@ def adaptive_chunk_size(content_length: int) -> int: return min(MAX_CHUNK, max(MIN_CHUNK, content_length // 4)) -# Adaptive worker controller (opt-in): slow-start then AIMD hill-climb over segment concurrency +# Adaptive worker controller (opt-in): AIMD hill-climb over segment concurrency ADAPTIVE_TICK = 4.0 # seconds between target re-evaluations -ADAPTIVE_START = 6 # initial worker target, clamped to cap -ADAPTIVE_STEP = 2 # additive increase per tick once slow-start ends +ADAPTIVE_STEP = 2 # additive increase per tick ADAPTIVE_MIN = 2 # never drop below this many workers ADAPTIVE_ERROR_BURST = 3 # errors within a tick that trigger multiplicative decrease ADAPTIVE_PLATEAU_GAIN = 1.10 # min speed ratio to justify keeping an increase @@ -180,17 +179,17 @@ class AdaptiveWorkerController: Pure and synchronous (no threads or sockets), so the policy is unit-testable. ``update`` is evaluated once per ``tick``, starting once half a tick of throughput - samples exists. It slow-starts (doubles the target) while every probe keeps paying - off, then drops to AIMD (+``ADAPTIVE_STEP``) after the first plateau or error burst: - a plateau reverts the last increase, an error burst halves the target with a - one-tick cooldown. Target is always clamped to ``[ADAPTIVE_MIN, cap]``. + 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, a probe upward + that plateaus is reverted, and otherwise it climbs by ``ADAPTIVE_STEP``. Target is + always clamped to ``[ADAPTIVE_MIN, cap]``. """ def __init__( self, cap: int, start: Optional[int] = None, tick: float = ADAPTIVE_TICK, window: float = SPEED_ROLLING_WINDOW ) -> None: self.cap = max(ADAPTIVE_MIN, cap) - base = ADAPTIVE_START if start is None else start + base = self.cap if start is None else start self.target = max(ADAPTIVE_MIN, min(base, self.cap)) self.tick = tick self.window = window @@ -202,7 +201,6 @@ class AdaptiveWorkerController: self._speed_before_increase = 0.0 self._target_before_increase = self.target # restore point for a plateau revert self._cooldown = False - self._ramping = True # slow-start: double per probe until the first plateau or error burst def prune(self, now: float) -> None: cutoff = now - self.window @@ -271,24 +269,20 @@ class AdaptiveWorkerController: self.target = max(ADAPTIVE_MIN, self.target // 2) self._cooldown = True self._last_action = None - self._ramping = False reason = "error_burst" elif inflight_plus_remaining is not None and inflight_plus_remaining < self.target: # 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: - # last probe upward did not pay off: revert it and hold. step size varies - # during slow-start, hence the recorded restore point rather than -STEP + # last probe upward did not pay off: revert it and hold self.target = max(ADAPTIVE_MIN, self._target_before_increase) self._last_action = None - self._ramping = False reason = "plateau" elif self.target < self.cap: self._speed_before_increase = speed_now self._target_before_increase = self.target - step = self.target if self._ramping else ADAPTIVE_STEP - self.target = min(self.cap, self.target + step) + self.target = min(self.cap, self.target + ADAPTIVE_STEP) self._last_action = "increase" reason = "increase" else: @@ -436,11 +430,18 @@ 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 request_range_start(headers: MutableMapping[str, Any]) -> Optional[int]: - """Start byte of the request's ``Range`` header; None when absent or not ``bytes=start-...``.""" +def request_range(headers: MutableMapping[str, Any]) -> Optional[tuple[int, Optional[int]]]: + """ + The request's ``Range`` as ``(start, end)``; None when absent or not ``bytes=start-...``. + + ``end`` is None for an open-ended ``bytes=start-``, which the caller cannot hold a + response to. + """ value = str(next((v for k, v in headers.items() if str(k).lower() == "range"), "")) - match = re.match(r"bytes=(\d+)-", value.strip()) - return int(match.group(1)) if match else None + match = re.match(r"bytes=(\d+)-(\d+)?", value.strip()) + if not match: + return None + return int(match.group(1)), int(match.group(2)) if match.group(2) else None def force_unblock_stream(stream: Any) -> None: @@ -596,6 +597,11 @@ def dispatch_parts( pool.shutdown(wait=False, cancel_futures=True) raise finally: + # abort outlives the shutdown so a leaked part exits at its next check instead of + # writing into save_path after the caller moved on; a consumer that abandons the + # generator (GeneratorExit) reaches here with abort still unset + if not completed: + abort.set() pool.shutdown(wait=worker_error, cancel_futures=True) if completed: control_file.unlink(missing_ok=True) @@ -803,16 +809,23 @@ def download( raise IOError(f"expected 206 for byte-range segment, got {stream.status_code}") if item_range and not part_mode and stream.status_code == 206: - # a 206 for a different range than the slice asked for (e.g. the whole - # parent) would be written as the segment and silently corrupt the merge - slice_start = request_range_start(req_headers) - if slice_start is not None: + # a 206 for a different range than the slice asked for (e.g. the whole parent, + # or the short range a size-capping CDN chose to serve) would be written as the + # whole segment and silently corrupt the merge + slice_range = request_range(req_headers) + if slice_range is not None: + slice_start, slice_end = slice_range content_range = parse_content_range( stream.headers.get("Content-Range") or stream.headers.get("content-range") ) - if content_range is None or content_range[0] != slice_start: + if ( + content_range is None + or content_range[0] != slice_start + or (slice_end is not None and content_range[1] != slice_end) + ): raise IOError( - f"byte-range segment got Content-Range {content_range!r}, expected start {slice_start}" + f"byte-range segment got Content-Range {content_range!r}, " + f"expected {slice_start}-{slice_end if slice_end is not None else ''}" ) if resumed and content_encoded: @@ -969,7 +982,10 @@ def download( if abort.wait(delay): return else: - time.sleep(delay) + # no batch abort on this path, but a cross-track cancel must still wake a + # parked worker rather than stall teardown for the full backoff + if DOWNLOAD_CANCELLED.wait(delay): + return attempts += 1 @@ -997,14 +1013,19 @@ def build_session_spec(session: Optional[Any]) -> Optional[dict[str, Any]]: "kind": "rnet", "impersonate": name, "headers": dict(session.headers), - "cookies": session.cookies.get_dict(), + "cookies": session.cookies.get_dict_by_domain(), "proxy": session.proxies.get("all") or session.proxies.get("https") or session.proxies.get("http"), } return None -def rebuild_session(spec: dict[str, Any]) -> Optional[Any]: - """Reconstruct a session inside a child process from a spec built by ``build_session_spec``.""" +def rebuild_session(spec: dict[str, Any], max_workers: int) -> Optional[Any]: + """Reconstruct a session inside a child process from a spec built by ``build_session_spec``. + + ``max_workers`` sizes the connection pool to what this child actually runs; the rebuilt + session belongs to one child alone, so mounting an adapter here is safe where remounting a + caller-provided session in ``requests()`` is not. + """ kind = spec["kind"] if kind == "requests": rs = Session() @@ -1013,6 +1034,9 @@ def rebuild_session(spec: dict[str, Any]) -> Optional[Any]: rs.cookies.update(spec["cookies"]) if spec.get("proxies"): rs.proxies.update(spec["proxies"]) + adapter = HTTPAdapter(pool_connections=max_workers, pool_maxsize=max_workers, pool_block=True) + rs.mount("https://", adapter) + rs.mount("http://", adapter) return rs if kind == "rnet": from unshackle.core.session import session as make_session @@ -1020,8 +1044,13 @@ def rebuild_session(spec: dict[str, Any]) -> Optional[Any]: ns = make_session(spec["impersonate"]) if spec.get("headers"): ns.headers.update(spec["headers"]) - if spec.get("cookies"): - ns.cookies.update(spec["cookies"]) + for domain, jar in spec.get("cookies", {}).items(): + if domain: + for name, value in jar.items(): + ns.cookies.set(name, value, domain=domain) + else: + # no domain to scope to; update() reproduces the parent's localhost-scoped state + ns.cookies.update(jar) if spec.get("proxy"): ns.proxies.update({"all": spec["proxy"]}) return ns @@ -1047,7 +1076,7 @@ def mp_worker(queue: Any, kwargs: dict[str, Any]) -> None: if name in ("READ_TIMEOUT", "RETRY_WAIT", "HEDGE_MIN_WAIT"): setattr(sys.modules[__name__], name, float(value)) spec = kwargs.pop("_session_spec") - kwargs["session"] = rebuild_session(spec) + kwargs["session"] = rebuild_session(spec, int(kwargs.get("max_workers") or 1)) for event in requests(**kwargs): queue.put(event) except Exception: @@ -1426,9 +1455,9 @@ def requests( save_path = url_item.get("save_path") if save_path: sp = Path(save_path) - for target in (sp, sp.with_name(f"{sp.name}.!dev")): + for stray in (sp, sp.with_name(f"{sp.name}.!dev")): try: - target.unlink(missing_ok=True) + stray.unlink(missing_ok=True) except OSError: pass if not ranged_used: @@ -1815,7 +1844,9 @@ def requests( raise finally: # losers must close their handles before merge sweeps *.!dev (WinError 32); - # no wait on cancel/fail: merge never runs and workers may be blocked in reads. + # no wait on cancel/fail: the merge below never runs and workers may be blocked + # in reads. The caller's own cleanup does still sweep on failure, so a worker + # parked in a read can briefly hold a handle there (Windows). # 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() diff --git a/unshackle/core/manifests/dash.py b/unshackle/core/manifests/dash.py index df9a903f..0f11a7ef 100644 --- a/unshackle/core/manifests/dash.py +++ b/unshackle/core/manifests/dash.py @@ -512,7 +512,10 @@ class DASH: status_update["downloaded"] = f"DASH {downloaded}" progress(**status_update) - if not collapse_single_url: + if collapse_single_url: + with open(save_path, "r+b") as collapsed: + collapsed.truncate(int(segments[-1][1].split("-")[1]) + 1) + else: # Verify output directory exists and contains files if not save_dir.exists(): error_msg = f"Output directory does not exist: {save_dir}" @@ -626,8 +629,9 @@ class DASH: prefix is the init segment, so the whole resource can be downloaded in one pass. Returns False for subtitle tracks, an unknown init length, mixed URLs, malformed - ranges, a gap between the init prefix and the first media byte, and ranges that - overlap or regress in document order. + ranges, and anything but one gapless run in document order starting at the init + prefix. A gap would pull bytes the segmented merge path excludes (a period the + manifest filters out) into the collapsed file, so it fails closed. """ if is_subtitle or init_len is None or not segments: return False @@ -648,12 +652,9 @@ class DASH: return False ranges.append((start, end)) - if min(start for start, _ in ranges) != init_len: - return False - prev_end = init_len - 1 for start, end in ranges: - if start <= prev_end: + if start != prev_end + 1: return False prev_end = end return True diff --git a/unshackle/core/session.py b/unshackle/core/session.py index 74cb2f26..5e42eb7c 100644 --- a/unshackle/core/session.py +++ b/unshackle/core/session.py @@ -4,6 +4,7 @@ from __future__ import annotations import http import logging +import math import random import time from collections.abc import Iterator, MutableMapping @@ -405,6 +406,21 @@ class RnetCookieAdapter(MutableMapping): return dict(self._cookies.get(domain, {})) return dict(self._flat) + def get_dict_by_domain(self) -> dict[str, dict[str, str]]: + """Return cookies grouped by domain, domain-less ones under ``""``. + + rnet scopes each cookie to the host of the URL it was set on, so a copy made with + ``get_dict`` can only ever be replayed against one host. Cookies that arrived as a + plain dict are recorded flat only, and are reported under the empty domain to match + the localhost fallback in :meth:`flush_to_client`. + """ + grouped = {domain: dict(cookies) for domain, cookies in self._cookies.items() if cookies} + scoped = {name for cookies in grouped.values() for name in cookies} + leftovers = {name: value for name, value in self._flat.items() if name not in scoped} + if leftovers: + grouped[""] = {**grouped.get("", {}), **leftovers} + return grouped + def clear(self, domain: Optional[str] = None, path: Optional[str] = None, name: Optional[str] = None) -> None: """Remove cookies (requests RequestsCookieJar compat). @@ -616,11 +632,22 @@ class RnetSession: if response: retry_after = response.headers.get("Retry-After") if retry_after: + wait: Optional[float] = None try: - return float(retry_after) + wait = float(retry_after) except ValueError: - if retry_date := parsedate_to_datetime(retry_after): - return (retry_date - datetime.now(timezone.utc)).total_seconds() + try: + retry_date = parsedate_to_datetime(retry_after) + if retry_date.tzinfo is None: + retry_date = retry_date.replace(tzinfo=timezone.utc) + wait = (retry_date - datetime.now(timezone.utc)).total_seconds() + except Exception: + # parsedate_to_datetime itself raises ValueError on malformed dates + # (Python >= 3.10); an unusable header must not escape the retry loop + wait = None + if wait is not None and math.isfinite(wait): + # a hostile Retry-After (e.g. 86400) would otherwise park the caller for a day + return min(wait, self.max_backoff) if attempt == 0: return 0.0