From 47d9e3d3c92bb15c09ece63bc89fe11ee0dc7405 Mon Sep 17 00:00:00 2001 From: imSp4rky Date: Mon, 2 Nov 2026 14:59:04 -0700 Subject: [PATCH] perf(api): move blocking file and hash work off the event loop Run session-tracks serialization, request cache IO, session cache dir removal (now outside the store lock), history reads and gzip over 1 MiB on worker threads; cache the per-key PBKDF2 namespace and key_id. --- tests/remote/unit/test_compression.py | 22 +++++++++ unshackle/core/api/compression.py | 7 ++- unshackle/core/api/handlers.py | 71 +++++++++++++++++---------- unshackle/core/api/session_store.py | 49 +++++++++++------- unshackle/core/api/stats.py | 1 + 5 files changed, 105 insertions(+), 45 deletions(-) diff --git a/tests/remote/unit/test_compression.py b/tests/remote/unit/test_compression.py index 67993a29..2c494b94 100644 --- a/tests/remote/unit/test_compression.py +++ b/tests/remote/unit/test_compression.py @@ -62,3 +62,25 @@ def test_compresses_large_json_when_accepted() -> None: decompressed = gzip.decompress(resp.body) assert json.loads(decompressed) == big assert resp.content_length == len(resp.body) + + +def test_compresses_a_body_over_the_thread_size_on_a_worker_thread(monkeypatch) -> None: + from unshackle.core.api import compression + + threaded = [] + real = asyncio.to_thread + + async def to_thread(fn, *args, **kwargs): + threaded.append(fn) + return await real(fn, *args, **kwargs) + + monkeypatch.setattr(compression.asyncio, "to_thread", to_thread) + monkeypatch.setattr(compression, "GZIP_THREAD_SIZE", 1024) + big = {"data": "x" * 4096} + + async def handler(req): # noqa: ARG001 + return web.json_response(big) + + resp = run(compression_middleware(_FakeReq(), handler)) + assert threaded == [gzip.compress] + assert json.loads(gzip.decompress(resp.body)) == big diff --git a/unshackle/core/api/compression.py b/unshackle/core/api/compression.py index bb5a62a9..f2515b34 100644 --- a/unshackle/core/api/compression.py +++ b/unshackle/core/api/compression.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import gzip import zlib from typing import Any @@ -11,6 +12,7 @@ from aiohttp import web # Cap the inflated size of a peer-supplied compressed blob (cookies, session cache, # serialized manifests) so a tiny payload cannot expand without bound (zip-bomb DoS). MAX_INFLATE_BYTES = 64 * 1024 * 1024 # 64 MiB +GZIP_THREAD_SIZE = 1024 * 1024 def safe_inflate(data: bytes, max_size: int = MAX_INFLATE_BYTES) -> bytes: @@ -50,6 +52,9 @@ async def compression_middleware(request: web.Request, handler: Any) -> web.Stre if not level: return response - response.body = gzip.compress(body, compresslevel=level) + if len(body) > GZIP_THREAD_SIZE: + response.body = await asyncio.to_thread(gzip.compress, body, compresslevel=level) + else: + response.body = gzip.compress(body, compresslevel=level) response.headers["Content-Encoding"] = "gzip" return response diff --git a/unshackle/core/api/handlers.py b/unshackle/core/api/handlers.py index 894c6383..dc11c3de 100644 --- a/unshackle/core/api/handlers.py +++ b/unshackle/core/api/handlers.py @@ -2,6 +2,7 @@ import asyncio import base64 import copy import enum +import functools import json import logging import re @@ -328,9 +329,14 @@ def server_login_material( def api_key_namespace(request: Optional[web.Request]) -> str: """Name of the caller's cache directories, derived from its API key so no caller can guess another's.""" + return key_namespace(request.headers.get("X-Secret-Key", "anonymous") if request else "anonymous") + + +@functools.lru_cache(maxsize=256) +def key_namespace(api_key: str) -> str: + """The PBKDF2 behind :func:`api_key_namespace`, kept per key: 100k rounds cost about 20 ms a call.""" import hashlib - api_key = request.headers.get("X-Secret-Key", "anonymous") if request else "anonymous" return hashlib.pbkdf2_hmac("sha256", api_key.encode(), b"unshackle-session-ns", 100_000).hex()[:12] @@ -423,16 +429,16 @@ class RequestCache: write_client_cache(cache_data, self.tag) self.client_auth = client_login or bool(cache_data) - def json_response(self, payload: Dict[str, Any]) -> web.Response: + async def json_response(self, payload: Dict[str, Any]) -> web.Response: """Answer with ``payload``, plus the updated ``cache`` for a client that logged in itself.""" if self.tag and self.client_auth: - cache_data = collect_cache_files(self.tag) + cache_data = await asyncio.to_thread(collect_cache_files, self.tag) if cache_data: payload["cache"] = cache_data return web.json_response(payload) - def cleanup(self) -> None: - SessionStore.cleanup_cache_dir(self.tag) + async def cleanup(self) -> None: + await SessionStore.remove_cache_dirs([self.tag]) def remote_list(name: Any) -> list[dict]: @@ -1617,6 +1623,7 @@ def extract_track_manifests(tracks: Tracks) -> List[Dict[str, Any]]: compression_level = app_config.serve.get("compression_level", 1) fragments: List[Dict[str, Any]] = [] + representations: Dict[int, int] = {} for track in list(tracks.videos) + list(tracks.audio) + list(tracks.subtitles): dash = (getattr(track, "data", None) or {}).get("dash") or {} @@ -1624,8 +1631,10 @@ def extract_track_manifests(tracks: Tracks) -> List[Dict[str, Any]]: adaptation_set = dash.get("adaptation_set") if manifest is None or adaptation_set is None: continue - if adaptation_set.getroottree().getroot() is manifest and len(manifest.findall(".//Representation")) > 1: - continue + if adaptation_set.getroottree().getroot() is manifest: + count = once(representations, id(manifest), lambda: len(manifest.findall(".//Representation"))) + if count > 1: + continue root = etree.Element("MPD", dict(manifest.attrib)) for child in manifest: @@ -1924,7 +1933,7 @@ async def search_handler(data: Dict[str, Any], request: Optional[web.Request] = request_cache = RequestCache() try: results = await asyncio.to_thread(run_service_search, data, normalized_service, query, request, request_cache) - return request_cache.json_response({"results": results, "count": len(results)}) + return await request_cache.json_response({"results": results, "count": len(results)}) except APIError: raise except (Exception, SystemExit) as e: @@ -1934,7 +1943,7 @@ async def search_handler(data: Dict[str, Any], request: Optional[web.Request] = e, context={"operation": "search", "service": normalized_service}, debug_mode=debug_mode ) finally: - request_cache.cleanup() + await request_cache.cleanup() async def list_titles_handler(data: Dict[str, Any], request: Optional[web.Request] = None) -> web.Response: @@ -1967,7 +1976,7 @@ async def list_titles_handler(data: Dict[str, Any], request: Optional[web.Reques title_list = [stamp_service_flags(serialize_title(t), service_instance) for t in listed_titles(titles)] - return request_cache.json_response({"titles": title_list}) + return await request_cache.json_response({"titles": title_list}) except APIError: raise @@ -1982,7 +1991,7 @@ async def list_titles_handler(data: Dict[str, Any], request: Optional[web.Reques finally: if lease is not None: await release_remote_service(lease) - request_cache.cleanup() + await request_cache.cleanup() async def list_tracks_handler(data: Dict[str, Any], request: Optional[web.Request] = None) -> web.Response: @@ -2121,7 +2130,7 @@ async def list_tracks_handler(data: Dict[str, Any], request: Optional[web.Reques response: Dict[str, Any] = {"episodes": episodes_data} if failed_episodes: response["unavailable_episodes"] = failed_episodes - return request_cache.json_response(response) + return await request_cache.json_response(response) else: raise APIError( APIErrorCode.NO_CONTENT, @@ -2152,7 +2161,7 @@ async def list_tracks_handler(data: Dict[str, Any], request: Optional[web.Reques "subtitles": [serialize_subtitle_track(t) for t in tracks.subtitles], } - return request_cache.json_response(response) + return await request_cache.json_response(response) except APIError: raise @@ -2167,7 +2176,7 @@ async def list_tracks_handler(data: Dict[str, Any], request: Optional[web.Reques finally: if lease is not None: await release_remote_service(lease) - request_cache.cleanup() + await request_cache.cleanup() VALID_VCODECS = [choice.upper() for choice in VIDEO_CODEC_LIST.choices] @@ -3505,11 +3514,15 @@ async def download_history_handler(data: Dict[str, Any], request: Optional[web.R allowed = get_allowed_services(request) if allowed is None: - history = read_job_history(limit=limit, service=data.get("service"), owner=owner_id(caller_key(request))) + history = await asyncio.to_thread( + read_job_history, limit=limit, service=data.get("service"), owner=owner_id(caller_key(request)) + ) else: # Read unbounded, drop entries outside the caller's allowlist, then apply limit. allowed_upper = {a.upper() for a in allowed} - entries = read_job_history(limit=0, service=data.get("service"), owner=owner_id(caller_key(request))) + entries = await asyncio.to_thread( + read_job_history, limit=0, service=data.get("service"), owner=owner_id(caller_key(request)) + ) history = [e for e in entries if str(e.get("service") or "").upper() in allowed_upper][:limit] return web.json_response({"history": history, "count": len(history)}) @@ -4219,8 +4232,9 @@ async def session_tracks_handler( video_tracks = sorted(tracks.videos, key=lambda t: t.bitrate or 0, reverse=True) audio_tracks = sorted(tracks.audio, key=lambda t: t.bitrate or 0, reverse=True) - manifests = extract_manifests(tracks) - track_manifests = extract_track_manifests(tracks) + manifests, track_manifests = await asyncio.to_thread( + lambda: (extract_manifests(tracks), extract_track_manifests(tracks)) + ) session_headers, session_cookies = client_session_auth(session) @@ -4255,12 +4269,17 @@ async def session_tracks_handler( else: server_cdm_type = "playready" - original_ids = original_audio_ids(audio_tracks, title) - video = [serialize_video_track(t, include_url=True) for t in video_tracks] - audio = [serialize_audio_track(t, include_url=True, is_original=t.id in original_ids) for t in audio_tracks] - subtitles = [serialize_subtitle_track(t, include_url=True) for t in tracks.subtitles] - for d in video + audio + subtitles: - d["data"] = scrub_track_data(d.get("data"), session) + def serialize() -> tuple[list, list, list, list]: + """The track and attachment lists, off the event loop: an attachment can be a 20 MB file.""" + original_ids = original_audio_ids(audio_tracks, title) + video = [serialize_video_track(t, include_url=True) for t in video_tracks] + audio = [serialize_audio_track(t, include_url=True, is_original=t.id in original_ids) for t in audio_tracks] + subtitles = [serialize_subtitle_track(t, include_url=True) for t in tracks.subtitles] + for d in video + audio + subtitles: + d["data"] = scrub_track_data(d.get("data"), session) + return video, audio, subtitles, [a for a in map(serialize_attachment, tracks.attachments) if a] + + video, audio, subtitles, attachments = await asyncio.to_thread(serialize) return web.json_response( { @@ -4272,7 +4291,7 @@ async def session_tracks_handler( {"timestamp": ch.timestamp, "name": ch.name} for ch in session.chapters_by_title.get(str(title_id), []) ], - "attachments": [a for a in map(serialize_attachment, tracks.attachments) if a], + "attachments": attachments, "manifests": manifests, "track_manifests": track_manifests, "session_headers": session_headers, @@ -4955,7 +4974,7 @@ def load_server_vaults(service_name: str) -> Any: return vaults -def once(shared: Dict[str, Any], key: Any, make: Callable[[], Any]) -> Any: +def once(shared: Dict[Any, Any], key: Any, make: Callable[[], Any]) -> Any: """``make()`` the first time ``key`` is asked for in ``shared``, the stored value after that.""" if key not in shared: shared[key] = make() diff --git a/unshackle/core/api/session_store.py b/unshackle/core/api/session_store.py index aef68104..fda282fc 100644 --- a/unshackle/core/api/session_store.py +++ b/unshackle/core/api/session_store.py @@ -164,6 +164,7 @@ class SessionStore: server_account: Optional[str] = None, ) -> SessionEntry: """Make a new remote session with an authenticated service instance.""" + evicted_tag: Optional[str] = None async with self._lock: max_sessions = self.max_sessions if max_sessions is not None and len(self._sessions) >= max_sessions: @@ -173,7 +174,7 @@ class SessionStore: evicted = self._sessions.pop(oldest_id) if evicted.input_bridge: evicted.input_bridge.cancel() - self.cleanup_cache_dir(evicted.cache_tag) + evicted_tag = evicted.cache_tag _publish("delete", evicted, "evicted") _schedule_pending_reload(evicted.service_tag) @@ -189,7 +190,8 @@ class SessionStore: self._sessions[session_id] = entry _publish("create", entry) log.info(f"Created session {sanitize_log(session_id)} for service {sanitize_log(service_tag)}") - return entry + await self.remove_cache_dirs([evicted_tag]) + return entry async def get(self, session_id: str) -> Optional[SessionEntry]: """Get a remote session by ID, returns None if not found or expired.""" @@ -198,17 +200,19 @@ class SessionStore: if entry is None: return None + expired = False if not entry.in_use and entry.auth_status not in (AuthStatus.AUTHENTICATING, AuthStatus.PENDING_INPUT): elapsed = (datetime.now(timezone.utc) - entry.last_accessed).total_seconds() if elapsed > self.ttl: log.info(f"Session {sanitize_log(session_id)} expired (elapsed={elapsed:.0f}s, ttl={self.ttl}s)") - self.cleanup_cache_dir(entry.cache_tag) _publish("delete", self._sessions.pop(session_id), "expired") _schedule_pending_reload(entry.service_tag) - return None - - entry.touch() - return entry + expired = True + if not expired: + entry.touch() + return entry + await self.remove_cache_dirs([entry.cache_tag]) + return None def peek(self, session_id: str) -> Optional[SessionEntry]: """A remote session entry without touching it, for read-only observers. @@ -222,15 +226,15 @@ class SessionStore: """Delete a remote session. Returns True if it existed.""" async with self._lock: entry = self._sessions.pop(session_id, None) - if entry: - if entry.input_bridge: - entry.input_bridge.cancel() - self.cleanup_cache_dir(entry.cache_tag) - _publish("delete", entry, "closed") - log.info(f"Deleted session {sanitize_log(session_id)}") - _schedule_pending_reload(entry.service_tag) - return True - return False + if not entry: + return False + if entry.input_bridge: + entry.input_bridge.cancel() + _publish("delete", entry, "closed") + log.info(f"Deleted session {sanitize_log(session_id)}") + _schedule_pending_reload(entry.service_tag) + await self.remove_cache_dirs([entry.cache_tag]) + return True async def cleanup_expired(self) -> int: """Remove all expired sessions. Returns count of removed sessions.""" @@ -247,18 +251,20 @@ class SessionStore: elif elapsed > self.ttl: expired.append(sid) ended_tags: set[str] = set() + cache_tags: list[Optional[str]] = [] for sid in expired: entry = self._sessions.pop(sid) if entry.input_bridge: entry.input_bridge.cancel() - self.cleanup_cache_dir(entry.cache_tag) + cache_tags.append(entry.cache_tag) _publish("delete", entry, "expired") ended_tags.add(entry.service_tag) for tag in ended_tags: _schedule_pending_reload(tag) if expired: log.info(f"Cleaned up {len(expired)} expired sessions") - return len(expired) + await self.remove_cache_dirs(cache_tags) + return len(expired) async def start_cleanup_loop(self) -> None: """Start periodic cleanup of expired sessions.""" @@ -292,6 +298,13 @@ class SessionStore: if count: log.info(f"Cancelled bridges for {count} active session(s)") + @classmethod + async def remove_cache_dirs(cls, cache_tags: List[Optional[str]]) -> None: + """Remove the cache directories of ended remote sessions on a worker thread, outside the store lock.""" + for tag in cache_tags: + if tag: + await asyncio.to_thread(cls.cleanup_cache_dir, tag) + @staticmethod def cleanup_cache_dir(cache_tag: Optional[str]) -> None: """Remove the remote session cache directory and empty parents.""" diff --git a/unshackle/core/api/stats.py b/unshackle/core/api/stats.py index 1d09ce19..71fb1782 100644 --- a/unshackle/core/api/stats.py +++ b/unshackle/core/api/stats.py @@ -72,6 +72,7 @@ def mask_key(key: Optional[str]) -> str: return key[:4] + "…" +@lru_cache(maxsize=1024) def key_id(key: Optional[str]) -> str: """A stable id for an API key. A hash prefix, so two unnamed keys never collide the way ``mask_key`` does, and the id carries no key material."""