feat(dl): add --speed-limit and serve global_speed_limit

Shared token bucket across all download threads; decimal suffixes (5M = 5.0 MB/s) or plain bytes/sec
This commit is contained in:
imSp4rky
2026-07-31 12:48:46 -06:00
parent e9a40a3af3
commit 734df87bad
9 changed files with 313 additions and 10 deletions
+1
View File
@@ -172,6 +172,7 @@ Keep only certain track types, or skip certain track types. Attachments are alwa
|---|---|---| |---|---|---|
| `--workers` | downloader default | Per-track download threads. | | `--workers` | downloader default | Per-track download threads. |
| `--downloads` | `1` | Number of tracks downloaded concurrently. | | `--downloads` | `1` | Number of tracks downloaded concurrently. |
| `--speed-limit` | unlimited | Cap total download speed across all threads and tracks, e.g. `500k`, `5M`, `1.5G` or plain bytes/sec. Values are bytes, not bits (`5M` = 5.0 MB/s). `off` disables a configured limit. |
| `--no-cache` | off | Bypass the title cache. | | `--no-cache` | off | Bypass the title cache. |
| `--reset-cache` | off | Clear the title cache. | | `--reset-cache` | off | Clear the title cache. |
| `--list` | off | List available/would-be-downloaded tracks; do not download. | | `--list` | off | List available/would-be-downloaded tracks; do not download. |
+2
View File
@@ -245,6 +245,7 @@ is the [REST API](../dev/rest-api/index.md) section; these are the config keys.
| `max_sessions` | int | `100` | Maximum concurrent sessions. | | `max_sessions` | int | `100` | Maximum concurrent sessions. |
| `history_limit` | int | `100` | How many finished jobs to retain in history. | | `history_limit` | int | `100` | How many finished jobs to retain in history. |
| `compression_level` | int | `1` | gzip level for responses. | | `compression_level` | int | `1` | gzip level for responses. |
| `global_speed_limit` | str | *(unlimited)* | Server-wide download speed cap, e.g. `10M`, `1.5G` or plain bytes/sec (same format as `speed_limit`). One shared budget across all concurrent jobs; per-job speed limits are ignored while it is set. |
| `cdm_overrides` | dict | *(unset)* | Allowed per-request CDM overrides. | | `cdm_overrides` | dict | *(unset)* | Allowed per-request CDM overrides. |
| `allow_job_credentials` | bool | `false` | Permit clients to supply credentials per job. | | `allow_job_credentials` | bool | `false` | Permit clients to supply credentials per job. |
| `devices` | list | *(auto)* | Widevine devices offered; auto-filled from `directories.wvds`. | | `devices` | list | *(auto)* | Widevine devices offered; auto-filled from `directories.wvds`. |
@@ -322,6 +323,7 @@ Common keys (a useful subset; every `dl` flag works):
| `no_subs` / `no_audio` / `no_chapters` | bool | `false` | Skip that track type. | | `no_subs` / `no_audio` / `no_chapters` | bool | `false` | Skip that track type. |
| `downloads` | int | `1` | Tracks downloaded concurrently. | | `downloads` | int | `1` | Tracks downloaded concurrently. |
| `workers` | int | *(downloader default)* | Threads per track. | | `workers` | int | *(downloader default)* | Threads per track. |
| `speed_limit` | str | *(unlimited)* | Total download speed cap across all downloads combined, e.g. `500k`, `5M`, `1.5G` or plain bytes/sec. Values are bytes, not bits (`5M` = 5.0 MB/s). |
| `slow` | str | *(unset)* | Inter-title delay, e.g. `"20-40"`. | | `slow` | str | *(unset)* | Inter-title delay, e.g. `"20-40"`. |
| `best_available` | bool | `false` | Fall back to best quality if the request is unavailable. | | `best_available` | bool | `false` | Fall back to best quality if the request is unavailable. |
| `proxy` | str | *(unset)* | Default proxy URI or 2-letter country. | | `proxy` | str | *(unset)* | Default proxy URI or 2-letter country. |
View File
+118
View File
@@ -0,0 +1,118 @@
import importlib
import inspect
import threading
import time
import pytest
from unshackle.core.constants import DOWNLOAD_CANCELLED
from unshackle.core.downloaders.requests import TokenBucket, format_speed, parse_speed_limit, set_speed_limit
requests_downloader = importlib.import_module("unshackle.core.downloaders.requests")
@pytest.mark.parametrize(
("value", "expected"),
[
("500k", 500_000),
("5M", 5_000_000),
("10MB", 10_000_000),
("10mb/s", 10_000_000),
("1.5G", 1_500_000_000),
("1048576", 1_048_576),
(5_000_000, 5_000_000.0),
(None, None),
("off", None),
("none", None),
("unlimited", None),
("0", None),
(0, None),
],
)
def test_parse_speed_limit(value, expected):
assert parse_speed_limit(value) == expected
@pytest.mark.parametrize(
"value",
["fast", "10TB", "10MiB", "-5M", "M", "1..5M", ".", "1.2.3", float("inf"), float("nan"), -5, -1.5],
)
def test_parse_speed_limit_rejects(value):
with pytest.raises(ValueError):
parse_speed_limit(value)
@pytest.mark.parametrize(
("rate", "expected"),
[
(500, "500 bytes/s"),
(512_000, "512.0 kB/s"),
(39 * 1024**2, "40.9 MB/s"),
(1.5 * 1024**3, "1.6 GB/s"),
],
)
def test_format_speed(rate, expected):
assert format_speed(rate) == expected
def test_result_accepts_a_missing_speed_limit():
"""The API/serve path calls dl.result() by keyword without speed_limit."""
from unshackle.commands.dl import dl
assert inspect.signature(dl.result).parameters["speed_limit"].default is None
def test_bucket_holds_aggregate_rate_across_threads():
rate = 10_000_000
bucket = TokenBucket(rate)
chunk = 100_000
totals = [0, 0, 0, 0]
start = time.monotonic()
stop = start + 1.0
def worker(i: int) -> None:
while time.monotonic() < stop:
bucket.consume(chunk)
totals[i] += chunk
threads = [threading.Thread(target=worker, args=(i,)) for i in range(4)]
for t in threads:
t.start()
for t in threads:
t.join()
achieved = sum(totals) / (time.monotonic() - start)
# exceeding the limit is a real failure; undershoot is scheduler noise on a loaded box
assert achieved <= rate * 1.1
assert achieved >= rate * 0.5
def test_bucket_chunk_larger_than_capacity_does_not_deadlock():
bucket = TokenBucket(1_000_000)
start = time.monotonic()
bucket.consume(1_100_000)
assert 1.0 < time.monotonic() - start < 2.0
def test_locked_limit_ignores_unlocked_calls():
"""serve's global_speed_limit must survive the per-job set in dl.result()."""
try:
set_speed_limit(5_000_000, lock=True)
set_speed_limit(None)
set_speed_limit(1_000_000)
limiter = requests_downloader._speed_limiter
assert limiter is not None and limiter.rate == 5_000_000
finally:
requests_downloader._speed_limit_locked = False
requests_downloader._speed_limiter = None
def test_bucket_wait_aborts_on_cancellation():
"""Debt sleeps grow with thread count; cancelling must not wait them out."""
bucket = TokenBucket(1_000)
DOWNLOAD_CANCELLED.set()
try:
start = time.monotonic()
bucket.consume(1_000_000)
assert time.monotonic() - start < 1.0
finally:
DOWNLOAD_CANCELLED.clear()
+43
View File
@@ -0,0 +1,43 @@
import pytest
from unshackle.core.downloaders.requests import SPEED_ROLLING_WINDOW, SpeedWindow
def test_rate_reflects_recent_window_not_full_history():
"""After a slow start (shared bucket), the display must show the current rate."""
win = SpeedWindow(0.0)
total = 0
for t in range(0, 60):
total += 5_000_000
win.rate(float(t), total)
for t in range(60, 120):
total += 10_000_000
rate = win.rate(float(t), total)
assert rate == pytest.approx(10_000_000, rel=0.01)
def test_first_report_is_the_since_start_average():
"""Bytes are credited in lumps; the first lump must not read as a whole window's worth."""
win = SpeedWindow(0.0)
assert win.rate(2.5, 25_000_000) == pytest.approx(10_000_000)
def test_zero_span_returns_none():
win = SpeedWindow(0.0)
assert win.rate(0.0, 1_000_000) is None
def test_stall_returns_none():
"""Once the seed ages out, a total that stops growing has no delta to report."""
win = SpeedWindow(0.0)
for t in range(0, SPEED_ROLLING_WINDOW + 6):
rate = win.rate(float(t), 1_000_000)
assert rate is None
def test_window_drops_old_samples():
win = SpeedWindow(0.0)
for t in range(0, 3 * SPEED_ROLLING_WINDOW):
win.rate(float(t), t * 1_000_000)
assert win.samples[0][0] >= 2 * SPEED_ROLLING_WINDOW - 1
assert len(win.samples) == SPEED_ROLLING_WINDOW + 1
+21
View File
@@ -45,6 +45,7 @@ from unshackle.core.config import config, resolve_cdm_name, resolve_decryption
from unshackle.core.console import GradientPulseBarColumn, SyncLive, console from unshackle.core.console import GradientPulseBarColumn, SyncLive, console
from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack, context_settings from unshackle.core.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack, context_settings
from unshackle.core.credential import Credential from unshackle.core.credential import Credential
from unshackle.core.downloaders import format_speed, parse_speed_limit, set_speed_limit
from unshackle.core.drm import DRM_T, ClearKeyCENC, MonaLisa, PlayReady, Widevine from unshackle.core.drm import DRM_T, ClearKeyCENC, MonaLisa, PlayReady, Widevine
from unshackle.core.events import events from unshackle.core.events import events
from unshackle.core.music import ( from unshackle.core.music import (
@@ -691,6 +692,16 @@ class dl:
help="Max workers/threads to download with per-track. Default depends on the downloader.", help="Max workers/threads to download with per-track. Default depends on the downloader.",
) )
@click.option("--downloads", type=int, default=1, help="Amount of tracks to download concurrently.") @click.option("--downloads", type=int, default=1, help="Amount of tracks to download concurrently.")
@click.option(
"--speed-limit",
"speed_limit",
type=str,
default=None,
metavar="SPEED",
help="Cap total download speed across all downloads combined, e.g. 500k, 5M, 1.5G or plain "
"bytes/sec (5M = 5.0 MB/s; bytes, not bits). "
"'off' disables a limit configured in the dl: config block.",
)
@click.option( @click.option(
"-o", "-o",
"--output", "--output",
@@ -1258,6 +1269,7 @@ class dl:
downloads: int, downloads: int,
worst: bool, worst: bool,
best_available: bool, best_available: bool,
speed_limit: Optional[str] = None,
split_audio: Optional[bool] = None, split_audio: Optional[bool] = None,
merge_video: Optional[bool] = None, merge_video: Optional[bool] = None,
real_video_bitrate: bool = False, real_video_bitrate: bool = False,
@@ -1275,6 +1287,15 @@ class dl:
if skip_dl: if skip_dl:
DOWNLOAD_LICENCE_ONLY.set() DOWNLOAD_LICENCE_ONLY.set()
try:
speed_limit_bps = parse_speed_limit(speed_limit)
except ValueError as e:
self.log.error(str(e))
sys.exit(1)
set_speed_limit(speed_limit_bps)
if speed_limit_bps:
self.log.info(f"Speed limit: {format_speed(speed_limit_bps)}")
if export: if export:
config.directories.exports.mkdir(parents=True, exist_ok=True) config.directories.exports.mkdir(parents=True, exist_ok=True)
export_path = config.directories.exports / f"export_{self.service}_{int(time.time())}.json" export_path = config.directories.exports / f"export_{self.service}_{int(time.time())}.json"
+9
View File
@@ -10,6 +10,7 @@ from unshackle.core.api import cors_middleware, setup_routes, setup_swagger
from unshackle.core.api.compression import compression_middleware from unshackle.core.api.compression import compression_middleware
from unshackle.core.config import config from unshackle.core.config import config
from unshackle.core.constants import context_settings from unshackle.core.constants import context_settings
from unshackle.core.downloaders import format_speed, parse_speed_limit, set_speed_limit
@click.command( @click.command(
@@ -146,6 +147,14 @@ def serve(
if remote_only: if remote_only:
api_only = True api_only = True
try:
global_speed_limit = parse_speed_limit(config.serve.get("global_speed_limit"))
except ValueError as e:
raise click.ClickException(f"serve.global_speed_limit: {e}")
if global_speed_limit:
set_speed_limit(global_speed_limit, lock=True)
log.info(f"Global speed limit: {format_speed(global_speed_limit)} (shared by all jobs)")
global_services = config.serve.get("services") global_services = config.serve.get("services")
if global_services: if global_services:
log.info(f"Global service allowlist: {', '.join(global_services)}") log.info(f"Global service allowlist: {', '.join(global_services)}")
+2 -2
View File
@@ -1,3 +1,3 @@
from .requests import requests from .requests import format_speed, parse_speed_limit, requests, set_speed_limit
__all__ = ("requests",) __all__ = ("requests", "format_speed", "parse_speed_limit", "set_speed_limit")
+117 -8
View File
@@ -1,8 +1,10 @@
import math import math
import os import os
import re
import statistics import statistics
import threading import threading
import time import time
from collections import deque
from concurrent.futures import FIRST_COMPLETED, wait from concurrent.futures import FIRST_COMPLETED, wait
from concurrent.futures.thread import ThreadPoolExecutor from concurrent.futures.thread import ThreadPoolExecutor
from http.cookiejar import CookieJar from http.cookiejar import CookieJar
@@ -42,6 +44,104 @@ RANGE_PARALLEL_MIN_SIZE = 64 * 1024 * 1024
RANGE_PARALLEL_PART_SIZE = 16 * 1024 * 1024 RANGE_PARALLEL_PART_SIZE = 16 * 1024 * 1024
class SpeedWindow:
"""Rolling-window download rate over the last SPEED_ROLLING_WINDOW seconds."""
def __init__(self, start: float) -> None:
# zero-byte seed keeps early readings from exceeding the true rate
self.samples: deque[tuple[float, int]] = deque([(start, 0)])
def rate(self, now: float, total: int) -> Optional[float]:
self.samples.append((now, total))
while now - self.samples[0][0] > SPEED_ROLLING_WINDOW:
self.samples.popleft()
span = now - self.samples[0][0]
delta = total - self.samples[0][1]
if span <= 0 or delta <= 0:
return None
return delta / span
class TokenBucket:
"""Aggregate byte-rate limiter shared by every download thread."""
def __init__(self, rate: float) -> None:
self.rate = rate
self.capacity = rate
self.tokens = 0.0
self.ts = time.monotonic()
self.lock = threading.Lock()
def consume(self, n: int) -> None:
# tokens may go negative; waiting for n > capacity would deadlock
with self.lock:
now = time.monotonic()
self.tokens = min(self.capacity, self.tokens + (now - self.ts) * self.rate)
self.ts = now
self.tokens -= n
wait = -self.tokens / self.rate if self.tokens < 0 else 0.0
if wait > 0:
DOWNLOAD_CANCELLED.wait(wait)
_speed_limiter: Optional[TokenBucket] = None
_speed_limit_locked = False
SPEED_UNITS = {
"": 1,
"b": 1,
"k": 1_000,
"kb": 1_000,
"m": 1_000_000,
"mb": 1_000_000,
"g": 1_000_000_000,
"gb": 1_000_000_000,
}
def parse_speed_limit(value: Union[str, int, float, None]) -> Optional[float]:
"""
Parse a human-readable speed limit into bytes/sec, or None for unlimited.
Accepts plain numbers (exact bytes/sec) or a decimal suffix with optional
/s, case-insensitive: "500k", "5M", "10MB/s", "1.5G", 5000000. Suffixes
match the displayed speeds (5M = 5.0 MB/s); values are bytes, not bits.
"", 0, "off", "none" and "unlimited" mean no limit.
"""
if value is None:
return None
if isinstance(value, (int, float)):
if not math.isfinite(value) or value < 0:
raise ValueError(f"Invalid speed limit {value!r}: must be a positive number of bytes/sec.")
return float(value) or None
text = value.strip().lower().removesuffix("/s")
if text in ("", "0", "off", "none", "unlimited"):
return None
match = re.fullmatch(r"(\d*\.?\d+)\s*([kmg]b?|b)?", text)
if not match:
raise ValueError(f"Invalid speed limit {value!r}: use e.g. 500k, 5M, 1.5G, plain bytes/sec, or 'off'.")
return float(match.group(1)) * SPEED_UNITS[match.group(2) or ""]
def format_speed(bytes_per_sec: float) -> str:
"""Format a byte rate in the same decimal units the progress bars show."""
return f"{filesize.decimal(int(bytes_per_sec))}/s"
def set_speed_limit(bytes_per_sec: Optional[float], lock: bool = False) -> None:
"""
Set (or clear) the global download speed limit shared by all threads.
A locked limit (serve's global_speed_limit) wins: later unlocked calls,
like the per-job one in dl.result(), become no-ops for the process.
"""
global _speed_limiter, _speed_limit_locked
if _speed_limit_locked and not lock:
return
_speed_limit_locked = lock
_speed_limiter = TokenBucket(bytes_per_sec) if bytes_per_sec else None
def _adaptive_chunk_size(content_length: int) -> int: def _adaptive_chunk_size(content_length: int) -> int:
"""Pick chunk size based on content length. Benchmarked sweet spot: 512KB-4MB.""" """Pick chunk size based on content length. Benchmarked sweet spot: 512KB-4MB."""
if content_length <= 0: if content_length <= 0:
@@ -132,7 +232,8 @@ def _dispatch_parts(
yield {"total": total_size} yield {"total": total_size}
total_bytes = 0 total_bytes = 0
start_time = last_report = time.time() last_report = time.time()
speed_window = SpeedWindow(last_report)
completed = False completed = False
worker_error = False worker_error = False
@@ -153,7 +254,9 @@ def _dispatch_parts(
now = time.time() now = time.time()
if now - last_report > 0.5 and total_bytes > 0: if now - last_report > 0.5 and total_bytes > 0:
yield {"downloaded": f"{filesize.decimal(math.ceil(total_bytes / (now - start_time)))}/s"} rate = speed_window.rate(now, total_bytes)
if rate:
yield {"downloaded": f"{filesize.decimal(math.ceil(rate))}/s"}
last_report = now last_report = now
done, pending = wait(pending, timeout=0.1, return_when=FIRST_COMPLETED) done, pending = wait(pending, timeout=0.1, return_when=FIRST_COMPLETED)
@@ -300,7 +403,10 @@ def download(
except ValueError: except ValueError:
content_length = 0 content_length = 0
limiter = _speed_limiter
chunk_size = _adaptive_chunk_size(content_length) chunk_size = _adaptive_chunk_size(content_length)
if limiter:
chunk_size = min(chunk_size, max(8192, int(limiter.rate / 4)))
total_size = (resume_offset + content_length) if resumed and content_length > 0 else 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 not segmented and not part_mode:
@@ -350,6 +456,8 @@ def download(
except Exception: except Exception:
pass pass
return return
if limiter is not None:
limiter.consume(len(chunk))
_write(chunk) _write(chunk)
download_size = len(chunk) download_size = len(chunk)
written += download_size written += download_size
@@ -543,7 +651,7 @@ def requests(
url_item = urls[0] url_item = urls[0]
try: try:
ranged_used = False ranged_used = False
if max_workers > 1: if max_workers > 1 and _speed_limiter is None:
total_size, supports_ranges = _probe_ranged(url_item["url"], session) total_size, supports_ranges = _probe_ranged(url_item["url"], session)
if supports_ranges and total_size >= RANGE_PARALLEL_MIN_SIZE: if supports_ranges and total_size >= RANGE_PARALLEL_MIN_SIZE:
try: try:
@@ -581,6 +689,7 @@ def requests(
total_bytes = 0 total_bytes = 0
start_time = time.time() start_time = time.time()
last_speed_report = start_time last_speed_report = start_time
speed_window = SpeedWindow(start_time)
last_hedge_check = 0.0 last_hedge_check = 0.0
hedge_median = 0.0 # cached; recomputed only when seg_durations grows hedge_median = 0.0 # cached; recomputed only when seg_durations grows
hedge_median_len = 0 hedge_median_len = 0
@@ -664,10 +773,9 @@ def requests(
# Yield speed every 0.5s (throttled to avoid spamming Rich) # Yield speed every 0.5s (throttled to avoid spamming Rich)
now = time.time() now = time.time()
if now - last_speed_report > 0.5 and total_bytes > 0: if now - last_speed_report > 0.5 and total_bytes > 0:
elapsed = now - start_time rate = speed_window.rate(now, total_bytes)
if elapsed > 0: if rate:
download_speed = math.ceil(total_bytes / elapsed) yield dict(downloaded=f"{filesize.decimal(math.ceil(rate))}/s")
yield dict(downloaded=f"{filesize.decimal(download_speed)}/s")
last_speed_report = now last_speed_report = now
# hedge stuck segments once spare workers exist (pending < max_workers); # hedge stuck segments once spare workers exist (pending < max_workers);
@@ -675,6 +783,7 @@ def requests(
if ( if (
len(pending) < max_workers len(pending) < max_workers
and seg_durations and seg_durations
and _speed_limiter is None
and now - last_hedge_check > 0.5 and now - last_hedge_check > 0.5
and not DOWNLOAD_CANCELLED.is_set() and not DOWNLOAD_CANCELLED.is_set()
): ):
@@ -768,4 +877,4 @@ def requests(
) )
__all__ = ("requests",) __all__ = ("requests", "format_speed", "parse_speed_limit", "set_speed_limit")