mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-11 19:45:06 +00:00
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:
@@ -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.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -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`.
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
@@ -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",)
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user