mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-12 12:05:07 +00:00
fix(downloader): harden ranged reads, teardown and multiprocess sessions
Validates the 206 range end, bounds the DASH collapse, and preserves cookie domains and pool size in spawned children.
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
</MPD>"""
|
||||
|
||||
|
||||
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"""<?xml version="1.0"?>
|
||||
<MPD xmlns="urn:mpeg:dash:schema:mpd:2011" type="static" profiles="urn:mpeg:dash:profile:isoff-on-demand:2011">
|
||||
<BaseURL>{base(srv)}</BaseURL>
|
||||
<Period id="0">
|
||||
<AdaptationSet contentType="video" mimeType="video/mp4" lang="en">
|
||||
<Representation id="v0" codecs="avc1.640028" bandwidth="1000000" width="1920" height="1080">
|
||||
<BaseURL>media.mp4</BaseURL>
|
||||
<SegmentList timescale="1">
|
||||
<Initialization range="0-{INIT_LEN - 1}"/>
|
||||
<SegmentURL mediaRange="{INIT_LEN}-5119"/>
|
||||
<SegmentURL mediaRange="5120-{CLIP_END}"/>
|
||||
</SegmentList>
|
||||
</Representation>
|
||||
</AdaptationSet>
|
||||
</Period>
|
||||
</MPD>"""
|
||||
|
||||
|
||||
def mixed_mpd(srv) -> str:
|
||||
# Two SegmentURLs on distinct paths -> mixed URLs -> predicate False -> segmented path.
|
||||
return f"""<?xml version="1.0"?>
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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"}
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user