kenzuya/unshackle
1
0
Fork 2
mirror of https://github.com/unshackle-dl/unshackle.git synced 2026-11-11 19:45:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(downloader): cut per-segment overhead, bound stalled connections

Browse Source 
Segment tmp-file+rename scheme, request timeouts, hedged tail requests, streamed ISM/DASH merges, rnet client knobs, curl_impersonate -> network.
This commit is contained in:
imSp4rky
2026-07-02 09:20:03 -06:00
parent 9835f45b1b
commit ef2e880c7f
6 changed files with 300 additions and 162 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/NETWORK_CONFIG.md
+28 -4
View File
@@ -243,13 +243,15 @@ fingerprinting. `RnetSession` is a drop-in `requests.Session` replacement and is
- Native rnet proxy support (HTTP, HTTPS, SOCKS5) — used by all proxy providers below.
- Cookie-jar and `requests`-style `data=` / `json=` / `headers=` kwargs for compatibility.
The legacy `curl_cffi` backend has been removed. The config key is still spelled
`curl_impersonate` for backward compatibility, but its value now selects an rnet preset.
The legacy `curl_cffi` backend has been removed.
### curl_impersonate (dict)
### network (dict)
HTTP session settings: which client to impersonate and how the underlying rnet
client behaves.
```yaml
curl_impersonate:
network:
browser: Chrome131 # exact rnet.Impersonate preset name
```
@@ -257,4 +259,26 @@ curl_impersonate:
`Edge101`, `Firefox133`, `Safari18`, `OkHttp4_12`). See the rnet README for the full list.
Default when unset: `Chrome131`.
The same section accepts optional rnet client tuning keys, forwarded verbatim:
```yaml
network:
browser: Chrome131
http1_only: true # force HTTP/1.1 instead of negotiated HTTP/2
# http2_only: true
# pool_max_size: 64 # max total connections in the pool
# pool_max_idle_per_host: 32
# tcp_nodelay: true
```
!!! note "Renamed from `curl_impersonate`"
This section was previously named `curl_impersonate` (after the removed backend).
The old spelling still works but is deprecated; rename it to `network`.
!!! tip "When to force HTTP/1.1"
HTTP/2 is the right default on large CDNs. Force `http1_only: true` when a host
throttles per connection or has a slow origin behind HTTP/2 flow control — in
benchmarks this gained 30–50% on such hosts, while costing up to 27% on fast
CDNs. Set it only when measurements support it.
---
unshackle/core/config.py
+8 -1
View File
@@ -57,7 +57,14 @@ class Config:
self.dl: dict = kwargs.get("dl") or {}
self.cdm: dict = kwargs.get("cdm") or {}
self.chapter_fallback_name: str = kwargs.get("chapter_fallback_name") or ""
self.curl_impersonate: dict = kwargs.get("curl_impersonate") or {}
self.network: dict = kwargs.get("network") or kwargs.get("curl_impersonate") or {}
self.curl_impersonate: dict = self.network
if "curl_impersonate" in kwargs and "network" not in kwargs:
warnings.warn(
"the 'curl_impersonate' config section is deprecated, rename it to 'network'.",
DeprecationWarning,
stacklevel=2,
)
self.remote_cdm: list[dict] = kwargs.get("remote_cdm") or []
self.credentials: dict = kwargs.get("credentials") or {}
self.firefox_cookies: dict = kwargs.get("firefox_cookies") or {}
unshackle/core/downloaders/requests.py
+216 -136
View File
@@ -1,12 +1,14 @@
import math
import os
import statistics
import threading
import time
from concurrent.futures import FIRST_COMPLETED, wait
from concurrent.futures.thread import ThreadPoolExecutor
from http.cookiejar import CookieJar
from pathlib import Path
from queue import Empty, Queue
from typing import Any, Generator, MutableMapping, Optional, Union
from typing import Any, Callable, Generator, MutableMapping, Optional, Union
from requests import Session
from requests.adapters import HTTPAdapter
@@ -19,6 +21,14 @@ MAX_ATTEMPTS = 5
RETRY_WAIT = 2
PROGRESS_WINDOW = 2
# read timeout bounds the gap between chunks so a quiet connection errors instead of hanging
CONNECT_TIMEOUT = 10
READ_TIMEOUT = 30
# re-request a tail segment stuck past max(HEDGE_FACTOR * median segment time, HEDGE_MIN_WAIT)
HEDGE_FACTOR = 3
HEDGE_MIN_WAIT = 5.0
# Adaptive chunk sizing — benchmarked optimal range
MIN_CHUNK = 524_288 # 512KB
MAX_CHUNK = 4_194_304 # 4MB
@@ -51,6 +61,10 @@ def _is_rnet_session(session: Any) -> bool:
def _probe_ranged(url: str, session: Any, **kwargs: Any) -> tuple[int, bool]:
headers = {**(kwargs.get("headers") or {}), "Range": "bytes=0-0"}
rest = {k: v for k, v in kwargs.items() if k != "headers"}
if _is_rnet_session(session):
rest.setdefault("read_timeout", READ_TIMEOUT)
else:
rest.setdefault("timeout", (CONNECT_TIMEOUT, READ_TIMEOUT))
try:
resp = session.get(url, stream=True, headers=headers, **rest)
except Exception:
@@ -179,6 +193,7 @@ def download(
segmented: bool = False,
part_offset: Optional[int] = None,
part_end: Optional[int] = None,
claimed: Optional[Callable[[], bool]] = None,
**kwargs: Any,
) -> Generator[dict[str, Any], None, None]:
"""
@@ -207,6 +222,8 @@ def download(
no truncate, no skip-if-exists, no control file; emits only `advance`
events; retries resume mid-part via Range.
part_end: Inclusive end byte of the part. Required when `part_offset` is set.
claimed: Optional predicate checked before every attempt; when it returns True
the download stops silently (another worker already delivered this file).
kwargs: Any extra keyword arguments to pass to the session.get() call. Use this
for one-time request changes like a header, cookie, or proxy. For example,
to request Byte-ranges use e.g., `headers={"Range": "bytes=0-128"}`.
@@ -214,161 +231,170 @@ def download(
session = session or Session()
part_mode = part_offset is not None and part_end is not None
save_dir = save_path.parent
control_file = save_path.with_name(f"{save_path.name}.!dev")
save_dir.mkdir(parents=True, exist_ok=True)
# partial data lives at the .!dev name and is renamed into place on completion:
# bare save_path = done, bare .!dev = resumable, no per-segment marker files
tmp_file = save_path.with_name(f"{save_path.name}.!dev")
if not segmented:
save_path.parent.mkdir(parents=True, exist_ok=True)
resume_offset = 0
if not part_mode:
if control_file.exists() and save_path.exists():
resume_offset = save_path.stat().st_size
elif control_file.exists():
control_file.unlink()
elif save_path.exists():
yield dict(file_downloaded=save_path, written=save_path.stat().st_size)
return
control_file.write_bytes(b"")
if save_path.exists():
if tmp_file.exists():
# legacy layout: empty marker beside a partial save_path, restart clean
tmp_file.unlink()
save_path.unlink()
else:
yield dict(file_downloaded=save_path, written=save_path.stat().st_size)
return
elif tmp_file.exists():
resume_offset = tmp_file.stat().st_size
_time = time.time
use_raw = _is_requests_session(session)
attempts = 1
completed = False
written = 0
try:
while True:
if not part_mode:
written = 0
last_speed_refresh = _time()
while True:
if claimed is not None and claimed():
return
if not part_mode:
written = 0
last_speed_refresh = _time()
try:
use_rnet = _is_rnet_session(session)
try:
use_rnet = _is_rnet_session(session)
request_kwargs = dict(kwargs)
if part_mode:
req_headers = dict(request_kwargs.get("headers", {}) or {})
req_headers["Range"] = f"bytes={part_offset + written}-{part_end}"
request_kwargs["headers"] = req_headers
elif resume_offset > 0:
req_headers = dict(request_kwargs.get("headers", {}) or {})
req_headers["Range"] = f"bytes={resume_offset}-"
request_kwargs["headers"] = req_headers
request_kwargs = dict(kwargs)
if use_rnet:
request_kwargs.setdefault("read_timeout", READ_TIMEOUT)
else:
request_kwargs.setdefault("timeout", (CONNECT_TIMEOUT, READ_TIMEOUT))
if part_mode:
req_headers = dict(request_kwargs.get("headers", {}) or {})
req_headers["Range"] = f"bytes={part_offset + written}-{part_end}"
request_kwargs["headers"] = req_headers
elif resume_offset > 0:
req_headers = dict(request_kwargs.get("headers", {}) or {})
req_headers["Range"] = f"bytes={resume_offset}-"
request_kwargs["headers"] = req_headers
stream = session.get(url, stream=True, **request_kwargs)
stream.raise_for_status()
stream = session.get(url, stream=True, **request_kwargs)
stream.raise_for_status()
resumed = (not part_mode) and resume_offset > 0 and stream.status_code == 206
if (not part_mode) and resume_offset > 0 and not resumed:
resume_offset = 0
if part_mode and stream.status_code != 206:
raise IOError(f"expected 206 for ranged part, got {stream.status_code}")
if use_rnet:
content_length = stream.content_length or 0
else:
try:
content_length = int(stream.headers.get("Content-Length", "0"))
if stream.headers.get("Content-Encoding", "").lower() in ["gzip", "deflate", "br"]:
content_length = 0
except ValueError:
resumed = (not part_mode) and resume_offset > 0 and stream.status_code == 206
if (not part_mode) and resume_offset > 0 and not resumed:
resume_offset = 0
if part_mode and stream.status_code != 206:
raise IOError(f"expected 206 for ranged part, got {stream.status_code}")
if use_rnet:
content_length = stream.content_length or 0
else:
try:
content_length = int(stream.headers.get("Content-Length", "0"))
if stream.headers.get("Content-Encoding", "").lower() in ["gzip", "deflate", "br"]:
content_length = 0
except ValueError:
content_length = 0
chunk_size = _adaptive_chunk_size(content_length)
total_size = (resume_offset + content_length) if resumed and content_length > 0 else content_length
chunk_size = _adaptive_chunk_size(content_length)
total_size = (resume_offset + content_length) if resumed and content_length > 0 else content_length
if not segmented and not part_mode:
if total_size > 0:
yield dict(total=total_size)
else:
yield dict(total=None)
if resumed and resume_offset > 0:
yield dict(advance=resume_offset)
if part_mode:
file_mode = "r+b"
file_buffering = 0
if not segmented and not part_mode:
if total_size > 0:
yield dict(total=total_size)
else:
file_mode = "ab" if resumed else "wb"
file_buffering = 1_048_576
with open(save_path, file_mode, buffering=file_buffering) as f:
if part_mode:
f.seek(part_offset + written)
elif not resumed and content_length > 0:
f.truncate(content_length)
f.seek(0)
_write = f.write
if use_rnet:
chunks = stream.stream()
elif use_raw:
chunks = iter(lambda: stream.raw.read(chunk_size), b"")
else:
chunks = stream.iter_content(chunk_size=chunk_size)
_data_accumulated = 0
_bytes_since_yield = 0
emit_progress = (not segmented) or part_mode
for chunk in chunks:
if DOWNLOAD_CANCELLED.is_set():
break
_write(chunk)
download_size = len(chunk)
written += download_size
if emit_progress:
_bytes_since_yield += download_size
_data_accumulated += download_size
now = _time()
time_since = now - last_speed_refresh
if time_since > PROGRESS_WINDOW:
yield dict(advance=_bytes_since_yield)
_bytes_since_yield = 0
if not part_mode:
download_speed = math.ceil(_data_accumulated / (time_since or 1))
yield dict(downloaded=f"{filesize.decimal(download_speed)}/s")
last_speed_refresh = now
_data_accumulated = 0
if emit_progress and _bytes_since_yield > 0:
yield dict(advance=_bytes_since_yield)
try:
stream.close()
except Exception:
pass
if not part_mode and not resumed and content_length > 0 and written != content_length:
f.truncate(written)
yield dict(total=None)
if resumed and resume_offset > 0:
yield dict(advance=resume_offset)
if part_mode:
file_mode = "r+b"
file_buffering = 0
else:
file_mode = "ab" if resumed else "wb"
file_buffering = 1_048_576
with open(save_path if part_mode else tmp_file, file_mode, buffering=file_buffering) as f:
if part_mode:
expected = part_end - part_offset + 1
if written < expected:
raise IOError(f"Failed to read part {part_offset}-{part_end}: got {written}/{expected}")
elif not segmented and content_length and written < content_length:
raise IOError(f"Failed to read {content_length} bytes from the track URI.")
f.seek(part_offset + written)
elif not resumed and content_length > 0:
f.truncate(content_length)
f.seek(0)
_write = f.write
if use_rnet:
chunks = stream.stream()
elif use_raw:
chunks = iter(lambda: stream.raw.read(chunk_size), b"")
else:
chunks = stream.iter_content(chunk_size=chunk_size)
_data_accumulated = 0
_bytes_since_yield = 0
emit_progress = (not segmented) or part_mode
for chunk in chunks:
if DOWNLOAD_CANCELLED.is_set():
break
_write(chunk)
download_size = len(chunk)
written += download_size
if emit_progress:
_bytes_since_yield += download_size
_data_accumulated += download_size
now = _time()
time_since = now - last_speed_refresh
if time_since > PROGRESS_WINDOW:
yield dict(advance=_bytes_since_yield)
_bytes_since_yield = 0
if not part_mode:
download_speed = math.ceil(_data_accumulated / (time_since or 1))
yield dict(downloaded=f"{filesize.decimal(download_speed)}/s")
last_speed_refresh = now
_data_accumulated = 0
if emit_progress and _bytes_since_yield > 0:
yield dict(advance=_bytes_since_yield)
if not part_mode:
yield dict(file_downloaded=save_path, written=resume_offset + written)
if segmented:
yield dict(advance=1)
completed = True
break
except Exception:
try:
stream.close()
except Exception:
pass
if DOWNLOAD_CANCELLED.is_set() or attempts == MAX_ATTEMPTS:
if part_mode and not DOWNLOAD_CANCELLED.is_set():
raise
return
if not part_mode and save_path.exists():
resume_offset = save_path.stat().st_size
time.sleep(RETRY_WAIT)
attempts += 1
finally:
if completed and not part_mode:
control_file.unlink(missing_ok=True)
if not part_mode and not resumed and content_length > 0 and written != content_length:
f.truncate(written)
if DOWNLOAD_CANCELLED.is_set() and (not content_length or written < content_length):
# cancelled mid-stream: keep the partial .!dev for resume, never finalize
return
if part_mode:
expected = part_end - part_offset + 1
if written < expected:
raise IOError(f"Failed to read part {part_offset}-{part_end}: got {written}/{expected}")
elif not segmented and content_length and written < content_length:
raise IOError(f"Failed to read {content_length} bytes from the track URI.")
if not part_mode:
os.replace(tmp_file, save_path)
yield dict(file_downloaded=save_path, written=resume_offset + written)
if segmented:
yield dict(advance=1)
break
except Exception:
try:
stream.close()
except Exception:
pass
if DOWNLOAD_CANCELLED.is_set() or attempts == MAX_ATTEMPTS:
if part_mode and not DOWNLOAD_CANCELLED.is_set():
raise
return
if not part_mode:
resume_offset = tmp_file.stat().st_size if tmp_file.exists() else 0
time.sleep(RETRY_WAIT)
attempts += 1
def requests(
@@ -411,7 +437,7 @@ def requests(
cookies: A mapping of Cookie Key/Values or a Cookie Jar to use for all downloads.
proxy: An optional proxy URI to route connections through for all downloads.
max_workers: The maximum amount of threads to use for downloads. Defaults to
min(12,(cpu_count+4)).
min(16, cpu_count + 4).
session: An optional requests.Session or RnetSession to use. If provided,
it will be used directly (preserving TLS fingerprinting). If None, a new
requests.Session with HTTPAdapter connection pooling will be created.
@@ -458,6 +484,8 @@ def requests(
output_dir / filename.format(i=i, ext=get_extension(url["url"] if isinstance(url, dict) else url))
]
]
# once per batch; download() skips it for segments
output_dir.mkdir(parents=True, exist_ok=True)
# Use provided session or create a new optimized requests.Session
# When a session is provided (e.g., service's RnetSession), don't mutate headers/cookies/proxy —
@@ -545,15 +573,50 @@ def requests(
pool = ThreadPoolExecutor(max_workers=max_workers)
event_queue: Queue[dict[str, Any]] = Queue()
def _download_worker(url_item: dict[str, Any]) -> None:
# hedging: first finisher claims the segment in seg_done under the lock,
# the loser discards its own output and emits nothing
seg_lock = threading.Lock()
seg_start: dict[int, float] = {}
seg_done: set[int] = set()
seg_durations: list[float] = []
hedged: set[int] = set()
def _download_worker(index: int, url_item: dict[str, Any], hedge: bool = False) -> None:
item = dict(url_item)
save_path = item.pop("save_path")
if save_path.exists(): # finished in a previous run
with seg_lock:
if index in seg_done:
return
seg_done.add(index)
event_queue.put(dict(file_downloaded=save_path, written=save_path.stat().st_size))
event_queue.put(dict(advance=1))
return
# each racer writes to its own target so a loser never touches the final file;
# the .!dev suffix keeps strays visible to the manifest parsers' cleanup
target = save_path.with_name(f"{save_path.name}.{'h' if hedge else 'p'}.!dev")
if not hedge:
seg_start[index] = time.time()
for event in download(
session=session,
segmented=segmented_batch,
**url_item,
save_path=target,
claimed=lambda: index in seg_done,
**item,
):
if "file_downloaded" in event:
with seg_lock:
if index in seg_done:
target.unlink(missing_ok=True)
return
seg_done.add(index)
if not hedge:
seg_durations.append(time.time() - seg_start[index])
os.replace(target, save_path)
event = dict(event, file_downloaded=save_path)
event_queue.put(event)
futures = [pool.submit(_download_worker, url) for url in urls]
futures = [pool.submit(_download_worker, i, url) for i, url in enumerate(urls)]
pending = set(futures)
pending_advance = 0
@@ -592,6 +655,23 @@ def requests(
yield dict(downloaded=f"{filesize.decimal(download_speed)}/s")
last_speed_report = now
# hedge stuck segments once spare workers exist (pending < max_workers)
if len(pending) < max_workers and seg_durations and not DOWNLOAD_CANCELLED.is_set():
threshold = max(HEDGE_FACTOR * statistics.median(seg_durations), HEDGE_MIN_WAIT)
with seg_lock:
stuck = [
i
for i, started in seg_start.items()
if i not in seg_done and i not in hedged and now - started > threshold
]
for i in stuck[: max_workers - len(pending)]:
hedged.add(i)
pending.add(pool.submit(_download_worker, i, urls[i], True))
# all segments claimed; superseded losers exit via their claimed() check
if len(seg_done) == len(urls):
break
# Wait efficiently for next future completion (OS condition variable)
completed, pending = wait(pending, timeout=0.1, return_when=FIRST_COMPLETED)
for future in completed:
unshackle/core/manifests/dash.py
+9 -9
View File
@@ -548,27 +548,27 @@ class DASH:
)
raise FileNotFoundError(error_msg)
is_text_subtitle = (
not drm and isinstance(track, Subtitle) and track.codec not in (Subtitle.Codec.fVTT, Subtitle.Codec.fTTML)
)
with open(save_path, "wb") as f:
if init_data:
f.write(init_data)
if len(segments_to_merge) > 1:
progress(downloaded="Merging", completed=0, total=len(segments_to_merge))
for segment_file in segments_to_merge:
segment_data = segment_file.read_bytes()
if (
not drm
and isinstance(track, Subtitle)
and track.codec not in (Subtitle.Codec.fVTT, Subtitle.Codec.fTTML)
):
segment_data = try_ensure_utf8(segment_data)
if is_text_subtitle:
segment_data = try_ensure_utf8(segment_file.read_bytes())
segment_data = (
segment_data.decode("utf8")
.replace("&lrm;", html.unescape("&lrm;"))
.replace("&rlm;", html.unescape("&rlm;"))
.encode("utf8")
)
f.write(segment_data)
f.flush()
f.write(segment_data)
else:
with open(segment_file, "rb") as src:
shutil.copyfileobj(src, f, 1024 * 1024)
segment_file.unlink()
progress(advance=1)
unshackle/core/manifests/ism.py
+20 -10
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import base64
import hashlib
import html
import shutil
import struct
import urllib.parse
from functools import partial
@@ -500,19 +501,20 @@ class ISM:
)
raise FileNotFoundError(error_msg)
is_text_subtitle = (
not session_drm
and isinstance(track, Subtitle)
and track.codec not in (Subtitle.Codec.fVTT, Subtitle.Codec.fTTML)
)
with open(save_path, "wb") as f:
first_segment = segments_to_merge[0].read_bytes() if segments_to_merge else None
init_segment = ISM._init_segment(track, session_drm, first_segment)
if init_segment:
f.write(init_segment)
for index, segment_file in enumerate(segments_to_merge):
# First segment was already read for the init synthesis — reuse it.
segment_data = first_segment if index == 0 and first_segment else segment_file.read_bytes()
if (
not session_drm
and isinstance(track, Subtitle)
and track.codec not in (Subtitle.Codec.fVTT, Subtitle.Codec.fTTML)
):
if is_text_subtitle:
# first segment was already read for the init synthesis, reuse it
segment_data = first_segment if index == 0 and first_segment else segment_file.read_bytes()
segment_data = try_ensure_utf8(segment_data)
segment_data = (
segment_data.decode("utf8")
@@ -520,8 +522,12 @@ class ISM:
.replace("&rlm;", html.unescape("&rlm;"))
.encode("utf8")
)
f.write(segment_data)
f.flush()
f.write(segment_data)
elif index == 0 and first_segment:
f.write(first_segment)
else:
with open(segment_file, "rb") as src:
shutil.copyfileobj(src, f, 1024 * 1024)
segment_file.unlink()
progress(advance=1)
@@ -535,7 +541,11 @@ class ISM:
events.emit(events.Types.TRACK_DECRYPTED, track=track, drm=session_drm, segment=None)
progress(downloaded="Decrypting", advance=100)
save_dir.rmdir()
try:
save_dir.rmdir()
except OSError:
# a superseded hedge download may still drop a .!dev file here
shutil.rmtree(save_dir, ignore_errors=True)
progress(downloaded="Downloaded")
unshackle/core/session.py
+19 -2
View File
@@ -485,7 +485,20 @@ class RnetSession:
self.log = logging.getLogger(self.__class__.__name__)
client_kwargs: dict[str, Any] = {}
for key in ("impersonate", "timeout", "proxies", "verify", "redirect"):
for key in (
"impersonate",
"timeout",
"connect_timeout",
"read_timeout",
"proxies",
"verify",
"redirect",
"pool_max_idle_per_host",
"pool_max_size",
"http1_only",
"http2_only",
"tcp_nodelay",
):
if key in session_kwargs:
client_kwargs[key] = session_kwargs.pop(key)
if "proxy" in session_kwargs:
@@ -720,11 +733,15 @@ def session(
session("Edge101", max_retries=3) # Edge 101 with custom retry
"""
if browser is None:
browser = config.curl_impersonate.get("browser", "Chrome131")
browser = config.network.get("browser", "Chrome131")
impersonate = _resolve_impersonate(browser)
session_kwargs: dict[str, Any] = {"impersonate": impersonate}
# optional rnet client knobs, see docs/NETWORK_CONFIG.md
for key in ("http1_only", "http2_only", "pool_max_idle_per_host", "pool_max_size", "tcp_nodelay"):
if key in config.network:
session_kwargs[key] = config.network[key]
session_kwargs.update(kwargs)
session_obj = RnetSession(**session_kwargs)