mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-12 20:15:10 +00:00
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:
@@ -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. |
|
||||||
|
|||||||
@@ -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. |
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -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
|
||||||
@@ -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"
|
||||||
|
|||||||
@@ -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)}")
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
Reference in New Issue
Block a user