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.
This commit is contained in:
imSp4rky
2026-09-04 09:13:31 -06:00
parent 1997a92c08
commit b391ffc80a
6 changed files with 63 additions and 37 deletions
+3 -1
View File
@@ -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.
---
+3 -2
View File
@@ -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`.
+10 -5
View File
@@ -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:
+1 -1
View File
@@ -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}")
+44 -28
View File
@@ -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",)
+2
View File
@@ -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)
)