From b391ffc80a647637224f70dfef8acf72ec532e61 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Fri, 4 Sep 2026 09:13:31 -0600 Subject: [PATCH] perf(kv): push keys to vaults concurrently, log one summary per title Vaults.add_key and add_keys now push to every vault at the same time, so one unreachable vault costs its own timeout instead of the sum. MySQL defaults connect_timeout to vault_timeout. --- docs/guide/vaults.md | 4 +- docs/reference/configuration/vaults.md | 5 +- unshackle/commands/dl.py | 15 ++++-- unshackle/core/api/handlers.py | 2 +- unshackle/core/vaults.py | 72 ++++++++++++++++---------- unshackle/vaults/MySQL.py | 2 + 6 files changed, 63 insertions(+), 37 deletions(-) diff --git a/docs/guide/vaults.md b/docs/guide/vaults.md index cead2ee1..3a2968c4 100644 --- a/docs/guide/vaults.md +++ b/docs/guide/vaults.md @@ -63,7 +63,9 @@ Network-backed vaults (`HTTP` and `API`) accept an optional per-vault `timeout` vault_timeout: 10.0 # default, in seconds ``` -A `timeout` set inline on an individual vault always wins over this global value. The `SQLite` and `MySQL` backends do not take a `timeout` field. +A `timeout` set inline on an individual vault always wins over this global value. The `SQLite` and `MySQL` backends do not take a `timeout` field. `MySQL` uses `vault_timeout` as its `connect_timeout` instead, and falls back to the pymysql default of 10 seconds when `vault_timeout` is `0`. To use a different connect timeout, set `connect_timeout` on the vault. + +unshackle pushes content keys to all vaults at the same time, so one slow or unreachable vault costs only its own timeout and does not delay the others. --- diff --git a/docs/reference/configuration/vaults.md b/docs/reference/configuration/vaults.md index 904aea6d..70442e8c 100644 --- a/docs/reference/configuration/vaults.md +++ b/docs/reference/configuration/vaults.md @@ -13,7 +13,7 @@ name) and a `name`, plus backend-specific keys. | `type` | Purpose | Required keys | Notes | |--------|---------|--------------|-------| | `SQLite` | Local SQLite database | `name`, `path` | Loaded **critically**; a failure aborts the run. | -| `MySQL` | Remote MySQL database | `name`, `host`, `database`, `username` | Extra keys (e.g. `password`, `port`) forwarded to pymysql. | +| `MySQL` | Remote MySQL database | `name`, `host`, `database`, `username` | Extra keys (e.g. `password`, `port`) forwarded to pymysql. `connect_timeout` defaults to `vault_timeout`. | | `API` | RESTful JSON API | `name`, `uri`, `token` | Optional `headers` map is sent with every request. Honours `vault_timeout`. | | `HTTP` | HTTP API with modes | `name`, `host`, one of `password`/`api_key`, and `username` in query mode | `api_mode`: `query` (default), `json`, `decrypt_labs`. Honours `vault_timeout`. | @@ -44,4 +44,5 @@ key_vaults: - **Type:** `float`  ยท  **Default:** `10.0` Timeout in seconds for vault operations. Injected automatically into any backend whose -constructor accepts a `timeout` parameter (a per-vault `timeout` still wins). +constructor accepts a `timeout` parameter (a per-vault `timeout` still wins). The `MySQL` backend +uses it as the default `connect_timeout`. diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index 656cd895..5cd9310b 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -1211,6 +1211,7 @@ class dl: self.cdm_override = ctx.params.get("cdm_name") cdm_only = ctx.params.get("cdm_only") + self.vault_cache_tally: Optional[tuple[int, int]] = None if cdm_only: self.vaults = Vaults(self.vault_service) @@ -3955,15 +3956,19 @@ class dl: self.log.error(f"Vault write failed: {exc!r}") def wait_vault_writes(self) -> None: - """Block until every queued vault write has run.""" + """Block until every queued vault write has run, then log the total once.""" self.VAULT_WRITER.submit(lambda: None).result() + if not (self.vault_cache_tally and self.vault_cache_tally[0]): + self.vault_cache_tally = None + return + keys, successful_caches = self.vault_cache_tally + self.vault_cache_tally = None + self.log.info(f"Cached {keys} Key{'' if keys == 1 else 's'} to {successful_caches}/{len(self.vaults)} Vaults") def cache_keys_to_vaults(self, content_keys: dict[UUID, str]) -> None: successful_caches = self.vaults.add_keys(content_keys) - self.log.info( - f"Cached {len(content_keys)} Key{'' if len(content_keys) == 1 else 's'} to " - f"{successful_caches}/{len(self.vaults)} Vaults" - ) + keys, caches = self.vault_cache_tally or (0, successful_caches) + self.vault_cache_tally = (keys + len(content_keys), min(caches, successful_caches)) @classmethod def drm_lock(cls, drm: DRM_T) -> Lock: diff --git a/unshackle/core/api/handlers.py b/unshackle/core/api/handlers.py index b4acfdde..ed64f81c 100644 --- a/unshackle/core/api/handlers.py +++ b/unshackle/core/api/handlers.py @@ -3566,7 +3566,7 @@ def cache_to_vaults(keys: Dict[str, str], service_name: str) -> None: key_map = {UUID(hex=kid): key for kid, key in keys.items()} cached = vaults.add_keys(key_map) if cached: - log.info(f"Cached {cached} key(s) to {cached} server vault(s)") + log.info(f"Cached {len(key_map)} key(s) to {cached}/{len(vaults)} server vault(s)") except (Exception, SystemExit) as e: log.warning(f"Failed to cache keys to vaults: {e}") diff --git a/unshackle/core/vaults.py b/unshackle/core/vaults.py index e6745d61..e46a1a07 100644 --- a/unshackle/core/vaults.py +++ b/unshackle/core/vaults.py @@ -115,42 +115,58 @@ class Vaults: 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.""" + """Add a KID:KEY to all Vaults, optionally with an exclusion. + + This method pushes to all Vaults at the same time, so one unreachable Vault costs only its + own timeout and does not delay the others. + """ + vaults = [vault for vault in self.vaults if vault != excluding and not vault.no_push] + if not vaults: + return 0 + with ThreadPoolExecutor(len(vaults), thread_name_prefix="vault-push") as pool: + return sum(pool.map(lambda vault: self.push_key(vault, kid, key), vaults)) + + def push_key(self, vault: Vault, kid: Union[UUID, str], key: str) -> int: + """Push one Content Key to one Vault. Returns 1 if the Vault stored it, 0 otherwise.""" dl = get_debug_logger() - success = 0 - for vault in self.vaults: - if vault != excluding and not vault.no_push: - try: - added = vault.add_key(self.service, kid, key) - success += added - if dl: - dl.log_vault_query(vault.name, "add_key", kid=str(kid), success=bool(added)) - except (PermissionError, NotImplementedError): - pass - except Exception as e: - log.warning(f"Failed to add key to Vault '{vault.name}': {e}") - if dl: - dl.log_vault_query(vault.name, "add_key", kid=str(kid), success=False, error=e) - return success + try: + added = vault.add_key(self.service, kid, key) + if dl: + dl.log_vault_query(vault.name, "add_key", kid=str(kid), success=bool(added)) + return int(added) + except (PermissionError, NotImplementedError): + return 0 + except Exception as e: + log.warning(f"Failed to add key to Vault '{vault.name}': {e}") + if dl: + dl.log_vault_query(vault.name, "add_key", kid=str(kid), success=False, error=e) + return 0 def add_keys(self, kid_keys: dict[Union[UUID, str], str]) -> int: """ Add multiple KID:KEYs to all Vaults. The Vaults skip duplicate Content Keys. This method absorbs and ignores the PermissionError a Vault raises when it cannot make Tables. This method also skips Vaults with no_push=True. + + This method pushes to all Vaults at the same time, so one unreachable Vault costs only its + own timeout and does not delay the others. """ - success = 0 - for vault in self.vaults: - if not vault.no_push: - try: - # Count each vault that successfully processes the keys (whether new or existing) - vault.add_keys(self.service, kid_keys) - success += 1 - except (PermissionError, NotImplementedError): - pass - except Exception as e: - log.warning(f"Failed to add keys to Vault '{vault.name}': {e}") - return success + vaults = [vault for vault in self.vaults if not vault.no_push] + if not vaults: + return 0 + with ThreadPoolExecutor(len(vaults), thread_name_prefix="vault-push") as pool: + return sum(pool.map(lambda vault: self.push_keys(vault, kid_keys), vaults)) + + def push_keys(self, vault: Vault, kid_keys: dict[Union[UUID, str], str]) -> int: + """Push Content Keys to one Vault. Returns 1 if the Vault accepted them, 0 on any failure.""" + try: + vault.add_keys(self.service, kid_keys) + return 1 + except (PermissionError, NotImplementedError): + return 0 + except Exception as e: + log.warning(f"Failed to add keys to Vault '{vault.name}': {e}") + return 0 __all__ = ("Vaults",) diff --git a/unshackle/vaults/MySQL.py b/unshackle/vaults/MySQL.py index d2c43eda..8c8a0c52 100644 --- a/unshackle/vaults/MySQL.py +++ b/unshackle/vaults/MySQL.py @@ -5,6 +5,7 @@ from uuid import UUID import pymysql from pymysql.cursors import DictCursor +from unshackle.core.config import config from unshackle.core.vault import Vault @@ -18,6 +19,7 @@ class MySQL(Vault): """ super().__init__(name, no_push) self.slug = f"{host}:{database}:{username}" + kwargs.setdefault("connect_timeout", config.vault_timeout or 10) self.conn_factory = ConnectionFactory( dict(host=host, db=database, user=username, cursorclass=DictCursor, **kwargs) )