From 537a1dc97387ad4c5f08bfa10d995a6b79c3e166 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Thu, 20 Aug 2026 11:24:12 -0600 Subject: [PATCH] perf(dl): parallelise mux, licensing, and vault lookups; size the shared pool Grow the service session pool to downloads * workers before track downloads, run mkvmerge jobs concurrently (muxing.concurrency, default 4), lock DRM preparation per key set instead of globally, and query local vaults first then remote vaults together. --- docs/reference/configuration/download.md | 2 + tests/core/test_vaults_lookup.py | 62 ++++++++++++++++++++ unshackle/commands/dl.py | 54 +++++++++++++---- unshackle/core/downloaders/__init__.py | 4 +- unshackle/core/downloaders/requests.py | 7 ++- unshackle/core/service.py | 25 ++++++++ unshackle/core/tracks/chapters.py | 7 ++- unshackle/core/vault.py | 3 + unshackle/core/vaults.py | 74 ++++++++++++++++-------- unshackle/vaults/SQLite.py | 2 + 10 files changed, 199 insertions(+), 41 deletions(-) create mode 100644 tests/core/test_vaults_lookup.py diff --git a/docs/reference/configuration/download.md b/docs/reference/configuration/download.md index 6efe68f3..de5f599d 100644 --- a/docs/reference/configuration/download.md +++ b/docs/reference/configuration/download.md @@ -232,10 +232,12 @@ Matroska (MKV) muxing options. | `default_language` | dict | `{}` | Force which language is flagged *default* per track type, e.g. `{audio: en, subtitle: en}`. | | `merge_audio` | bool | `true` | Merge all audio into one file. `--split-audio` on the CLI flips this off. | | `merge_video` | bool | `false` | Merge video tracks that share height, range, and codec into one file, so only language varies inside a file. `--merge-video` on the CLI flips this on. | +| `concurrency` | int | `4` | How many mkvmerge processes run at the same time when a title muxes more than one output file (one per video variant or audio codec). Each process uses one CPU core, so this shortens the mux step on an SSD. Set it to `1` on a hard disk, where parallel writes are slower. | ```yaml muxing: set_title: true + concurrency: 4 default_language: audio: en subtitle: en diff --git a/tests/core/test_vaults_lookup.py b/tests/core/test_vaults_lookup.py new file mode 100644 index 00000000..f7f441d3 --- /dev/null +++ b/tests/core/test_vaults_lookup.py @@ -0,0 +1,62 @@ +import threading +import time +from typing import Iterator, Optional, Union +from uuid import UUID + +from unshackle.core.vault import Vault +from unshackle.core.vaults import Vaults + + +class FakeVault(Vault): + def __init__(self, name: str, key: Optional[str], delay: float = 0.0, local: bool = False): + super().__init__(name) + self.key, self.delay, self.local = key, delay, local + self.calls = 0 + self.threads: set[int] = set() + + def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]: + self.calls += 1 + self.threads.add(threading.get_ident()) + time.sleep(self.delay) + return self.key + + def get_keys(self, service: str) -> Iterator[tuple[str, str]]: + return iter(()) + + def add_key(self, service: str, kid: Union[UUID, str], key: str) -> bool: + return False + + def add_keys(self, service: str, kid_keys: dict) -> int: + return 0 + + def get_services(self) -> Iterator[str]: + return iter(()) + + +def make(*vaults: FakeVault) -> Vaults: + v = Vaults("SVC") + v.vaults = list(vaults) + return v + + +def test_local_hit_skips_remote() -> None: + local = FakeVault("local", "aa" * 16, local=True) + remote = FakeVault("remote", "bb" * 16) + assert make(remote, local).get_key("kid") == ("aa" * 16, local) + assert remote.calls == 0 + + +def test_remote_vaults_run_together_and_first_hit_wins() -> None: + local = FakeVault("local", None, local=True) + slow = FakeVault("slow", "cc" * 16, delay=1.0) + fast = FakeVault("fast", "dd" * 16, delay=0.05) + start = time.monotonic() + key, vault = make(local, slow, fast).get_key("kid") + assert (key, vault) == ("dd" * 16, fast) + assert time.monotonic() - start < 0.9 # did not wait for the slow vault + assert slow.calls == 1 and local.calls == 1 + assert slow.threads != fast.threads + + +def test_miss_everywhere() -> None: + assert make(FakeVault("a", None), FakeVault("b", "0" * 32)).get_key("kid") == (None, None) diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 0e3afea5..4d92267e 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -46,13 +46,13 @@ from unshackle.core.config import config, resolve_cdm_name, resolve_decryption from unshackle.core.console import GradientPulseBarColumn, SyncLive, console, listing_panel 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.downloaders import default_max_workers, 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.providers.anilist import parse_anilist_ref from unshackle.core.providers.tvdb import SEASON_TYPES, parse_int from unshackle.core.proxies import Basic, ExpressVPN, Gluetun, Hola, NordVPN, ProtonVPN, SurfsharkVPN, WindscribeVPN -from unshackle.core.service import Service +from unshackle.core.service import Service, grow_session_pool from unshackle.core.services import Services from unshackle.core.temp import with_task_temp from unshackle.core.title_cacher import get_account_hash @@ -907,7 +907,8 @@ class dl: def cli(ctx: click.Context, **kwargs: Any) -> dl: return dl(ctx, **kwargs) - DRM_TABLE_LOCK = Lock() + DRM_LOCKS: dict[str, Lock] = {} + DRM_LOCKS_GUARD = Lock() VAULT_WRITE_LOCK = Lock() EXPORT_LOCK = Lock() LICENSE_KEY_CACHE: dict[UUID, str] = {} @@ -2846,6 +2847,7 @@ class dl: self.log.debug(f"Skipped subtitle {track.id} was already absent from the track list.") skipped_before = len(self.skipped_subtitles) + grow_session_pool(service.session, downloads * (workers or default_max_workers())) download_tracks_in_passes( title.tracks, downloads, @@ -3266,10 +3268,13 @@ class dl: try: with SyncLive(Padding(progress, (0, 5, 1, 5)), console=console, refresh_per_second=20): mux_failed = False - for mux_index, (task_id, task_tracks, audio_codec) in enumerate(multiplex_tasks, start=1): + audio_expected = not video_only and not no_audio + + def mux_one( + mux_index: int, task_id: TaskID, task_tracks: Tracks + ) -> tuple[Path, int, list[str]]: progress.start_task(task_id) - audio_expected = not video_only and not no_audio - muxed_path, return_code, errors = task_tracks.mux( + return task_tracks.mux( str(title), progress=partial(progress.update, task_id=task_id), delete=False, @@ -3278,6 +3283,19 @@ class dl: skip_subtitles=skip_subtitle_mux, output_path=unique_mux_output(task_tracks, mux_index), ) + + with ThreadPoolExecutor( + max(1, min(len(multiplex_tasks), int(config.muxing.get("concurrency", 4)))) + ) as mux_pool: + mux_results = list( + mux_pool.map( + lambda args: mux_one(*args), + [(i, t, tt) for i, (t, tt, _) in enumerate(multiplex_tasks, start=1)], + ) + ) + for (task_id, task_tracks, audio_codec), (muxed_path, return_code, errors) in zip( + multiplex_tasks, mux_results + ): muxed_paths.append(muxed_path) muxed_audio_codecs[muxed_path] = audio_codec if return_code >= 2: @@ -3291,7 +3309,6 @@ class dl: self.log.warning(line) if return_code >= 2: mux_failed = True - break if mux_failed: sys.exit(1) @@ -3827,6 +3844,19 @@ class dl: f"{successful_caches}/{len(self.vaults)} Vaults" ) + @classmethod + def drm_lock(cls, drm: DRM_T) -> Lock: + """Return the lock for one DRM's set of content keys. + + Tracks that share KIDs still take turns, so the second one finds the content keys in + LICENSE_KEY_CACHE and skips the challenge. Tracks with different KIDs send a challenge + and query the key vaults at the same time. + """ + kids = getattr(drm, "kids", None) or [] + key = ",".join(sorted(getattr(k, "hex", str(k)) for k in kids)) or getattr(drm, "content_id", None) + with cls.DRM_LOCKS_GUARD: + return cls.DRM_LOCKS.setdefault(key or type(drm).__name__, Lock()) + def prepare_drm( self, drm: DRM_T, @@ -3871,7 +3901,7 @@ class dl: keys=[{"kid": k.hex, "key": v} for k, v in drm.content_keys.items()], remote=True, ) - with self.DRM_TABLE_LOCK: + with self.drm_lock(drm): pssh_str = "" expected_class = "PlayReady" if server_drm_type == "playready" else "Widevine" matching_drm = next( @@ -3949,7 +3979,7 @@ class dl: ) pending_vault_writes: list[Callable[[], Any]] = [] - with self.DRM_TABLE_LOCK: + with self.drm_lock(drm): pssh_display = self.truncate_pssh_for_display(drm.pssh.dumps(), "Widevine") cek_tree = Tree(Text.assemble(("Widevine", "cyan"), (f"({pssh_display})", "text"), overflow="fold")) pre_existing_tree = next( @@ -4144,7 +4174,7 @@ class dl: ) pending_vault_writes = [] - with self.DRM_TABLE_LOCK: + with self.drm_lock(drm): pssh_display = self.truncate_pssh_for_display(drm.pssh_b64 or "", "PlayReady") cek_tree = Tree( Text.assemble( @@ -4303,7 +4333,7 @@ class dl: elif isinstance(drm, ClearKeyCENC): pending_vault_writes = [] - with self.DRM_TABLE_LOCK: + with self.drm_lock(drm): cek_tree = Tree(Text.assemble(("ClearKey", "cyan"), overflow="fold")) pre_existing_tree = next( (x for x in table.columns[0].cells if isinstance(x, Tree) and x.label == cek_tree.label), None @@ -4442,7 +4472,7 @@ class dl: self.flush_vault_writes(pending_vault_writes) elif isinstance(drm, MonaLisa): - with self.DRM_TABLE_LOCK: + with self.drm_lock(drm): display_id = drm.content_id or drm.pssh pssh_display = self.truncate_pssh_for_display(display_id, "MonaLisa") cek_tree = Tree(Text.assemble(("MonaLisa", "cyan"), (f"({pssh_display})", "text"), overflow="fold")) diff --git a/unshackle/core/downloaders/__init__.py b/unshackle/core/downloaders/__init__.py index bc68b6fa..390fab82 100644 --- a/unshackle/core/downloaders/__init__.py +++ b/unshackle/core/downloaders/__init__.py @@ -1,3 +1,3 @@ -from .requests import format_speed, parse_speed_limit, requests, set_speed_limit +from .requests import default_max_workers, format_speed, parse_speed_limit, requests, set_speed_limit -__all__ = ("requests", "format_speed", "parse_speed_limit", "set_speed_limit") +__all__ = ("requests", "default_max_workers", "format_speed", "parse_speed_limit", "set_speed_limit") diff --git a/unshackle/core/downloaders/requests.py b/unshackle/core/downloaders/requests.py index b6a8b946..e01f9650 100644 --- a/unshackle/core/downloaders/requests.py +++ b/unshackle/core/downloaders/requests.py @@ -1019,6 +1019,11 @@ def build_session_spec(session: Optional[Any]) -> Optional[dict[str, Any]]: return None +def default_max_workers() -> int: + """Return the per-track worker count to use when the caller gives none.""" + return min(16, (os.cpu_count() or 1) + 4) + + def rebuild_session(spec: dict[str, Any], max_workers: int) -> Optional[Any]: """Reconstruct an HTTP session inside a child process from a spec built by ``build_session_spec``. @@ -1343,7 +1348,7 @@ def requests( urls = [urls] if not max_workers: - max_workers = min(16, (os.cpu_count() or 1) + 4) + max_workers = default_max_workers() # A spawned child re-imports this module and would start with no TokenBucket, so the # configured cap would silently stop applying; the cap wins over the fan-out, matching diff --git a/unshackle/core/service.py b/unshackle/core/service.py index 86b9665b..5ec0b36e 100644 --- a/unshackle/core/service.py +++ b/unshackle/core/service.py @@ -82,6 +82,31 @@ class TimeoutHTTPAdapter(HTTPAdapter): return super().send(request, **kwargs) +def grow_session_pool(session: Any, size: int) -> None: + """Grow a shared requests session's connection pool to ``size`` before track downloads start. + + The worker threads of every track draw on this one pool, because the downloader never + remounts an HTTP session the caller passes in (see downloaders/requests.py). The pool must hold + ``downloads * workers`` connections, or threads queue for a slot instead of reading. + Call this before any download thread exists: a remount races with other threads that call + ``get_adapter``. RnetSession does not block on its idle-pool cap, so this function skips it. + """ + if not isinstance(session, requests.Session): + return + adapter = session.get_adapter("https://") + if not isinstance(adapter, HTTPAdapter) or getattr(adapter, "_pool_maxsize", 0) >= size: + return + grown = TimeoutHTTPAdapter( + max_retries=adapter.max_retries, + pool_connections=size, + pool_maxsize=size, + pool_block=True, + timeout=getattr(adapter, "default_timeout", DEFAULT_TIMEOUT), + ) + session.mount("https://", grown) + session.mount("http://", grown) + + @dataclass class TrackRequest: """Holds what the user requested for video codec and range selection. diff --git a/unshackle/core/tracks/chapters.py b/unshackle/core/tracks/chapters.py index b58a6fa7..71898bd4 100644 --- a/unshackle/core/tracks/chapters.py +++ b/unshackle/core/tracks/chapters.py @@ -1,6 +1,8 @@ from __future__ import annotations +import os import re +import threading from abc import ABC from pathlib import Path from typing import Any, Iterable, Optional, Union @@ -120,7 +122,10 @@ class Chapters(SortedKeyList, ABC): path.parent.mkdir(parents=True, exist_ok=True) ogm_text = self.dumps(*args, **kwargs) - return path.write_text(ogm_text, encoding="utf8") + tmp = path.with_name(f"{path.name}.{os.getpid()}.{threading.get_ident()}.tmp") + written = tmp.write_text(ogm_text, encoding="utf8") + os.replace(tmp, path) + return written def add(self, value: Chapter) -> None: if not isinstance(value, Chapter): diff --git a/unshackle/core/vault.py b/unshackle/core/vault.py index d62f61d9..3abbe52a 100644 --- a/unshackle/core/vault.py +++ b/unshackle/core/vault.py @@ -4,6 +4,9 @@ from uuid import UUID class Vault(metaclass=ABCMeta): + # a vault on this machine; Vaults.get_key tries these first, before the network ones + local: bool = False + def __init__(self, name: str, no_push: bool = False): self.name = name self.no_push = no_push diff --git a/unshackle/core/vaults.py b/unshackle/core/vaults.py index e878e717..e6745d61 100644 --- a/unshackle/core/vaults.py +++ b/unshackle/core/vaults.py @@ -1,6 +1,7 @@ import inspect import logging import time +from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Iterator, Optional, Union from uuid import UUID @@ -56,39 +57,62 @@ class Vaults: self.vaults.append(vault) def get_key(self, kid: Union[UUID, str]) -> tuple[Optional[str], Optional[Vault]]: - """Get Key from the first Vault it can by KID (Key ID) and Service.""" + """Get Key from the first Vault it can by KID (Key ID) and Service. + + Local vaults go first, one at a time: a hit there costs nothing and skips the network. + The remaining vaults run at the same time, and the first one with the content key wins. + """ + local = [v for v in self.vaults if v.local] + remote = [v for v in self.vaults if not v.local] + for vault in local: + key = self.query_vault(vault, kid) + if key: + return key, vault + if not remote: + return None, None + pool = ThreadPoolExecutor(len(remote)) + try: + futures = {pool.submit(self.query_vault, vault, kid): vault for vault in remote} + for future in as_completed(futures): + key = future.result() + if key: + return key, futures[future] + return None, None + finally: + # a slow vault must not hold up a content key another vault already returned + pool.shutdown(wait=False, cancel_futures=True) + + def query_vault(self, vault: Vault, kid: Union[UUID, str]) -> Optional[str]: + """Ask one vault for a content key. A failure logs a warning and returns None.""" dl = get_debug_logger() - for vault in self.vaults: - start = time.monotonic() - try: - key = vault.get_key(kid, self.service) - except (PermissionError, NotImplementedError): - continue - except Exception as e: - log.warning(f"Failed to get key from Vault '{vault.name}': {e}") - if dl: - dl.log_vault_query( - vault.name, - "get_key", - kid=str(kid), - reachable=False, - error=e, - duration_ms=round((time.monotonic() - start) * 1000, 1), - ) - continue - found = bool(key and key.count("0") != len(key)) + start = time.monotonic() + try: + key = vault.get_key(kid, self.service) + except (PermissionError, NotImplementedError): + return None + except Exception as e: + log.warning(f"Failed to get key from Vault '{vault.name}': {e}") if dl: dl.log_vault_query( vault.name, "get_key", kid=str(kid), - key_found=found, - reachable=True, + reachable=False, + error=e, duration_ms=round((time.monotonic() - start) * 1000, 1), ) - if found: - return key, vault - return None, None + return None + found = bool(key and key.count("0") != len(key)) + if dl: + dl.log_vault_query( + vault.name, + "get_key", + kid=str(kid), + key_found=found, + reachable=True, + duration_ms=round((time.monotonic() - start) * 1000, 1), + ) + return key if found else None def add_key(self, kid: Union[UUID, str], key: str, excluding: Optional[Vault] = None) -> int: """Add a KID:KEY to all Vaults, optionally with an exclusion.""" diff --git a/unshackle/vaults/SQLite.py b/unshackle/vaults/SQLite.py index ab293ef9..187d0586 100644 --- a/unshackle/vaults/SQLite.py +++ b/unshackle/vaults/SQLite.py @@ -11,6 +11,8 @@ from unshackle.core.vault import Vault class SQLite(Vault): """Key Vault using a locally-accessed sqlite DB file.""" + local = True + def __init__(self, name: str, path: Union[str, Path], no_push: bool = False): super().__init__(name, no_push) self.path = Path(path).expanduser()