From 734df87badd46169e4163e14afb5ae1ab4fabe22 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Fri, 31 Jul 2026 12:48:46 -0600 Subject: [PATCH] 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 --- docs/guide/cli-reference.md | 1 + docs/reference/configuration.md | 2 + tests/downloaders/__init__.py | 0 tests/downloaders/test_speed_limit.py | 118 +++++++++++++++++++++++ tests/downloaders/test_speed_window.py | 43 +++++++++ unshackle/commands/dl.py | 21 +++++ unshackle/commands/serve.py | 9 ++ unshackle/core/downloaders/__init__.py | 4 +- unshackle/core/downloaders/requests.py | 125 +++++++++++++++++++++++-- 9 files changed, 313 insertions(+), 10 deletions(-) create mode 100644 tests/downloaders/__init__.py create mode 100644 tests/downloaders/test_speed_limit.py create mode 100644 tests/downloaders/test_speed_window.py diff --git a/docs/guide/cli-reference.md b/docs/guide/cli-reference.md index c19c9424..ce4d53dd 100644 --- a/docs/guide/cli-reference.md +++ b/docs/guide/cli-reference.md @@ -172,6 +172,7 @@ Keep only certain track types, or skip certain track types. Attachments are alwa |---|---|---| | `--workers` | downloader default | Per-track download threads. | | `--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. | | `--reset-cache` | off | Clear the title cache. | | `--list` | off | List available/would-be-downloaded tracks; do not download. | diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 22ab66da..45a00d26 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -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. | | `history_limit` | int | `100` | How many finished jobs to retain in history. | | `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. | | `allow_job_credentials` | bool | `false` | Permit clients to supply credentials per job. | | `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. | | `downloads` | int | `1` | Tracks downloaded concurrently. | | `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"`. | | `best_available` | bool | `false` | Fall back to best quality if the request is unavailable. | | `proxy` | str | *(unset)* | Default proxy URI or 2-letter country. | diff --git a/tests/downloaders/__init__.py b/tests/downloaders/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/downloaders/test_speed_limit.py b/tests/downloaders/test_speed_limit.py new file mode 100644 index 00000000..538447ca --- /dev/null +++ b/tests/downloaders/test_speed_limit.py @@ -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() diff --git a/tests/downloaders/test_speed_window.py b/tests/downloaders/test_speed_window.py new file mode 100644 index 00000000..46d555b2 --- /dev/null +++ b/tests/downloaders/test_speed_window.py @@ -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 diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 3695160f..ed25dd1b 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -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.constants import DOWNLOAD_CANCELLED, DOWNLOAD_LICENCE_ONLY, AnyTrack, context_settings 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.events import events 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.", ) @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( "-o", "--output", @@ -1258,6 +1269,7 @@ class dl: downloads: int, worst: bool, best_available: bool, + speed_limit: Optional[str] = None, split_audio: Optional[bool] = None, merge_video: Optional[bool] = None, real_video_bitrate: bool = False, @@ -1275,6 +1287,15 @@ class dl: if skip_dl: 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: config.directories.exports.mkdir(parents=True, exist_ok=True) export_path = config.directories.exports / f"export_{self.service}_{int(time.time())}.json" diff --git a/unshackle/commands/serve.py b/unshackle/commands/serve.py index 22ad10e3..65ccbde6 100644 --- a/unshackle/commands/serve.py +++ b/unshackle/commands/serve.py @@ -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.config import config from unshackle.core.constants import context_settings +from unshackle.core.downloaders import format_speed, parse_speed_limit, set_speed_limit @click.command( @@ -146,6 +147,14 @@ def serve( if remote_only: 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") if global_services: log.info(f"Global service allowlist: {', '.join(global_services)}") diff --git a/unshackle/core/downloaders/__init__.py b/unshackle/core/downloaders/__init__.py index 66d3b420..bc68b6fa 100644 --- a/unshackle/core/downloaders/__init__.py +++ b/unshackle/core/downloaders/__init__.py @@ -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") diff --git a/unshackle/core/downloaders/requests.py b/unshackle/core/downloaders/requests.py index 8a13b296..97aa3411 100644 --- a/unshackle/core/downloaders/requests.py +++ b/unshackle/core/downloaders/requests.py @@ -1,8 +1,10 @@ import math import os +import re import statistics import threading import time +from collections import deque from concurrent.futures import FIRST_COMPLETED, wait from concurrent.futures.thread import ThreadPoolExecutor from http.cookiejar import CookieJar @@ -42,6 +44,104 @@ RANGE_PARALLEL_MIN_SIZE = 64 * 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: """Pick chunk size based on content length. Benchmarked sweet spot: 512KB-4MB.""" if content_length <= 0: @@ -132,7 +232,8 @@ def _dispatch_parts( yield {"total": total_size} total_bytes = 0 - start_time = last_report = time.time() + last_report = time.time() + speed_window = SpeedWindow(last_report) completed = False worker_error = False @@ -153,7 +254,9 @@ def _dispatch_parts( now = time.time() 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 done, pending = wait(pending, timeout=0.1, return_when=FIRST_COMPLETED) @@ -300,7 +403,10 @@ def download( except ValueError: content_length = 0 + limiter = _speed_limiter 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 if not segmented and not part_mode: @@ -350,6 +456,8 @@ def download( except Exception: pass return + if limiter is not None: + limiter.consume(len(chunk)) _write(chunk) download_size = len(chunk) written += download_size @@ -543,7 +651,7 @@ def requests( url_item = urls[0] try: 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) if supports_ranges and total_size >= RANGE_PARALLEL_MIN_SIZE: try: @@ -581,6 +689,7 @@ def requests( total_bytes = 0 start_time = time.time() last_speed_report = start_time + speed_window = SpeedWindow(start_time) last_hedge_check = 0.0 hedge_median = 0.0 # cached; recomputed only when seg_durations grows hedge_median_len = 0 @@ -664,10 +773,9 @@ def requests( # Yield speed every 0.5s (throttled to avoid spamming Rich) now = time.time() if now - last_speed_report > 0.5 and total_bytes > 0: - elapsed = now - start_time - if elapsed > 0: - download_speed = math.ceil(total_bytes / elapsed) - yield dict(downloaded=f"{filesize.decimal(download_speed)}/s") + rate = speed_window.rate(now, total_bytes) + if rate: + yield dict(downloaded=f"{filesize.decimal(math.ceil(rate))}/s") last_speed_report = now # hedge stuck segments once spare workers exist (pending < max_workers); @@ -675,6 +783,7 @@ def requests( if ( len(pending) < max_workers and seg_durations + and _speed_limiter is None and now - last_hedge_check > 0.5 and not DOWNLOAD_CANCELLED.is_set() ): @@ -768,4 +877,4 @@ def requests( ) -__all__ = ("requests",) +__all__ = ("requests", "format_speed", "parse_speed_limit", "set_speed_limit")