kenzuya/unshackle
1
0
Fork 2
mirror of https://github.com/unshackle-dl/unshackle.git synced 2026-11-14 04:55:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(dl): push vault key writes to a background thread

Browse Source 
Tracks stayed in LICENSING while a key found in a local vault was pushed to remote vaults under a global lock. Queue those writes on one worker, snapshot content_keys at submit, log writer failures, move --export out of the DRM lock, and join the queue before the run exits.
This commit is contained in:
imSp4rky
2026-08-24 20:29:44 -06:00
parent 65353af5c8
commit 1543e4d22f
1 changed files with 24 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
unshackle/commands/dl.py
+24 -18
View File
@@ -953,7 +953,7 @@ class dl:
DRM_LOCKS: dict[str, Lock] = {}
DRM_LOCKS_GUARD = Lock()
VAULT_WRITE_LOCK = Lock()
VAULT_WRITER = ThreadPoolExecutor(1, thread_name_prefix="vault-writer")
EXPORT_LOCK = Lock()
LICENSE_KEY_CACHE: dict[UUID, str] = {}
@@ -2949,6 +2949,7 @@ class dl:
message="Download cancelled by user",
context={"title": str(title)},
)
self.wait_vault_writes()
return
except Exception as e: # noqa
# Reported and swallowed (no re-raise) so the CLI exits cleanly; flag it so the
@@ -2990,6 +2991,7 @@ class dl:
),
postscript,
)
self.wait_vault_writes()
return
if skip_dl:
@@ -3626,6 +3628,7 @@ class dl:
run_context["folder"] = str(folder)
dispatch("success", "run", run_context, postscript)
self.wait_vault_writes()
dl_time = time_elapsed_since(start_time)
console.print(Padding(f"Processed all titles in [progress.elapsed]{dl_time}", (0, 5, 1, 5)))
@@ -3916,13 +3919,19 @@ class dl:
export.write_text(json.dumps(doc, indent=4, ensure_ascii=False), encoding="utf8")
def flush_vault_writes(self, pending: list[Callable[[], Any]]) -> None:
if not pending:
return
with self.VAULT_WRITE_LOCK:
for write in pending:
write()
"""Hand the queued vault writes to the background writer so the track does not wait on them."""
for write in pending:
self.VAULT_WRITER.submit(write).add_done_callback(self.report_vault_write)
pending.clear()
def report_vault_write(self, future: futures.Future[Any]) -> None:
if exc := future.exception():
self.log.error(f"Vault write failed: {exc!r}")
def wait_vault_writes(self) -> None:
"""Block until every queued vault write has run."""
self.VAULT_WRITER.submit(lambda: None).result()
def cache_keys_to_vaults(self, content_keys: dict[UUID, str]) -> None:
successful_caches = self.vaults.add_keys(content_keys)
self.log.info(
@@ -4245,7 +4254,7 @@ class dl:
self.LICENSE_KEY_CACHE.update(drm.content_keys)
pending_vault_writes.append(partial(self.cache_keys_to_vaults, drm.content_keys))
pending_vault_writes.append(partial(self.cache_keys_to_vaults, dict(drm.content_keys)))
if track_kid and track_kid not in drm.content_keys:
msg = f"No Content Key for KID {track_kid.hex} was returned in the License"
@@ -4259,10 +4268,9 @@ class dl:
table.add_row()
table.add_row(cek_tree)
if export:
self.write_export(export, title, track, drm)
self.flush_vault_writes(pending_vault_writes)
if export:
self.write_export(export, title, track, drm)
elif isinstance(drm, PlayReady):
if self.debug_logger:
@@ -4418,7 +4426,7 @@ class dl:
self.LICENSE_KEY_CACHE.update(drm.content_keys)
pending_vault_writes.append(partial(self.cache_keys_to_vaults, drm.content_keys))
pending_vault_writes.append(partial(self.cache_keys_to_vaults, dict(drm.content_keys)))
if track_kid and track_kid not in drm.content_keys:
msg = f"No Content Key for KID {track_kid.hex} was returned in the License"
@@ -4432,10 +4440,9 @@ class dl:
table.add_row()
table.add_row(cek_tree)
if export:
self.write_export(export, title, track, drm)
self.flush_vault_writes(pending_vault_writes)
if export:
self.write_export(export, title, track, drm)
elif isinstance(drm, ClearKeyCENC):
pending_vault_writes = []
@@ -4558,7 +4565,7 @@ class dl:
self.LICENSE_KEY_CACHE.update(drm.content_keys)
pending_vault_writes.append(partial(self.cache_keys_to_vaults, drm.content_keys))
pending_vault_writes.append(partial(self.cache_keys_to_vaults, dict(drm.content_keys)))
if track_kid and track_kid not in drm.content_keys:
msg = f"No Content Key for KID {track_kid.hex} was returned in the License"
@@ -4572,10 +4579,9 @@ class dl:
table.add_row()
table.add_row(cek_tree)
if export:
self.write_export(export, title, track, drm)
self.flush_vault_writes(pending_vault_writes)
if export:
self.write_export(export, title, track, drm)
elif isinstance(drm, MonaLisa):
with self.drm_lock(drm):