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

@@ -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

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)

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"))

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")

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

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.

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):

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

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."""

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()