From 41bca681e5e3d1a49e8a5ccac5c6df5408dd2708 Mon Sep 17 00:00:00 2001 From: arabcoders Date: Mon, 26 Jan 2026 16:36:14 +0300 Subject: [PATCH 1/6] refactor: limit item add --- app/features/tasks/router.py | 3 +-- app/library/config.py | 4 ++++ app/library/downloads/item_adder.py | 18 +++++++++++++++++- app/library/downloads/playlist_processor.py | 4 ++-- app/library/downloads/queue_manager.py | 9 +++++++-- 5 files changed, 31 insertions(+), 7 deletions(-) diff --git a/app/features/tasks/router.py b/app/features/tasks/router.py index 6c086ee5..9b94c095 100644 --- a/app/features/tasks/router.py +++ b/app/features/tasks/router.py @@ -439,12 +439,11 @@ async def task_metadata(request: Request, repo: TasksRepository, config: Config, opts["proxy"] = proxy try: - from httpx_curl_cffi import AsyncCurlTransport, CurlOpt + from httpx_curl_cffi import AsyncCurlTransport opts["transport"] = AsyncCurlTransport( impersonate="chrome", default_headers=True, - curl_options={CurlOpt.FRESH_CONNECT: True}, ) opts.pop("headers", None) except Exception: diff --git a/app/library/config.py b/app/library/config.py index c139ed42..6bca0d7e 100644 --- a/app/library/config.py +++ b/app/library/config.py @@ -207,6 +207,9 @@ class Config(metaclass=Singleton): live_premiere_buffer: int = 5 """The buffer time in minutes to add to video duration to wait before starting premiere download.""" + add_items_concurrency: int = 4 + """The number of concurrent add items to be processed at same time.""" + playlist_items_concurrency: int = 4 """The number of concurrent playlist items to be processed at same time.""" @@ -274,6 +277,7 @@ class Config(metaclass=Singleton): "auto_clear_history_days", "default_pagination", "extract_info_concurrency", + "add_items_concurrency", "flaresolverr_max_timeout", "flaresolverr_client_timeout", "flaresolverr_cache_ttl", diff --git a/app/library/downloads/item_adder.py b/app/library/downloads/item_adder.py index af7aeaba..40036a41 100644 --- a/app/library/downloads/item_adder.py +++ b/app/library/downloads/item_adder.py @@ -76,7 +76,11 @@ async def add_item( async def add( - queue: "DownloadQueue", item: "Item", already: set | None = None, entry: dict | None = None + queue: "DownloadQueue", + item: "Item", + already: set | None = None, + entry: dict | None = None, + playlist: bool = False, ) -> dict[str, str]: """ Add an item to the download queue. @@ -86,11 +90,23 @@ async def add( item: Item to be added to the queue already: Set of already downloaded items entry: Entry associated with the item (if already extracted) + playlist: Whether the item is part of a playlist Returns: dict[str, str]: Status dict with "status" and optional "msg" keys """ + if playlist: + # If not done like this, it would lead to deadlocks when adding playlist items. + return await _add_impl(queue=queue, item=item, already=already, entry=entry) + + async with queue.add_sem: + return await _add_impl(queue=queue, item=item, already=already, entry=entry) + + +async def _add_impl( + queue: "DownloadQueue", item: "Item", already: set | None = None, entry: dict | None = None +) -> dict[str, str]: _preset: Preset | None = Presets.get_instance().get(item.preset) if item.has_cli(): diff --git a/app/library/downloads/playlist_processor.py b/app/library/downloads/playlist_processor.py index 2a7006d3..1ea39dcd 100644 --- a/app/library/downloads/playlist_processor.py +++ b/app/library/downloads/playlist_processor.py @@ -98,9 +98,9 @@ async def process_playlist( ): dct = merge_dict(merge_dict({"_type": "video"}, etr), entry) dct.pop("entries", None) - return await queue.add(item=newItem, entry=dct, already=already) + return await queue.add(item=newItem, entry=dct, already=already, playlist=True) - return await queue.add(item=newItem, already=already) + return await queue.add(item=newItem, already=already, playlist=True) finally: if acquired: queue.processors.release() diff --git a/app/library/downloads/queue_manager.py b/app/library/downloads/queue_manager.py index 88a35ce6..c4ab5b11 100644 --- a/app/library/downloads/queue_manager.py +++ b/app/library/downloads/queue_manager.py @@ -43,6 +43,8 @@ class DownloadQueue(metaclass=Singleton): "DataStore for the download queue." self.processors = asyncio.Semaphore(self.config.playlist_items_concurrency) "Semaphore to limit the number of concurrent processors." + self.add_sem = asyncio.Semaphore(self.config.add_items_concurrency) + "Semaphore to limit the number of concurrent add item operations." self.pool = PoolManager(queue=self, config=self.config) "Pool manager for coordinating download execution." @@ -221,7 +223,9 @@ class DownloadQueue(metaclass=Singleton): async def on_shutdown(self, _: web.Application): await self.pool.shutdown() - async def add(self, item: Item, already: set | None = None, entry: dict | None = None) -> dict[str, str]: + async def add( + self, item: Item, already: set | None = None, entry: dict | None = None, playlist: bool = False + ) -> dict[str, str]: """ Add an item to the download queue. @@ -229,12 +233,13 @@ class DownloadQueue(metaclass=Singleton): item: Item to be added to the queue already: Set of already downloaded items entry: Entry associated with the item (if already extracted) + playlist: Whether the item is part of a playlist Returns: dict[str, str]: Status dict with "status" and optional "msg" keys """ - return await add_impl(queue=self, item=item, already=already, entry=entry) + return await add_impl(queue=self, item=item, already=already, entry=entry, playlist=playlist) async def cancel(self, ids: list[str]) -> dict[str, str]: """ From 98f25027bb5814c09b7e39c8465eb39da98cea7f Mon Sep 17 00:00:00 2001 From: arabcoders Date: Mon, 26 Jan 2026 17:42:24 +0300 Subject: [PATCH 2/6] refactor: use shared http clients --- .../definitions/handlers/_base_handler.py | 56 ++-- .../tasks/definitions/handlers/generic.py | 82 +++--- app/features/tasks/router.py | 100 +++---- app/library/UpdateChecker.py | 49 ++-- app/library/httpx_client.py | 221 +++++++++++++++- app/main.py | 2 + app/routes/api/docs.py | 67 ++--- app/routes/api/images.py | 250 +++++++++--------- app/tests/test_update_checker.py | 70 ++--- 9 files changed, 527 insertions(+), 370 deletions(-) diff --git a/app/features/tasks/definitions/handlers/_base_handler.py b/app/features/tasks/definitions/handlers/_base_handler.py index 0ee1e274..57b05c4c 100644 --- a/app/features/tasks/definitions/handlers/_base_handler.py +++ b/app/features/tasks/definitions/handlers/_base_handler.py @@ -1,12 +1,12 @@ # flake8: noqa: ARG004 -from typing import Any - -import httpx -from yt_dlp.utils.networking import random_user_agent +from typing import TYPE_CHECKING, Any from app.features.tasks.definitions.results import HandleTask, TaskFailure, TaskResult from app.library.config import Config -from app.library.httpx_client import async_client +from app.library.httpx_client import Globals, build_request_headers, get_async_client, resolve_curl_transport + +if TYPE_CHECKING: + import httpx class BaseHandler: @@ -33,7 +33,7 @@ class BaseHandler: @staticmethod async def request( url: str, headers: dict | None = None, ytdlp_opts: dict | None = None, **kwargs - ) -> httpx.Response: + ) -> "httpx.Response": """ Make an HTTP request. @@ -50,31 +50,21 @@ class BaseHandler: headers = {} if not isinstance(headers, dict) else headers ytdlp_opts = {} if not isinstance(ytdlp_opts, dict) else ytdlp_opts - opts: dict[str, Any] = { - "headers": { - "User-Agent": random_user_agent(), - }, - } + use_curl = resolve_curl_transport() + request_headers = build_request_headers( + base_headers=headers, + user_agent=Globals.get_random_agent(), + use_curl=use_curl, + ) - try: - from httpx_curl_cffi import AsyncCurlTransport, CurlOpt - - opts["transport"] = AsyncCurlTransport( - impersonate="chrome", - default_headers=True, - curl_options={CurlOpt.FRESH_CONNECT: True}, - ) - opts["headers"].pop("User-Agent", None) - except Exception: - pass - - for k, v in headers.items(): - opts["headers"][k] = v - - if proxy := ytdlp_opts.get("proxy", None): - opts["proxy"] = proxy - - async with async_client(**opts) as client: - method = kwargs.pop("method", "GET").upper() - timeout = ytdlp_opts.get("timeout", ytdlp_opts.get("socket_timeout", 120)) - return await client.request(method=method, url=url, timeout=timeout, **kwargs) + proxy = ytdlp_opts.get("proxy", None) + client = get_async_client(proxy=proxy, use_curl=use_curl) + method = kwargs.pop("method", "GET").upper() + timeout = ytdlp_opts.get("timeout", ytdlp_opts.get("socket_timeout", 120)) + return await client.request( + method=method, + url=url, + headers=request_headers, + timeout=timeout, + **kwargs, + ) diff --git a/app/features/tasks/definitions/handlers/generic.py b/app/features/tasks/definitions/handlers/generic.py index 90e68a72..d2ac1f8b 100644 --- a/app/features/tasks/definitions/handlers/generic.py +++ b/app/features/tasks/definitions/handlers/generic.py @@ -14,7 +14,6 @@ from urllib.parse import urljoin import jmespath from parsel import Selector -from yt_dlp.utils.networking import random_user_agent from app.features.tasks.definitions.results import HandleTask, TaskFailure, TaskItem, TaskResult from app.features.tasks.definitions.schemas import ( @@ -23,7 +22,7 @@ from app.features.tasks.definitions.schemas import ( ) from app.library.cache import Cache from app.library.config import Config -from app.library.httpx_client import async_client +from app.library.httpx_client import Globals, build_request_headers, get_async_client, resolve_curl_transport from app.library.Utils import fetch_info, get_archive_id from ._base_handler import BaseHandler @@ -48,7 +47,7 @@ class GenericTaskHandler(BaseHandler): """Modification times of source files to detect changes.""" @classmethod - async def refresh_definitions(cls, force: bool = False) -> None: + async def refresh_definitions(cls, force: bool = False) -> list[TaskDefinition]: """ Refresh the cached task definitions if source files have changed. @@ -66,12 +65,7 @@ class GenericTaskHandler(BaseHandler): repo = TaskDefinitionsRepository.get_instance() models = await repo.list() - definitions: list[TaskDefinition] = [] - for model in models: - td = model_to_schema(model) - definitions.append(td) - - cls._definitions = definitions + cls._definitions = [model_to_schema(model) for model in models] return cls._definitions except Exception as exc: LOG.error(f"Failed to load task definitions from database: {exc}") @@ -97,10 +91,10 @@ class GenericTaskHandler(BaseHandler): try: for matcher in definition.match_url: - pattern_str = None + pattern_str: str | None = None if matcher.startswith("/") and matcher.endswith("/") and len(matcher) > 2: - pattern_str: str = matcher[1:-1] + pattern_str = matcher[1:-1] else: pattern_str = fnmatch.translate(matcher) @@ -274,53 +268,37 @@ class GenericTaskHandler(BaseHandler): """ headers: dict[str, str] = {**definition.definition.request.headers} - client_options: dict[str, Any] = { - "headers": { - "User-Agent": random_user_agent(), - } - } - - try: - from httpx_curl_cffi import AsyncCurlTransport, CurlOpt - - client_options["transport"] = AsyncCurlTransport( - impersonate="chrome", - default_headers=True, - curl_options={CurlOpt.FRESH_CONNECT: True}, - ) - client_options["headers"].pop("User-Agent", None) - except Exception: - pass - - if headers: - client_options["headers"].update(headers) - - if proxy := ytdlp_opts.get("proxy"): - client_options["proxy"] = proxy + use_curl = resolve_curl_transport() + request_headers = build_request_headers( + base_headers=headers, + user_agent=Globals.get_random_agent(), + use_curl=use_curl, + ) timeout_value: float | Any = definition.definition.request.timeout or ytdlp_opts.get("socket_timeout", 120) - async with async_client(**client_options) as client: - response: httpx.Response = await client.request( - method=definition.definition.request.method.upper(), - url=url, - params=definition.definition.request.params or None, - data=definition.definition.request.data, - json=definition.definition.request.json_data, - timeout=timeout_value, - ) - response.raise_for_status() + client = get_async_client(proxy=ytdlp_opts.get("proxy"), use_curl=use_curl) + response: httpx.Response = await client.request( + method=definition.definition.request.method.upper(), + url=url, + params=definition.definition.request.params or None, + data=definition.definition.request.data, + json=definition.definition.request.json_data, + timeout=timeout_value, + headers=request_headers, + ) + response.raise_for_status() - if "json" == definition.definition.response.type: - try: - json_data: dict[str, Any] = response.json() - except Exception as exc: - LOG.error(f"Failed to decode JSON response from '{url}': {exc}") - return response.text, None + if "json" == definition.definition.response.type: + try: + json_data: dict[str, Any] = response.json() + except Exception as exc: + LOG.error(f"Failed to decode JSON response from '{url}': {exc}") + return response.text, None - return response.text, json_data + return response.text, json_data - return response.text, None + return response.text, None @staticmethod async def _fetch_with_selenium( diff --git a/app/features/tasks/router.py b/app/features/tasks/router.py index 9b94c095..d7ffc43b 100644 --- a/app/features/tasks/router.py +++ b/app/features/tasks/router.py @@ -327,7 +327,8 @@ async def task_metadata(request: Request, repo: TasksRepository, config: Config, The response object. """ - task_id = request.match_info.get("id") + if not (task_id := request.match_info.get("id")): + return web.json_response(data={"error": "No task id."}, status=web.HTTPBadRequest.status_code) try: if not (model := await repo.get(int(task_id))): @@ -335,7 +336,6 @@ async def task_metadata(request: Request, repo: TasksRepository, config: Config, data={"error": f"Task '{task_id}' does not exist."}, status=web.HTTPNotFound.status_code ) - # Convert to extended Task with handler methods task = ExtendedTask.model_validate(model) (save_path, _) = get_file(config.download_path, task.folder) @@ -348,6 +348,11 @@ async def task_metadata(request: Request, repo: TasksRepository, config: Config, metadata, status, message = await task.fetch_metadata() if not status: return web.json_response(data={"error": message}, status=web.HTTPBadRequest.status_code) + if not isinstance(metadata, dict): + return web.json_response( + data={"error": "Failed to get metadata."}, + status=web.HTTPBadRequest.status_code, + ) if not task.folder: try: @@ -395,7 +400,7 @@ async def task_metadata(request: Request, repo: TasksRepository, config: Config, from app.yt_dlp_plugins.postprocessor.nfo_maker import NFOMakerPP - title: str = sanitize_filename(info.get("title")) + title: str = sanitize_filename(str(info.get("title") or "")) info_file: Path = save_path / f"{title} [{info.get('id')}].info.json" info_file.write_text(encoder.encode(metadata), encoding="utf-8") info["json_file"] = str(info_file.relative_to(config.download_path)) @@ -406,9 +411,7 @@ async def task_metadata(request: Request, repo: TasksRepository, config: Config, xml_content = "\n" xml_content += f" {NFOMakerPP._escape_text(info.get('title'))}\n" if info.get("description"): - xml_content += ( - f" {NFOMakerPP._escape_text(NFOMakerPP._clean_description(info.get('description')))}\n" - ) + xml_content += f" {NFOMakerPP._escape_text(NFOMakerPP._clean_description(str(info.get('description') or '')))}\n" if info.get("id"): xml_content += f" {NFOMakerPP._escape_text(info.get('id'))}\n" if info.get("id_type") and info.get("id"): @@ -425,52 +428,53 @@ async def task_metadata(request: Request, repo: TasksRepository, config: Config, xml_file.write_text(xml_content, encoding="utf-8") try: - from yt_dlp.utils.networking import random_user_agent - - from app.library.httpx_client import async_client + from app.library.httpx_client import ( + Globals, + build_request_headers, + get_async_client, + resolve_curl_transport, + ) ytdlp_args: dict = task.get_ytdlp_opts().get_all() - opts: dict[str, Any] = { - "headers": { - "User-Agent": request.headers.get("User-Agent", ytdlp_args.get("user_agent", random_user_agent())), - }, - } - if proxy := ytdlp_args.get("proxy"): - opts["proxy"] = proxy + use_curl = resolve_curl_transport() + request_headers = build_request_headers( + user_agent=request.headers.get("User-Agent", ytdlp_args.get("user_agent", Globals.get_random_agent())), + use_curl=use_curl, + ) - try: - from httpx_curl_cffi import AsyncCurlTransport - - opts["transport"] = AsyncCurlTransport( - impersonate="chrome", - default_headers=True, - ) - opts.pop("headers", None) - except Exception: - pass - - async with async_client(**opts) as client: - for key in info.get("thumbnails", {}): - try: - url = info["thumbnails"][key] - LOG.info(f"Fetching thumbnail '{key}' from '{url}'") - if not url: - continue - - try: - validate_url(url, allow_internal=config.allow_internal_urls) - except ValueError: - LOG.warning(f"Invalid thumbnail url '{url}'") - continue - - resp = await client.request(method="GET", url=url, follow_redirects=True) - - img_file = save_path / f"{key}.jpg" - img_file.write_bytes(resp.content) - info["thumbnails"][key] = str(img_file.relative_to(config.download_path)) - except Exception as e: - LOG.warning(f"Failed to fetch thumbnail '{key}' from '{url}'. '{e!s}'") + client = get_async_client(proxy=ytdlp_args.get("proxy"), use_curl=use_curl) + thumbnails = info.get("thumbnails", {}) + if not isinstance(thumbnails, dict): + thumbnails = {} + info["thumbnails"] = thumbnails + for key in thumbnails: + url: str | None = None + try: + url = thumbnails.get(key) + LOG.info(f"Fetching thumbnail '{key}' from '{url}'") + if not url: continue + + try: + validate_url(url, allow_internal=config.allow_internal_urls) + except ValueError: + LOG.warning(f"Invalid thumbnail url '{url}'") + continue + + resp = await client.request( + method="GET", + url=url, + follow_redirects=True, + headers=request_headers, + ) + + img_file = save_path / f"{key}.jpg" + img_file.write_bytes(resp.content) + thumbnails[key] = str(img_file.relative_to(config.download_path)) + except Exception as e: + url_log = url or "unknown" + LOG.warning(f"Failed to fetch thumbnail '{key}' from '{url_log}'. '{e!s}'") + continue except Exception as e: LOG.warning(f"Failed to fetch thumbnails. '{e!s}'") diff --git a/app/library/UpdateChecker.py b/app/library/UpdateChecker.py index 72d3cdac..4d6eae2d 100644 --- a/app/library/UpdateChecker.py +++ b/app/library/UpdateChecker.py @@ -8,7 +8,7 @@ from aiohttp import web from .cache import Cache from .config import Config from .Events import EventBus, Events -from .httpx_client import async_client +from .httpx_client import get_async_client from .Scheduler import Scheduler from .Singleton import Singleton from .version import APP_VERSION @@ -147,36 +147,37 @@ class UpdateChecker(metaclass=Singleton): try: LOG.info(f"Checking for {name} updates...") - async with async_client(timeout=10.0) as client: - response = await client.get( - api_url, - headers={"Accept": "application/vnd.github+json"}, - ) + client = get_async_client(use_curl=False) + response = await client.get( + api_url, + headers={"Accept": "application/vnd.github+json"}, + timeout=10.0, + ) - if 200 != response.status_code: - LOG.warning(f"Failed to check for {name} updates: HTTP {response.status_code}") - return ("error", None) + if 200 != response.status_code: + LOG.warning(f"Failed to check for {name} updates: HTTP {response.status_code}") + return ("error", None) - data: dict[str, Any] = response.json() + data: dict[str, Any] = response.json() - latest_tag: str = data.get("tag_name", "") - if not latest_tag: - LOG.warning(f"No tag_name found in {name} GitHub release data.") - return ("error", None) + latest_tag: str = data.get("tag_name", "") + if not latest_tag: + LOG.warning(f"No tag_name found in {name} GitHub release data.") + return ("error", None) - compare_current: str = current_version.lstrip("v") if strip_v_prefix else current_version - compare_latest: str = latest_tag.lstrip("v") if strip_v_prefix else latest_tag + compare_current: str = current_version.lstrip("v") if strip_v_prefix else current_version + compare_latest: str = latest_tag.lstrip("v") if strip_v_prefix else latest_tag - if self._compare_versions(compare_current, compare_latest): - LOG.warning(f"{name} update available: {current_version} -> {latest_tag}") - result: tuple[str, str] = ("update_available", latest_tag) - await self._cache.aset(cache_key, result, self.CACHE_DURATION) - return result - - LOG.info(f"No {name} updates available.") - result: tuple[str, None] = ("up_to_date", None) + if self._compare_versions(compare_current, compare_latest): + LOG.warning(f"{name} update available: {current_version} -> {latest_tag}") + result = ("update_available", latest_tag) await self._cache.aset(cache_key, result, self.CACHE_DURATION) return result + + LOG.info(f"No {name} updates available.") + result = ("up_to_date", None) + await self._cache.aset(cache_key, result, self.CACHE_DURATION) + return result except Exception as e: LOG.exception(e) LOG.error(f"Error checking for {name} updates: {e!s}") diff --git a/app/library/httpx_client.py b/app/library/httpx_client.py index 78fd3a1b..a0c1a43d 100644 --- a/app/library/httpx_client.py +++ b/app/library/httpx_client.py @@ -1,17 +1,61 @@ from __future__ import annotations +import asyncio +import functools import logging -from typing import Any +import threading +from dataclasses import dataclass +from typing import Any, Literal, cast, overload import httpx from .cf_solver_shared import is_cf_challenge, solver -__all__: list[str] = ["async_client", "sync_client"] +__all__: list[str] = [ + "async_client", + "build_request_headers", + "close_shared_clients", + "get_async_client", + "get_sync_client", + "resolve_curl_transport", + "sync_client", +] LOG: logging.Logger = logging.getLogger("httpx_cf") +class Globals: + random_agent: str | None = None + + SHARED_ASYNC_CLIENTS: dict[_AsyncClientKey, httpx.AsyncClient] = {} + SHARED_SYNC_CLIENTS: dict[_SyncClientKey, httpx.Client] = {} + SHARED_CLIENT_LOCK = threading.Lock() + + def get_random_agent() -> str: + if Globals.random_agent: + return Globals.random_agent + + from yt_dlp.utils.networking import random_user_agent + + Globals.random_agent = random_user_agent() + return Globals.random_agent + + +@dataclass(frozen=True) +class _AsyncClientKey: + enable_cf: bool + proxy: str | None + use_curl: bool + curl_impersonate: str + curl_default_headers: bool + + +@dataclass(frozen=True) +class _SyncClientKey: + enable_cf: bool + proxy: str | None + + def _parse_cookie_header(cookie_header: str | None) -> dict[str, str]: cookies: dict[str, str] = {} @@ -37,17 +81,95 @@ def _merge_cookies(existing_header: str | None, new_cookies: list[dict[str, str] return merged +def _normalize_proxy(proxy: str | dict[str, str] | None) -> str | None: + if proxy is None: + return None + + if isinstance(proxy, str): + cleaned: str = proxy.strip() + return cleaned or None + + if isinstance(proxy, dict): + return "|".join(f"{key}={value}" for key, value in sorted(proxy.items())) + + return str(proxy) + + +@functools.lru_cache(maxsize=1) +def _curl_available() -> bool: + try: + import httpx_curl_cffi # noqa: F401 + + return True + except Exception: + return False + + +def resolve_curl_transport(use_curl: bool = True) -> bool: + return use_curl and _curl_available() + + +def _build_async_curl_transport( + use_curl: bool, + curl_impersonate: str, + curl_default_headers: bool, +) -> httpx.AsyncBaseTransport | None: + if not resolve_curl_transport(use_curl): + return None + + from httpx_curl_cffi import AsyncCurlTransport + + return AsyncCurlTransport( + impersonate=curl_impersonate, + default_headers=curl_default_headers, + ) + + +@overload +def _get_transport( + enable_cf: bool, + is_async: Literal[True], + transport: httpx.AsyncBaseTransport | None, +) -> httpx.AsyncBaseTransport: ... + + +@overload +def _get_transport( + enable_cf: bool, + is_async: Literal[False], + transport: httpx.BaseTransport | None, +) -> httpx.BaseTransport: ... + + def _get_transport( enable_cf: bool, is_async: bool, - transport: httpx.AsyncBaseTransport | None, -) -> httpx.AsyncBaseTransport: + transport: httpx.AsyncBaseTransport | httpx.BaseTransport | None, +) -> httpx.AsyncBaseTransport | httpx.BaseTransport: if enable_cf: - return CFAsyncTransport(base=transport) if is_async else CFTransport(base=transport) + if is_async: + async_transport = cast("httpx.AsyncBaseTransport | None", transport) + return CFAsyncTransport(base=async_transport) + + sync_transport = cast("httpx.BaseTransport | None", transport) + return CFTransport(base=sync_transport) return transport or (httpx.AsyncHTTPTransport() if is_async else httpx.HTTPTransport()) +def build_request_headers( + base_headers: dict[str, str] | None = None, + user_agent: str | None = None, + use_curl: bool = True, +) -> dict[str, str]: + headers: dict[str, str] = base_headers.copy() if isinstance(base_headers, dict) else {} + + if user_agent and not use_curl: + headers.setdefault("User-Agent", user_agent) + + return headers + + class CFAsyncTransport(httpx.AsyncBaseTransport): def __init__(self, base: httpx.AsyncBaseTransport | None = None): self.base: httpx.AsyncBaseTransport | httpx.AsyncHTTPTransport = base or httpx.AsyncHTTPTransport() @@ -82,7 +204,9 @@ class CFAsyncTransport(httpx.AsyncBaseTransport): return await self.base.handle_async_request(request) def close(self) -> None: - self.base.close() + close_fn = getattr(self.base, "close", None) + if callable(close_fn): + close_fn() class CFTransport(httpx.BaseTransport): @@ -134,8 +258,9 @@ def async_client(enable_cf: bool = True, **kwargs: Any) -> httpx.AsyncClient: httpx.AsyncClient: The configured httpx.AsyncClient instance. """ - transport = kwargs.pop("transport", None) - return httpx.AsyncClient(transport=_get_transport(enable_cf, is_async=True, transport=transport), **kwargs) + transport = cast("httpx.AsyncBaseTransport | None", kwargs.pop("transport", None)) + async_transport = _get_transport(enable_cf, is_async=True, transport=transport) + return httpx.AsyncClient(transport=async_transport, **kwargs) def sync_client(enable_cf: bool = True, **kwargs: Any) -> httpx.Client: @@ -150,5 +275,81 @@ def sync_client(enable_cf: bool = True, **kwargs: Any) -> httpx.Client: httpx.Client: The configured httpx.Client instance. """ - transport = kwargs.pop("transport", None) - return httpx.Client(transport=_get_transport(enable_cf, is_async=False, transport=transport), **kwargs) + transport = cast("httpx.BaseTransport | None", kwargs.pop("transport", None)) + sync_transport = _get_transport(enable_cf, is_async=False, transport=transport) + return httpx.Client(transport=sync_transport, **kwargs) + + +def get_async_client( + enable_cf: bool = True, + proxy: str | dict[str, str] | None = None, + use_curl: bool = True, + curl_impersonate: str = "chrome", + curl_default_headers: bool = True, +) -> httpx.AsyncClient: + proxy_key = _normalize_proxy(proxy) + use_curl = resolve_curl_transport(use_curl) + key = _AsyncClientKey( + enable_cf=enable_cf, + proxy=proxy_key, + use_curl=use_curl, + curl_impersonate=curl_impersonate, + curl_default_headers=curl_default_headers, + ) + + with Globals.SHARED_CLIENT_LOCK: + if key in Globals.SHARED_ASYNC_CLIENTS: + return Globals.SHARED_ASYNC_CLIENTS[key] + + transport = _build_async_curl_transport( + use_curl=use_curl, + curl_impersonate=curl_impersonate, + curl_default_headers=curl_default_headers, + ) + client = httpx.AsyncClient( + transport=cast( + "httpx.AsyncBaseTransport", + _get_transport(enable_cf, is_async=True, transport=transport), + ), + proxy=cast("Any", proxy), + ) + Globals.SHARED_ASYNC_CLIENTS[key] = client + return client + + +def get_sync_client( + enable_cf: bool = True, + proxy: str | dict[str, str] | None = None, +) -> httpx.Client: + proxy_key = _normalize_proxy(proxy) + key = _SyncClientKey(enable_cf=enable_cf, proxy=proxy_key) + + with Globals.SHARED_CLIENT_LOCK: + if key in Globals.SHARED_SYNC_CLIENTS: + return Globals.SHARED_SYNC_CLIENTS[key] + + client = httpx.Client( + transport=cast( + "httpx.BaseTransport", + _get_transport(enable_cf, is_async=False, transport=None), + ), + proxy=cast("Any", proxy), + ) + Globals.SHARED_SYNC_CLIENTS[key] = client + return client + + +async def close_shared_clients(_: Any | None = None) -> None: + async_clients = list(Globals.SHARED_ASYNC_CLIENTS.values()) + sync_clients = list(Globals.SHARED_SYNC_CLIENTS.values()) + Globals.SHARED_ASYNC_CLIENTS.clear() + Globals.SHARED_SYNC_CLIENTS.clear() + + for client in sync_clients: + try: + client.close() + except Exception: + pass + + if async_clients: + await asyncio.gather(*(client.aclose() for client in async_clients), return_exceptions=True) diff --git a/app/main.py b/app/main.py index 6dfe8dde..7ac1c9ee 100644 --- a/app/main.py +++ b/app/main.py @@ -27,6 +27,7 @@ from app.library.downloads import DownloadQueue from app.library.Events import EventBus, Events from app.library.HttpAPI import HttpAPI from app.library.HttpSocket import HttpSocket +from app.library.httpx_client import close_shared_clients from app.library.Scheduler import Scheduler from app.library.Services import Services from app.library.sqlite_store import SqliteStore @@ -124,6 +125,7 @@ class Main: get_task_definitions_repo().attach(self._app) DownloadQueue.get_instance().attach(self._app) UpdateChecker.get_instance().attach(self._app) + self._app.on_shutdown.append(close_shared_clients) EventBus.get_instance().emit( Events.LOADED, diff --git a/app/routes/api/docs.py b/app/routes/api/docs.py index f377f515..cbba2010 100644 --- a/app/routes/api/docs.py +++ b/app/routes/api/docs.py @@ -1,14 +1,13 @@ import logging import time from datetime import UTC, datetime -from typing import Any from aiohttp import web from aiohttp.web import Request, Response -from yt_dlp.utils.networking import random_user_agent from app.library.cache import Cache from app.library.config import Config +from app.library.httpx_client import Globals, build_request_headers, get_async_client, resolve_curl_transport from app.library.router import add_route, route from app.library.YTDLPOpts import YTDLPOpts @@ -65,48 +64,38 @@ async def get_doc(request: Request, config: Config, cache: Cache) -> Response: try: ytdlp_args: dict = YTDLPOpts.get_instance().preset(name=config.default_preset).get_all() - opts: dict[str, Any] = { + use_curl = resolve_curl_transport() + request_headers = build_request_headers( + user_agent=request.headers.get("User-Agent", ytdlp_args.get("user_agent", Globals.get_random_agent())), + use_curl=use_curl, + ) + proxy = ytdlp_args.get("proxy") + + client = get_async_client(proxy=proxy, use_curl=use_curl) + LOG.debug(f"Fetching doc from '{url}'.") + response = await client.request( + method="GET", + url=url, + follow_redirects=True, + headers=request_headers, + ) + dct = { + "body": response.content, "headers": { - "User-Agent": request.headers.get("User-Agent", ytdlp_args.get("user_agent", random_user_agent())), + "Content-Type": EXT_TO_MIME.get(file[file.rfind(".") :], "text/plain"), + "Pragma": "public", + "Access-Control-Allow-Origin": "*", + "Cache-Control": f"public, max-age={time.time() + 3600}", + "Expires": time.strftime( + "%a, %d %b %Y %H:%M:%S GMT", + datetime.fromtimestamp(time.time() + 3600, tz=UTC).timetuple(), + ), }, } - if proxy := ytdlp_args.get("proxy"): - opts["proxy"] = proxy - try: - from httpx_curl_cffi import AsyncCurlTransport, CurlOpt + cache.set(cache_key, dct, ttl=3600) - opts["transport"] = AsyncCurlTransport( - impersonate="chrome", - default_headers=True, - curl_options={CurlOpt.FRESH_CONNECT: True}, - ) - opts.pop("headers", None) - except Exception: - pass - - from app.library.httpx_client import async_client - - async with async_client(**opts) as client: - LOG.debug(f"Fetching doc from '{url}'.") - response = await client.request(method="GET", url=url, follow_redirects=True) - dct = { - "body": response.content, - "headers": { - "Content-Type": EXT_TO_MIME.get(file[file.rfind(".") :], "text/plain"), - "Pragma": "public", - "Access-Control-Allow-Origin": "*", - "Cache-Control": f"public, max-age={time.time() + 3600}", - "Expires": time.strftime( - "%a, %d %b %Y %H:%M:%S GMT", - datetime.fromtimestamp(time.time() + 3600, tz=UTC).timetuple(), - ), - }, - } - - cache.set(cache_key, dct, ttl=3600) - - return web.Response(**dct) + return web.Response(**dct) except Exception as e: LOG.error(f"Failed to request doc from '{url}'.'. '{e!s}'.") return web.json_response(data={"error": "Failed to get doc."}, status=web.HTTPInternalServerError.status_code) diff --git a/app/routes/api/images.py b/app/routes/api/images.py index 8cfcc74a..9d1b80bb 100644 --- a/app/routes/api/images.py +++ b/app/routes/api/images.py @@ -7,11 +7,11 @@ from urllib.parse import urlparse from aiohttp import web from aiohttp.web import Request, Response -from yt_dlp.utils.networking import random_user_agent from app.library.ag_utils import ag from app.library.cache import Cache from app.library.config import Config +from app.library.httpx_client import Globals, build_request_headers, get_async_client, resolve_curl_transport from app.library.router import route from app.library.Utils import validate_url from app.library.YTDLPOpts import YTDLPOpts @@ -45,52 +45,39 @@ async def get_thumbnail(request: Request, config: Config) -> Response: try: ytdlp_args: dict = YTDLPOpts.get_instance().preset(name=config.default_preset).get_all() - opts: dict[str, Any] = { - "headers": { - "User-Agent": request.headers.get("User-Agent", ytdlp_args.get("user_agent", random_user_agent())), + use_curl = resolve_curl_transport() + request_headers = build_request_headers( + user_agent=request.headers.get("User-Agent", ytdlp_args.get("user_agent", Globals.get_random_agent())), + use_curl=use_curl, + ) + proxy = ytdlp_args.get("proxy") + + client = get_async_client(proxy=proxy, use_curl=use_curl) + LOG.debug(f"Fetching thumbnail from '{url}'.") + response = await client.request( + method="GET", + url=url, + follow_redirects=True, + headers=request_headers, + ) + + if response.status_code != web.HTTPOk.status_code: + LOG.error(f"Failed to fetch thumbnail from '{url}'. Status code: {response.status_code}.") + return web.json_response(data={"error": "failed to retrieve the thumbnail."}, status=response.status_code) + + return web.Response( + body=response.content, + headers={ + "Content-Type": response.headers.get("Content-Type"), + "Pragma": "public", + "Access-Control-Allow-Origin": "*", + "Cache-Control": f"public, max-age={time.time() + 31536000}", + "Expires": time.strftime( + "%a, %d %b %Y %H:%M:%S GMT", + datetime.fromtimestamp(time.time() + 31536000, tz=UTC).timetuple(), + ), }, - } - - if proxy := ytdlp_args.get("proxy"): - opts["proxy"] = proxy - - try: - from httpx_curl_cffi import AsyncCurlTransport, CurlOpt - - opts["transport"] = AsyncCurlTransport( - impersonate="chrome", - default_headers=True, - curl_options={CurlOpt.FRESH_CONNECT: True}, - ) - opts.pop("headers", None) - except Exception: - pass - - from app.library.httpx_client import async_client - - async with async_client(**opts) as client: - LOG.debug(f"Fetching thumbnail from '{url}'.") - response = await client.request(method="GET", url=url, follow_redirects=True) - - if response.status_code != web.HTTPOk.status_code: - LOG.error(f"Failed to fetch thumbnail from '{url}'. Status code: {response.status_code}.") - return web.json_response( - data={"error": "failed to retrieve the thumbnail."}, status=response.status_code - ) - - return web.Response( - body=response.content, - headers={ - "Content-Type": response.headers.get("Content-Type"), - "Pragma": "public", - "Access-Control-Allow-Origin": "*", - "Cache-Control": f"public, max-age={time.time() + 31536000}", - "Expires": time.strftime( - "%a, %d %b %Y %H:%M:%S GMT", - datetime.fromtimestamp(time.time() + 31536000, tz=UTC).timetuple(), - ), - }, - ) + ) except Exception as e: LOG.error(f"Error fetching thumbnail from '{url}'. '{e}'.") return web.json_response( @@ -126,95 +113,100 @@ async def get_background(request: Request, config: Config, cache: Cache) -> Resp CACHE_KEY = "random_background" if cache.has(CACHE_KEY) and not request.query.get("force", False): - data = await cache.aget(CACHE_KEY) - return web.Response( - body=data.get("content"), - headers={ - "X-Cache": "HIT", - "X-Cache-TTL": str(await cache.attl(CACHE_KEY)), - "X-Image-Via": data.get("backend"), - **data.get("headers"), - }, - ) - - ytdlp_args: dict = YTDLPOpts.get_instance().preset(name=config.default_preset).get_all() - opts: dict[str, Any] = { - "headers": { - "User-Agent": request.headers.get("User-Agent", ytdlp_args.get("user_agent", random_user_agent())), - }, - } - - if proxy := ytdlp_args.get("proxy"): - opts["proxy"] = proxy - - try: - from httpx_curl_cffi import AsyncCurlTransport, CurlOpt - - opts["transport"] = AsyncCurlTransport( - impersonate="chrome", - default_headers=True, - curl_options={CurlOpt.FRESH_CONNECT: True}, - ) - opts.pop("headers", None) - except Exception: - pass - - from app.library.httpx_client import async_client - - async with async_client(**opts) as client: - if backend.startswith("https://www.bing.com/HPImageArchive.aspx"): - if not cache.has(CACHE_KEY_BING): - response = await client.request(method="GET", url=backend) - if response.status_code != web.HTTPOk.status_code: - return web.json_response( - data={"error": "failed to retrieve the random background image."}, - status=web.HTTPInternalServerError.status_code, - ) - - img_url: str | None = ag(response.json(), "images.0.url") - if not img_url: - return web.json_response( - data={"error": "failed to retrieve the random background image."}, - status=web.HTTPInternalServerError.status_code, - ) - - backend = f"https://www.bing.com{img_url}" - await cache.aset(key=CACHE_KEY_BING, value=backend, ttl=3600 * 24) - else: - backend: str = await cache.aget(CACHE_KEY_BING) - - LOG.debug(f"Requesting random picture from '{backend!s}'.") - - response = await client.request(method="GET", url=backend, follow_redirects=True) - - if response.status_code != web.HTTPOk.status_code: - return web.json_response( - data={"error": "failed to retrieve the random background image."}, - status=web.HTTPInternalServerError.status_code, + cached_data = await cache.aget(CACHE_KEY) + if isinstance(cached_data, dict): + cached_headers = cached_data.get("headers") + if not isinstance(cached_headers, dict): + cached_headers = {} + cached_headers = {str(key): str(value) for key, value in cached_headers.items()} + cached_backend = cached_data.get("backend") + if not isinstance(cached_backend, str): + cached_backend = "" if cached_backend is None else str(cached_backend) + return web.Response( + body=cached_data.get("content"), + headers={ + "X-Cache": "HIT", + "X-Cache-TTL": str(await cache.attl(CACHE_KEY)), + "X-Image-Via": cached_backend, + **cached_headers, + }, ) - data: dict[str, Any] = { - "content": response.content, - "backend": urlparse(backend).netloc, - "headers": { - "Content-Type": response.headers.get("Content-Type", "image/jpeg"), - "Content-Length": str(len(response.content)), - }, - } + ytdlp_args: dict = YTDLPOpts.get_instance().preset(name=config.default_preset).get_all() + use_curl = resolve_curl_transport() + request_headers = build_request_headers( + user_agent=request.headers.get("User-Agent", ytdlp_args.get("user_agent", Globals.get_random_agent())), + use_curl=use_curl, + ) + proxy = ytdlp_args.get("proxy") - await cache.aset(key=CACHE_KEY, value=data, ttl=3600) + client = get_async_client(proxy=proxy, use_curl=use_curl) + if backend.startswith("https://www.bing.com/HPImageArchive.aspx"): + if not cache.has(CACHE_KEY_BING): + response = await client.request(method="GET", url=backend, headers=request_headers) + if response.status_code != web.HTTPOk.status_code: + return web.json_response( + data={"error": "failed to retrieve the random background image."}, + status=web.HTTPInternalServerError.status_code, + ) - LOG.debug(f"Random background image from '{backend!s}' cached.") + img_url: str | None = ag(response.json(), "images.0.url") + if not img_url: + return web.json_response( + data={"error": "failed to retrieve the random background image."}, + status=web.HTTPInternalServerError.status_code, + ) - return web.Response( - body=data.get("content"), - headers={ - "X-Cache": "MISS", - "X-Cache-TTL": "3600", - "X-Image-Via": data.get("backend"), - **data.get("headers"), - }, + backend = f"https://www.bing.com{img_url}" + await cache.aset(key=CACHE_KEY_BING, value=backend, ttl=3600 * 24) + else: + backend = await cache.aget(CACHE_KEY_BING) + + if not isinstance(backend, str) or not backend: + return web.json_response( + data={"error": "failed to retrieve the random background image."}, + status=web.HTTPInternalServerError.status_code, ) + + LOG.debug(f"Requesting random picture from '{backend!s}'.") + + response = await client.request( + method="GET", + url=backend, + follow_redirects=True, + headers=request_headers, + ) + + if response.status_code != web.HTTPOk.status_code: + return web.json_response( + data={"error": "failed to retrieve the random background image."}, + status=web.HTTPInternalServerError.status_code, + ) + + data: dict[str, Any] = { + "content": response.content, + "backend": urlparse(backend).netloc, + "headers": { + "Content-Type": response.headers.get("Content-Type", "image/jpeg"), + "Content-Length": str(len(response.content)), + }, + } + response_headers: dict[str, str] = data["headers"] + image_via = str(data.get("backend") or "") + + await cache.aset(key=CACHE_KEY, value=data, ttl=3600) + + LOG.debug(f"Random background image from '{backend!s}' cached.") + + return web.Response( + body=data.get("content"), + headers={ + "X-Cache": "MISS", + "X-Cache-TTL": "3600", + "X-Image-Via": image_via, + **response_headers, + }, + ) except Exception as e: LOG.error(f"Failed to request random background image from '{backend!s}'.'. '{e!s}'.") return web.json_response( diff --git a/app/tests/test_update_checker.py b/app/tests/test_update_checker.py index f876b865..33cf3c79 100644 --- a/app/tests/test_update_checker.py +++ b/app/tests/test_update_checker.py @@ -153,7 +153,7 @@ class TestUpdateChecker: assert "" == config.yt_new_version, "Should not update yt_new_version when disabled" @pytest.mark.asyncio - @patch("app.library.UpdateChecker.async_client") + @patch("app.library.UpdateChecker.get_async_client") async def test_check_for_updates_finds_newer_version(self, mock_client): """Test that check_for_updates detects when a newer version is available.""" from app.library.config import Config @@ -173,9 +173,9 @@ class TestUpdateChecker: mock_ytdlp_response.json.return_value = {"tag_name": "9999.12.31"} mock_get = AsyncMock(side_effect=[mock_app_response, mock_ytdlp_response]) - mock_context = AsyncMock() - mock_context.__aenter__.return_value.get = mock_get - mock_client.return_value = mock_context + mock_http = MagicMock() + mock_http.get = mock_get + mock_client.return_value = mock_http with patch("app.library.UpdateChecker.APP_VERSION", "v1.0.0"): checker = UpdateChecker.get_instance(config=config) @@ -190,7 +190,7 @@ class TestUpdateChecker: assert checker._job_id is None, "Should stop scheduled task after finding app update" @pytest.mark.asyncio - @patch("app.library.UpdateChecker.async_client") + @patch("app.library.UpdateChecker.get_async_client") async def test_check_for_updates_no_update_available(self, mock_client): """Test that check_for_updates correctly handles when no update is available.""" from app.library.config import Config @@ -210,9 +210,9 @@ class TestUpdateChecker: mock_ytdlp_response.json.return_value = {"tag_name": "2020.01.01"} mock_get = AsyncMock(side_effect=[mock_app_response, mock_ytdlp_response]) - mock_context = AsyncMock() - mock_context.__aenter__.return_value.get = mock_get - mock_client.return_value = mock_context + mock_http = MagicMock() + mock_http.get = mock_get + mock_client.return_value = mock_http checker = UpdateChecker.get_instance(config=config) checker._job_id = "test-job" @@ -228,7 +228,7 @@ class TestUpdateChecker: assert "test-job" == checker._job_id, "Should keep scheduled task running" @pytest.mark.asyncio - @patch("app.library.UpdateChecker.async_client") + @patch("app.library.UpdateChecker.get_async_client") async def test_check_for_updates_handles_http_error(self, mock_client): """Test that check_for_updates handles HTTP errors gracefully.""" from app.library.config import Config @@ -242,9 +242,9 @@ class TestUpdateChecker: mock_response = MagicMock() mock_response.status_code = 404 - mock_context = AsyncMock() - mock_context.__aenter__.return_value.get = AsyncMock(return_value=mock_response) - mock_client.return_value = mock_context + mock_http = MagicMock() + mock_http.get = AsyncMock(return_value=mock_response) + mock_client.return_value = mock_http checker = UpdateChecker.get_instance(config=config) @@ -258,7 +258,7 @@ class TestUpdateChecker: assert "" == config.yt_new_version, "Should not set yt_new_version on HTTP error" @pytest.mark.asyncio - @patch("app.library.UpdateChecker.async_client") + @patch("app.library.UpdateChecker.get_async_client") async def test_check_for_updates_handles_exception(self, mock_client): """Test that check_for_updates handles exceptions gracefully.""" from app.library.config import Config @@ -269,9 +269,9 @@ class TestUpdateChecker: config.new_version = "" config.yt_new_version = "" - mock_context = AsyncMock() - mock_context.__aenter__.return_value.get = AsyncMock(side_effect=Exception("Network error")) - mock_client.return_value = mock_context + mock_http = MagicMock() + mock_http.get = AsyncMock(side_effect=Exception("Network error")) + mock_client.return_value = mock_http checker = UpdateChecker.get_instance(config=config) @@ -337,10 +337,10 @@ class TestUpdateChecker: mock_ytdlp_response.json.return_value = {"tag_name": "2026.01.01"} mock_get = AsyncMock(side_effect=[mock_app_response, mock_ytdlp_response]) - mock_context = AsyncMock() - mock_context.__aenter__.return_value.get = mock_get + mock_http = MagicMock() + mock_http.get = mock_get - with patch("app.library.UpdateChecker.async_client", return_value=mock_context): + with patch("app.library.UpdateChecker.get_async_client", return_value=mock_http): with patch("app.library.UpdateChecker.APP_VERSION", "v1.0.0"): checker = UpdateChecker.get_instance(config=config) await checker.check_for_updates() @@ -383,7 +383,7 @@ class TestUpdateChecker: loop.close() @pytest.mark.asyncio - @patch("app.library.UpdateChecker.async_client") + @patch("app.library.UpdateChecker.get_async_client") async def test_check_ytdlp_version_finds_newer_version(self, mock_client): """Test that yt-dlp check detects when a newer version is available.""" from app.library.config import Config @@ -397,9 +397,9 @@ class TestUpdateChecker: mock_response.status_code = 200 mock_response.json.return_value = {"tag_name": "9999.12.31"} - mock_context = AsyncMock() - mock_context.__aenter__.return_value.get = AsyncMock(return_value=mock_response) - mock_client.return_value = mock_context + mock_http = MagicMock() + mock_http.get = AsyncMock(return_value=mock_response) + mock_client.return_value = mock_http checker = UpdateChecker.get_instance(config=config) status, new_version = await checker._check_ytdlp_version() @@ -409,7 +409,7 @@ class TestUpdateChecker: assert "9999.12.31" == config.yt_new_version, "Should store new yt-dlp version tag" @pytest.mark.asyncio - @patch("app.library.UpdateChecker.async_client") + @patch("app.library.UpdateChecker.get_async_client") async def test_check_ytdlp_version_no_update_available(self, mock_client): """Test that yt-dlp check correctly handles when no update is available.""" from app.library.config import Config @@ -423,9 +423,9 @@ class TestUpdateChecker: mock_response.status_code = 200 mock_response.json.return_value = {"tag_name": "2020.01.01"} - mock_context = AsyncMock() - mock_context.__aenter__.return_value.get = AsyncMock(return_value=mock_response) - mock_client.return_value = mock_context + mock_http = MagicMock() + mock_http.get = AsyncMock(return_value=mock_response) + mock_client.return_value = mock_http checker = UpdateChecker.get_instance(config=config) status, new_version = await checker._check_ytdlp_version() @@ -435,7 +435,7 @@ class TestUpdateChecker: assert "" == config.yt_new_version, "Should clear yt_new_version when no update available" @pytest.mark.asyncio - @patch("app.library.UpdateChecker.async_client") + @patch("app.library.UpdateChecker.get_async_client") async def test_check_ytdlp_version_handles_http_error(self, mock_client): """Test that yt-dlp check handles HTTP errors gracefully.""" from app.library.config import Config @@ -448,9 +448,9 @@ class TestUpdateChecker: mock_response = MagicMock() mock_response.status_code = 500 - mock_context = AsyncMock() - mock_context.__aenter__.return_value.get = AsyncMock(return_value=mock_response) - mock_client.return_value = mock_context + mock_http = MagicMock() + mock_http.get = AsyncMock(return_value=mock_response) + mock_client.return_value = mock_http checker = UpdateChecker.get_instance(config=config) status, new_version = await checker._check_ytdlp_version() @@ -460,7 +460,7 @@ class TestUpdateChecker: assert "" == config.yt_new_version, "Should not set yt_new_version on HTTP error" @pytest.mark.asyncio - @patch("app.library.UpdateChecker.async_client") + @patch("app.library.UpdateChecker.get_async_client") async def test_check_for_updates_caches_separately(self, mock_client): """Test that app and yt-dlp checks are cached separately.""" from app.library.config import Config @@ -480,9 +480,9 @@ class TestUpdateChecker: mock_ytdlp_response.json.return_value = {"tag_name": "2026.01.01"} mock_get = AsyncMock(side_effect=[mock_app_response, mock_ytdlp_response]) - mock_context = AsyncMock() - mock_context.__aenter__.return_value.get = mock_get - mock_client.return_value = mock_context + mock_http = MagicMock() + mock_http.get = mock_get + mock_client.return_value = mock_http with patch("app.library.UpdateChecker.APP_VERSION", "v1.0.0"): checker = UpdateChecker.get_instance(config=config) From fe0a1c5b704b8ada89db3c649198b23128cf27a5 Mon Sep 17 00:00:00 2001 From: arabcoders Date: Mon, 26 Jan 2026 23:32:06 +0300 Subject: [PATCH 3/6] fix: offload info extraction to ProcessPoolExecutor --- FAQ.md | 1 - app/features/conditions/router.py | 4 +- .../tasks/definitions/handlers/generic.py | 5 +- .../tasks/definitions/handlers/rss.py | 5 +- app/features/tasks/definitions/results.py | 4 +- .../tests/test_generic_task_handler.py | 2 +- app/library/Utils.py | 158 +------ app/library/config.py | 8 - app/library/downloads/core.py | 13 +- app/library/downloads/extractor.py | 405 ++++++++++++++++++ app/library/downloads/item_adder.py | 34 +- app/library/downloads/playlist_processor.py | 98 ++--- app/library/downloads/queue_manager.py | 20 +- app/routes/api/yt_dlp.py | 16 +- app/tests/test_utils.py | 16 +- 15 files changed, 478 insertions(+), 311 deletions(-) create mode 100644 app/library/downloads/extractor.py diff --git a/FAQ.md b/FAQ.md index f1332746..42b56e70 100644 --- a/FAQ.md +++ b/FAQ.md @@ -46,7 +46,6 @@ or the `environment:` section in `compose.yaml` file. | YTP_PREVENT_LIVE_PREMIERE | Prevents the initial youtube premiere stream from being downloaded | `false` | | YTP_LIVE_PREMIERE_BUFFER | buffer time in minutes to add to video duration | `5` | | YTP_TASKS_HANDLER_TIMER | The cron expression for the tasks handler timer | `15 */1 * * *` | -| YTP_PLAYLIST_ITEMS_CONCURRENCY | The number of playlist items be to processed at same time | `1` | | YTP_TEMP_DISABLED | Disable temp files handling. | `false` | | YTP_DOWNLOAD_PATH_DEPTH | How many subdirectories to show in auto complete. | `1` | | YTP_ALLOW_INTERNAL_URLS | Allow requests to internal URLs | `false` | diff --git a/app/features/conditions/router.py b/app/features/conditions/router.py index b432c82f..1a01588e 100644 --- a/app/features/conditions/router.py +++ b/app/features/conditions/router.py @@ -87,10 +87,10 @@ async def conditions_test(request: Request, encoder: Encoder, cache: Cache, conf preset: str = params.get("preset", config.default_preset) key: str = cache.hash(url + str(preset)) if not cache.has(key): - from app.library.Utils import fetch_info + from app.library.downloads.extractor import fetch_info from app.library.YTDLPOpts import YTDLPOpts - data: dict | None = await fetch_info( + (data, _) = await fetch_info( config=YTDLPOpts.get_instance().preset(name=preset).get_all(), url=url, debug=False, diff --git a/app/features/tasks/definitions/handlers/generic.py b/app/features/tasks/definitions/handlers/generic.py index d2ac1f8b..6f4cfac3 100644 --- a/app/features/tasks/definitions/handlers/generic.py +++ b/app/features/tasks/definitions/handlers/generic.py @@ -22,8 +22,9 @@ from app.features.tasks.definitions.schemas import ( ) from app.library.cache import Cache from app.library.config import Config +from app.library.downloads.extractor import fetch_info from app.library.httpx_client import Globals, build_request_headers, get_async_client, resolve_curl_transport -from app.library.Utils import fetch_info, get_archive_id +from app.library.Utils import get_archive_id from ._base_handler import BaseHandler @@ -181,7 +182,7 @@ class GenericTaskHandler(BaseHandler): f"[{definition.name}]: '{task.name}': Unable to generate static archive id for '{url}' in feed. Doing real request to fetch yt-dlp archive id." ) - info = await fetch_info( + (info, _) = await fetch_info( config=task.get_ytdlp_opts().get_all(), url=url, no_archive=True, diff --git a/app/features/tasks/definitions/handlers/rss.py b/app/features/tasks/definitions/handlers/rss.py index 0f538a04..b7849b8d 100644 --- a/app/features/tasks/definitions/handlers/rss.py +++ b/app/features/tasks/definitions/handlers/rss.py @@ -6,7 +6,8 @@ from xml.etree.ElementTree import Element from app.features.tasks.definitions.results import HandleTask, TaskFailure, TaskItem, TaskResult from app.library.cache import Cache -from app.library.Utils import fetch_info, get_archive_id +from app.library.downloads.extractor import fetch_info +from app.library.Utils import get_archive_id from ._base_handler import BaseHandler @@ -179,7 +180,7 @@ class RssGenericHandler(BaseHandler): "Doing real request to fetch yt-dlp archive ID." ) - info = await fetch_info( + (info, _) = await fetch_info( config=params, url=url, no_archive=True, diff --git a/app/features/tasks/definitions/results.py b/app/features/tasks/definitions/results.py index e9d293a9..31d4fc3b 100644 --- a/app/features/tasks/definitions/results.py +++ b/app/features/tasks/definitions/results.py @@ -101,7 +101,7 @@ class HandleTask(TaskSchema): params_dict = params.get_all() - ie_info: dict | None = await fetch_info( + (ie_info, _) = await fetch_info( params_dict, self.url, no_archive=True, @@ -133,7 +133,7 @@ class HandleTask(TaskSchema): archive_file: Path = Path(archive_file) - ie_info: dict | None = await fetch_info(params, self.url, no_archive=True, follow_redirect=True) + (ie_info, _) = await fetch_info(params, self.url, no_archive=True, follow_redirect=True) if not ie_info or not isinstance(ie_info, dict): return (False, "Failed to extract information from URL.") diff --git a/app/features/tasks/definitions/tests/test_generic_task_handler.py b/app/features/tasks/definitions/tests/test_generic_task_handler.py index b55f78e2..eda8392e 100644 --- a/app/features/tasks/definitions/tests/test_generic_task_handler.py +++ b/app/features/tasks/definitions/tests/test_generic_task_handler.py @@ -339,7 +339,7 @@ async def test_generic_task_handler_inspect(monkeypatch): # Mock fetch_info to return valid info with required fields for archive ID generation async def fake_fetch_info(config, url, **kwargs): # noqa: ARG001 - return {"id": "test_video_1", "extractor_key": "Example"} + return ({"id": "test_video_1", "extractor_key": "Example"}, []) with patch("app.features.tasks.definitions.handlers.generic.fetch_info", side_effect=fake_fetch_info): task = HandleTask(id=1, name="Inspect", url="https://example.com/api") diff --git a/app/library/Utils.py b/app/library/Utils.py index b62ce94f..fbc14a47 100644 --- a/app/library/Utils.py +++ b/app/library/Utils.py @@ -1,7 +1,5 @@ -import asyncio import base64 import copy -import functools import glob import ipaddress import json @@ -22,8 +20,6 @@ from typing import Any from Crypto.Cipher import AES from yt_dlp.utils import age_restricted -from .LogWrapper import LogWrapper -from .mini_filter import match_str from .ytdlp import YTDLP, make_archive_id LOG: logging.Logger = logging.getLogger("Utils") @@ -88,9 +84,6 @@ TAG_REGEX: re.Pattern[str] = re.compile(r"%{([^:}]+)(?::([^}]*))?}c") DT_PATTERN: re.Pattern[str] = re.compile(r"^(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:[+-]\d{2}:\d{2}))\s?") "Regex to match ISO 8601 datetime strings." -EXTRACTORS_SEMAPHORE: asyncio.Semaphore | None = None -"Global semaphore for limiting concurrent extractor operations." - class StreamingError(Exception): """Raised when an error occurs during streaming.""" @@ -324,156 +317,7 @@ def calc_download_path(base_path: str | Path, folder: str | None = None, create_ return str(download_path) -def extract_info( - config: dict, - url: str, - debug: bool = False, - no_archive: bool = False, - follow_redirect: bool = False, - sanitize_info: bool = False, - **kwargs, -) -> dict: - """ - Extracts video information from the given URL. - - Args: - config (dict): Configuration options. - url (str): URL to extract information from. - debug (bool): Enable debug logging. - no_archive (bool): Disable download archive. - follow_redirect (bool): Follow URL redirects. - sanitize_info (bool): Sanitize the extracted information - **kwargs: Additional arguments. - - Returns: - dict: Video information. - - """ - params: dict = { - **config, - "simulate": True, - "color": "no_color", - "extract_flat": True, - "skip_download": True, - "ignoreerrors": True, - "ignore_no_formats_error": True, - } - - if debug: - params["verbose"] = True - else: - params["quiet"] = True - - log_wrapper = LogWrapper() - idDict: dict[str, str | None] = get_archive_id(url=url) - archive_id: str | None = f".{idDict['id']}" if idDict.get("id") else None - - log_wrapper.add_target( - target=logging.getLogger(f"yt-dlp{archive_id if archive_id else '.extract_info'}"), - level=logging.DEBUG if debug else logging.WARNING, - ) - - if "callback" in params: - if isinstance(params["callback"], dict): - log_wrapper.add_target( - target=params["callback"]["func"], - level=params["callback"]["level"] or logging.ERROR, - name=params["callback"]["name"] or "callback", - ) - else: - log_wrapper.add_target(target=params["callback"], level=logging.ERROR, name="callback") - - params.pop("callback", None) - - if log_wrapper.has_targets(): - if "logger" in params: - log_wrapper.add_target(target=params["logger"], level=logging.DEBUG) - - params["logger"] = log_wrapper - - if kwargs.get("no_log", False): - params["logger"] = LogWrapper() - params["quiet"] = True - params["no_warnings"] = True - - if no_archive and "download_archive" in params: - del params["download_archive"] - - data: dict[str, Any] | None = YTDLP(params=params).extract_info(url, download=False) - - if data and follow_redirect and "_type" in data and "url" == data["_type"]: - return extract_info( - config, - data["url"], - debug=debug, - no_archive=no_archive, - follow_redirect=follow_redirect, - sanitize_info=sanitize_info, - ) - - if not data: - return data - - data["is_premiere"] = match_str("media_type=video & duration & is_live", data) - if not data["is_premiere"]: - data["is_premiere"] = "video" == data.get("media_type") and "is_upcoming" == data.get("live_status") - - return YTDLP.sanitize_info(data, remove_private_keys=True) if sanitize_info else data - - -async def fetch_info( - config: dict, - url: str, - debug: bool = False, - no_archive: bool = False, - follow_redirect: bool = False, - sanitize_info: bool = False, - **kwargs, -) -> dict: - """ - Extracts video information from the given URL. - - Args: - config (dict): Configuration options. - url (str): URL to extract information from. - debug (bool): Enable debug logging. - no_archive (bool): Disable download archive. - follow_redirect (bool): Follow URL redirects. - sanitize_info (bool): Sanitize the extracted information - **kwargs: Additional arguments. - - Returns: - dict: Video information. - - """ - from .config import Config - - global EXTRACTORS_SEMAPHORE # noqa: PLW0603 - - conf = Config.get_instance() - if EXTRACTORS_SEMAPHORE is None: - EXTRACTORS_SEMAPHORE = asyncio.Semaphore(conf.extract_info_concurrency) - - async with EXTRACTORS_SEMAPHORE: - return await asyncio.wait_for( - fut=asyncio.get_running_loop().run_in_executor( - None, - functools.partial( - extract_info, - config=config, - url=url, - debug=debug, - no_archive=no_archive, - follow_redirect=follow_redirect, - sanitize_info=sanitize_info, - **kwargs, - ), - ), - timeout=conf.extract_info_timeout, - ) - - -def _is_safe_key(key: any) -> bool: +def _is_safe_key(key: Any) -> bool: """ Check if a dictionary key is safe for merging. diff --git a/app/library/config.py b/app/library/config.py index 6bca0d7e..3074e410 100644 --- a/app/library/config.py +++ b/app/library/config.py @@ -207,12 +207,6 @@ class Config(metaclass=Singleton): live_premiere_buffer: int = 5 """The buffer time in minutes to add to video duration to wait before starting premiere download.""" - add_items_concurrency: int = 4 - """The number of concurrent add items to be processed at same time.""" - - playlist_items_concurrency: int = 4 - """The number of concurrent playlist items to be processed at same time.""" - auto_clear_history_days: int = 0 """Number of days after which completed download history is automatically cleared. 0 to disable.""" @@ -271,13 +265,11 @@ class Config(metaclass=Singleton): "max_workers_per_extractor", "extract_info_timeout", "debugpy_port", - "playlist_items_concurrency", "download_path_depth", "download_info_expires", "auto_clear_history_days", "default_pagination", "extract_info_concurrency", - "add_items_concurrency", "flaresolverr_max_timeout", "flaresolverr_client_timeout", "flaresolverr_cache_ttl", diff --git a/app/library/downloads/core.py b/app/library/downloads/core.py index 416e80fd..be7213da 100644 --- a/app/library/downloads/core.py +++ b/app/library/downloads/core.py @@ -15,9 +15,10 @@ import yt_dlp.utils from app.library.config import Config from app.library.Events import EventBus, Events -from app.library.Utils import create_cookies_file, extract_info, extract_ytdlp_logs +from app.library.Utils import create_cookies_file, extract_ytdlp_logs from app.library.ytdlp import YTDLP +from .extractor import extract_info_sync from .hooks import HookHandlers, NestedLogger from .process_manager import ProcessManager from .status_tracker import StatusTracker @@ -156,21 +157,17 @@ class Download: if not self.info_dict or not isinstance(self.info_dict, dict): self.logger.info(f"Extracting info for '{self.info.url}'.") - self.logs = [] ie_params: dict = params.copy() - ie_params["callback"] = { - "func": lambda _, msg: self.logs.append(msg), - "level": logging.WARNING, - "name": "callback-logger", - } - info: dict = extract_info( + (info, logs) = extract_info_sync( config=ie_params, url=self.info.url, debug=self.debug, no_archive=not params.get("download_archive", False), follow_redirect=True, + capture_logs=logging.WARNING, ) + self.logs = [log["msg"] for log in logs] if info: self.info_dict = info diff --git a/app/library/downloads/extractor.py b/app/library/downloads/extractor.py new file mode 100644 index 00000000..3e11a7a6 --- /dev/null +++ b/app/library/downloads/extractor.py @@ -0,0 +1,405 @@ +import asyncio +import functools +import logging +import pickle +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path +from typing import Any + +from app.library.LogWrapper import LogWrapper +from app.library.mini_filter import match_str +from app.library.Singleton import Singleton +from app.library.ytdlp import YTDLP + +LOG: logging.Logger = logging.getLogger("downloads.extractor") + + +class ExtractorConfig: + """Configuration for the extractor.""" + + def __init__( + self, + concurrency: int = 4, + timeout: float = 60.0, + wait_threshold: float = 0.2, + ): + """ + Initialize extractor configuration. + + Args: + concurrency: Maximum number of concurrent extract operations + timeout: Timeout for extract operations in seconds + wait_threshold: Log warning if waiting exceeds this threshold in seconds + + """ + self.concurrency = concurrency + self.timeout = timeout + self.wait_threshold = wait_threshold + + +class ExtractorPool(metaclass=Singleton): + """ + Manages process pool and semaphore for video information extraction. + + This class uses the Singleton pattern to ensure only one instance exists. + """ + + def __init__(self): + """Initialize the extractor pool.""" + self._pool: ProcessPoolExecutor | None = None + self._semaphore: asyncio.Semaphore | None = None + self._config: ExtractorConfig | None = None + + @classmethod + def get_instance(cls) -> "ExtractorPool": + """ + Get the singleton instance. + + Returns: + ExtractorPool instance + + """ + return cls() + + def _ensure_initialized(self, config: ExtractorConfig) -> None: + """ + Ensure pool and semaphore are initialized. + + Args: + config: Extractor configuration + + """ + if self._config is None or self._config.concurrency != config.concurrency: + self._config = config + + if self._semaphore is None: + self._semaphore = asyncio.Semaphore(config.concurrency) + + if self._pool is None: + self._pool = ProcessPoolExecutor(max_workers=config.concurrency) + LOG.info("Initialized extractor process pool with %s workers", config.concurrency) + + def get_pool(self, config: ExtractorConfig) -> ProcessPoolExecutor: + """ + Get the process pool executor. + + Args: + config: Extractor configuration + + Returns: + ProcessPoolExecutor instance + + """ + self._ensure_initialized(config) + if self._pool is None: + msg = "Process pool not initialized" + raise RuntimeError(msg) + return self._pool + + def get_semaphore(self, config: ExtractorConfig) -> asyncio.Semaphore: + """ + Get the semaphore for limiting concurrency. + + Args: + config: Extractor configuration + + Returns: + asyncio.Semaphore instance + + """ + self._ensure_initialized(config) + if self._semaphore is None: + msg = "Semaphore not initialized" + raise RuntimeError(msg) + return self._semaphore + + async def shutdown(self) -> None: + """Shutdown the extractor pool and clean up resources.""" + if self._pool is not None: + try: + self._pool.shutdown(wait=False, cancel_futures=False) + LOG.debug("Extractor process pool shutdown complete") + except Exception as exc: + LOG.error("Error shutting down extractor process pool: %s", exc) + else: + self._pool = None + + self._semaphore = None + self._config = None + + +def _is_picklable(value: Any) -> bool: + """ + Check if a value can be pickled. + + Args: + value: Value to check + + Returns: + bool: True if value can be pickled, False otherwise + + """ + try: + pickle.dumps(value) + return True + except Exception: + return False + + +def _sanitize_picklable(value: Any) -> Any: + """ + Convert a value to a picklable format. + + Args: + value: Value to sanitize + + Returns: + Sanitized value that can be pickled + + """ + if isinstance(value, Path): + return str(value) + if isinstance(value, dict): + return {k: _sanitize_picklable(v) for k, v in value.items()} + if isinstance(value, (list, tuple, set)): + return [_sanitize_picklable(v) for v in value] + if _is_picklable(value): + return value + return str(value) + + +def _sanitize_config(config: dict[str, Any]) -> dict[str, Any]: + """ + Sanitize configuration for use in process pool. + + Removes unpicklable keys and converts values to picklable formats. + + Args: + config: Configuration dictionary + + Returns: + Sanitized configuration dictionary + + """ + sanitized: dict[str, Any] = {} + for key, value in config.items(): + if key in {"logger", "progress_hooks", "postprocessor_hooks", "postprocessors", "post_hooks"}: + continue + sanitized[key] = _sanitize_picklable(value) + + return sanitized + + +def _get_archive_id_dict(url: str) -> dict[str, str | None]: + """ + Get archive ID for logging purposes. + + Args: + url: URL to get archive ID for + + Returns: + Dictionary with id, ie_key, and archive_id + + """ + from app.library.Utils import get_archive_id + + return get_archive_id(url=url) + + +def extract_info_sync( + config: dict[str, Any], + url: str, + debug: bool = False, + no_archive: bool = False, + follow_redirect: bool = False, + sanitize_info: bool = False, + capture_logs: int | None = None, + **kwargs, +) -> tuple[dict[str, Any] | None, list[dict[str, Any]]]: + """ + Extract video information from a URL. + + Args: + config: yt-dlp configuration options + url: URL to extract information from + debug: Enable debug logging + no_archive: Disable download archive + follow_redirect: Follow URL redirects + sanitize_info: Sanitize the extracted information + capture_logs: If provided (e.g., logging.WARNING), capture logs at this level. + **kwargs: Additional arguments + + Returns: + tuple[dict | None, list[dict]]: Extracted information and captured logs. + + """ + params: dict[str, Any] = { + **config, + "simulate": True, + "color": "no_color", + "extract_flat": True, + "skip_download": True, + "ignoreerrors": True, + "ignore_no_formats_error": True, + } + + if debug: + params["verbose"] = True + else: + params["quiet"] = True + + log_wrapper = LogWrapper() + id_dict: dict[str, str | None] = _get_archive_id_dict(url=url) + archive_id: str | None = f".{id_dict['id']}" if id_dict.get("id") else None + + log_wrapper.add_target( + target=logging.getLogger(f"yt-dlp{archive_id if archive_id else '.extract_info'}"), + level=logging.DEBUG if debug else logging.WARNING, + ) + + captured_logs: list[dict[str, Any]] = kwargs.get("captured_logs", []) + if capture_logs is not None: + log_wrapper.add_target( + target=lambda level, msg: captured_logs.append({"level": level, "msg": msg}), + level=capture_logs, + name="log-capture", + ) + + if log_wrapper.has_targets(): + if "logger" in params: + log_wrapper.add_target(target=params["logger"], level=logging.DEBUG) + + params["logger"] = log_wrapper + + if kwargs.get("no_log", False): + params["logger"] = LogWrapper() + params["quiet"] = True + params["no_warnings"] = True + + if no_archive and "download_archive" in params: + del params["download_archive"] + + data: dict[str, Any] | None = YTDLP(params=params).extract_info(url, download=False) + + if data and follow_redirect and "_type" in data and "url" == data["_type"]: + return extract_info_sync( + config, + data["url"], + debug=debug, + no_archive=no_archive, + follow_redirect=follow_redirect, + sanitize_info=sanitize_info, + capture_logs=capture_logs, + captured_logs=captured_logs, + **kwargs, + ) + + if not data: + return (data, captured_logs) + + data["is_premiere"] = match_str("media_type=video & duration & is_live", data) + if not data["is_premiere"]: + data["is_premiere"] = "video" == data.get("media_type") and "is_upcoming" == data.get("live_status") + + result = YTDLP.sanitize_info(data, remove_private_keys=True) if sanitize_info else data + + return (result, captured_logs) + + +async def fetch_info( + config: dict[str, Any], + url: str, + debug: bool = False, + no_archive: bool = False, + follow_redirect: bool = False, + sanitize_info: bool = False, + capture_logs: int | None = None, + extractor_config: ExtractorConfig | None = None, + **kwargs, +) -> tuple[dict[str, Any] | None, list[dict[str, Any]]]: + """ + Extract video information from a URL. + + This function uses a process pool to avoid blocking the event loop. + If the process pool fails, it falls back to using a thread pool. + + Args: + config: yt-dlp configuration options + url: URL to extract information from + debug: Enable debug logging + no_archive: Disable download archive + follow_redirect: Follow URL redirects + sanitize_info: Sanitize the extracted information + capture_logs: If provided (e.g., logging.WARNING), capture logs + extractor_config: Configuration for the extractor + **kwargs: Additional arguments + + Returns: + tuple[dict | None, list[dict]]: Extracted information and captured logs. + + """ + if extractor_config is None: + from app.library.config import Config + + conf = Config.get_instance() + extractor_config = ExtractorConfig( + concurrency=conf.extract_info_concurrency, + timeout=conf.extract_info_timeout, + wait_threshold=0.2, + ) + + pool_manager: ExtractorPool = ExtractorPool.get_instance() + semaphore: asyncio.Semaphore = pool_manager.get_semaphore(extractor_config) + + await semaphore.acquire() + loop = asyncio.get_running_loop() + + safe_config = _sanitize_config(config) + + try: + try: + executor: ProcessPoolExecutor = pool_manager.get_pool(extractor_config) + + return await asyncio.wait_for( + fut=loop.run_in_executor( + executor, + functools.partial( + extract_info_sync, + config=safe_config, + url=url, + debug=debug, + no_archive=no_archive, + follow_redirect=follow_redirect, + sanitize_info=sanitize_info, + capture_logs=capture_logs, + **kwargs, + ), + ), + timeout=extractor_config.timeout, + ) + + except Exception as exc: + LOG.warning("extract_info process pool failed, falling back to thread pool url=%s error=%s", url, exc) + return await asyncio.wait_for( + fut=loop.run_in_executor( + None, + functools.partial( + extract_info_sync, + config=config, + url=url, + debug=debug, + no_archive=no_archive, + follow_redirect=follow_redirect, + sanitize_info=sanitize_info, + capture_logs=capture_logs, + **kwargs, + ), + ), + timeout=extractor_config.timeout, + ) + finally: + semaphore.release() + + +async def shutdown_extractor() -> None: + await ExtractorPool.get_instance().shutdown() diff --git a/app/library/downloads/item_adder.py b/app/library/downloads/item_adder.py index 40036a41..0db43f47 100644 --- a/app/library/downloads/item_adder.py +++ b/app/library/downloads/item_adder.py @@ -16,13 +16,13 @@ from app.library.Utils import ( archive_read, arg_converter, create_cookies_file, - fetch_info, get_extras, merge_dict, ytdlp_reject, ) from .core import Download +from .extractor import fetch_info from .playlist_processor import process_playlist from .video_processor import add_video @@ -76,11 +76,7 @@ async def add_item( async def add( - queue: "DownloadQueue", - item: "Item", - already: set | None = None, - entry: dict | None = None, - playlist: bool = False, + queue: "DownloadQueue", item: "Item", already: set | None = None, entry: dict | None = None ) -> dict[str, str]: """ Add an item to the download queue. @@ -90,23 +86,11 @@ async def add( item: Item to be added to the queue already: Set of already downloaded items entry: Entry associated with the item (if already extracted) - playlist: Whether the item is part of a playlist Returns: dict[str, str]: Status dict with "status" and optional "msg" keys """ - if playlist: - # If not done like this, it would lead to deadlocks when adding playlist items. - return await _add_impl(queue=queue, item=item, already=already, entry=entry) - - async with queue.add_sem: - return await _add_impl(queue=queue, item=item, already=already, entry=entry) - - -async def _add_impl( - queue: "DownloadQueue", item: "Item", already: set | None = None, entry: dict | None = None -) -> dict[str, str]: _preset: Preset | None = Presets.get_instance().get(item.preset) if item.has_cli(): @@ -137,16 +121,7 @@ async def _add_impl( already.add(item.url) try: - logs: list = [] - - yt_conf: dict = { - "callback": { - "func": lambda _, msg: logs.append(msg), - "level": logging.WARNING, - "name": "callback-logger", - }, - **item.get_ytdlp_opts().get_all(), - } + yt_conf: dict = item.get_ytdlp_opts().get_all() if yt_conf.get("external_downloader"): LOG.warning(f"Using external downloader '{yt_conf.get('external_downloader')}' for '{item.url}'.") @@ -209,12 +184,13 @@ async def _add_impl( if not entry: LOG.info(f"Extracting '{item.url}'{' with cookies' if yt_conf.get('cookiefile') else ''}.") - entry: dict | None = await fetch_info( + (entry, logs) = await fetch_info( config=yt_conf, url=item.url, debug=bool(queue.config.ytdlp_debug), no_archive=False, follow_redirect=True, + capture_logs=logging.WARNING, ) if not entry: diff --git a/app/library/downloads/playlist_processor.py b/app/library/downloads/playlist_processor.py index 1ea39dcd..0f3b638f 100644 --- a/app/library/downloads/playlist_processor.py +++ b/app/library/downloads/playlist_processor.py @@ -1,13 +1,10 @@ """Playlist processing.""" -import asyncio import logging from typing import TYPE_CHECKING, Any from app.library.Utils import merge_dict, ytdlp_reject -from .utils import handle_task_exception - if TYPE_CHECKING: from app.library.ItemDTO import Item @@ -59,87 +56,54 @@ async def process_playlist( "n_entries": len(entries), } - async def playlist_processor(i: int, etr: dict): - acquired = False - try: - item_name: str = ( - f"'{entry.get('title')}: {i}/{playlist_keys['n_entries']}' - '{etr.get('id')}: {etr.get('title')}'" - ) - LOG.debug(f"Waiting to acquire lock for {item_name}") - await queue.processors.acquire() - acquired = True - LOG.debug(f"Acquired lock for {item_name}") + async def process_item(i: int, etr: dict) -> dict[str, str]: + """Process a single playlist item.""" + item_name: str = ( + f"'{entry.get('title')}: {i}/{playlist_keys['n_entries']}' - '{etr.get('id')}: {etr.get('title')}'" + ) + LOG.info(f"Processing '{item_name}'.") - LOG.info(f"Processing '{item_name}'.") + _status, _msg = ytdlp_reject(entry=etr, yt_params=yt_params) + if not _status: + return {"status": "error", "msg": _msg} - _status, _msg = ytdlp_reject(entry=etr, yt_params=yt_params) - if not _status: - return {"status": "error", "msg": _msg} + extras: dict[str, Any] = { + **playlist_keys, + "playlist_index": i, + "playlist_index_number": i, + "playlist_autonumber": i, + } - extras: dict[str, Any] = { - **playlist_keys, - "playlist_index": i, - "playlist_index_number": i, - "playlist_autonumber": i, - } + for property in ("id", "title", "uploader", "uploader_id"): + if property in entry: + extras[f"playlist_{property}"] = entry.get(property) - for property in ("id", "title", "uploader", "uploader_id"): - if property in entry: - extras[f"playlist_{property}"] = entry.get(property) + extractor_key = entry.get("ie_key") or entry.get("extractor_key") or entry.get("extractor") or "" + if "thumbnail" not in etr and "youtube" in str(extractor_key).lower(): + extras["thumbnail"] = "https://img.youtube.com/vi/{id}/maxresdefault.jpg".format(**etr) - extractor_key = entry.get("ie_key") or entry.get("extractor_key") or entry.get("extractor") or "" - if "thumbnail" not in etr and "youtube" in str(extractor_key).lower(): - extras["thumbnail"] = "https://img.youtube.com/vi/{id}/maxresdefault.jpg".format(**etr) + newItem: Item = item.new_with(url=etr.get("url") or etr.get("webpage_url"), extras=extras) - newItem: Item = item.new_with(url=etr.get("url") or etr.get("webpage_url"), extras=extras) + if ("video" == etr.get("_type") and etr.get("url")) or ( + "formats" in etr and isinstance(etr["formats"], list) and len(etr["formats"]) > 0 + ): + dct = merge_dict(merge_dict({"_type": "video"}, etr), entry) + dct.pop("entries", None) + return await queue.add(item=newItem, entry=dct, already=already) - if ("video" == etr.get("_type") and etr.get("url")) or ( - "formats" in etr and isinstance(etr["formats"], list) and len(etr["formats"]) > 0 - ): - dct = merge_dict(merge_dict({"_type": "video"}, etr), entry) - dct.pop("entries", None) - return await queue.add(item=newItem, entry=dct, already=already, playlist=True) - - return await queue.add(item=newItem, already=already, playlist=True) - finally: - if acquired: - queue.processors.release() + return await queue.add(item=newItem, already=already) max_downloads: int = -1 ytdlp_opts: dict[str, Any] = item.get_ytdlp_opts().get_all() if ytdlp_opts.get("max_downloads") and isinstance(ytdlp_opts.get("max_downloads"), int): max_downloads: int = ytdlp_opts.get("max_downloads") - batch_size: int = max(50, int(queue.config.playlist_items_concurrency) * 10) - - async def run_batch(batch: list[tuple[int, dict]]) -> list[dict]: - tasks: list[asyncio.Task] = [] - for i, etr in batch: - task = asyncio.create_task( - playlist_processor(i, etr), - name=f"playlist_processor_{etr.get('id')}_{i}", - ) - task.add_done_callback(lambda t: handle_task_exception(t, LOG)) - tasks.append(task) - - batch_results: list[dict] = await asyncio.gather(*tasks) - await asyncio.sleep(0) - return batch_results - - results: list[dict] = [] - batch: list[tuple[int, dict]] = [] - + results: list[dict[str, str]] = [] for i, etr in enumerate(entries, start=1): if max_downloads > 0 and i > max_downloads: break - batch.append((i, etr)) - if len(batch) >= batch_size: - results.extend(await run_batch(batch)) - batch.clear() - - if batch: - results.extend(await run_batch(batch)) + results.append(await process_item(i, etr)) log_msg: str = f"Playlist '{playlist_name}' processing completed with '{len(results)}' entries." if max_downloads > 0 and len(entries) > max_downloads: diff --git a/app/library/downloads/queue_manager.py b/app/library/downloads/queue_manager.py index c4ab5b11..1230d208 100644 --- a/app/library/downloads/queue_manager.py +++ b/app/library/downloads/queue_manager.py @@ -1,4 +1,3 @@ -import asyncio import functools import glob import logging @@ -9,6 +8,7 @@ from typing import TYPE_CHECKING from aiohttp import web from app.library.config import Config +from app.library.downloads.extractor import shutdown_extractor from app.library.Events import EventBus, Events from app.library.ItemDTO import Item, ItemDTO from app.library.Scheduler import Scheduler @@ -41,10 +41,6 @@ class DownloadQueue(metaclass=Singleton): "DataStore for the completed downloads." self.queue = DataStore(type=StoreType.QUEUE, connection=SqliteStore.get_instance()) "DataStore for the download queue." - self.processors = asyncio.Semaphore(self.config.playlist_items_concurrency) - "Semaphore to limit the number of concurrent processors." - self.add_sem = asyncio.Semaphore(self.config.add_items_concurrency) - "Semaphore to limit the number of concurrent add item operations." self.pool = PoolManager(queue=self, config=self.config) "Pool manager for coordinating download execution." @@ -52,7 +48,7 @@ class DownloadQueue(metaclass=Singleton): def get_instance(config: Config | None = None) -> "DownloadQueue": return DownloadQueue(config=config) - def attach(self, _: web.Application) -> None: + def attach(self, app: web.Application) -> None: Services.get_instance().add("queue", self) async def event_handler(_, __): @@ -79,7 +75,7 @@ class DownloadQueue(metaclass=Singleton): id=delete_old_history.__name__, ) - # app.on_shutdown.append(self.on_shutdown) + app.on_shutdown.append(self.on_shutdown) async def test(self) -> bool: await self.done.test() @@ -221,11 +217,10 @@ class DownloadQueue(metaclass=Singleton): return self.pool.is_paused() async def on_shutdown(self, _: web.Application): - await self.pool.shutdown() + # await self.pool.shutdown() + await shutdown_extractor() - async def add( - self, item: Item, already: set | None = None, entry: dict | None = None, playlist: bool = False - ) -> dict[str, str]: + async def add(self, item: Item, already: set | None = None, entry: dict | None = None) -> dict[str, str]: """ Add an item to the download queue. @@ -233,13 +228,12 @@ class DownloadQueue(metaclass=Singleton): item: Item to be added to the queue already: Set of already downloaded items entry: Entry associated with the item (if already extracted) - playlist: Whether the item is part of a playlist Returns: dict[str, str]: Status dict with "status" and optional "msg" keys """ - return await add_impl(queue=self, item=item, already=already, entry=entry, playlist=playlist) + return await add_impl(queue=self, item=item, already=already, entry=entry) async def cancel(self, ids: list[str]) -> dict[str, str]: """ diff --git a/app/routes/api/yt_dlp.py b/app/routes/api/yt_dlp.py index 996e0248..f5fdd7d6 100644 --- a/app/routes/api/yt_dlp.py +++ b/app/routes/api/yt_dlp.py @@ -10,6 +10,7 @@ from aiohttp.web import Request, Response from app.features.presets.service import Presets from app.library.cache import Cache from app.library.config import Config +from app.library.downloads.extractor import fetch_info from app.library.encoder import Encoder from app.library.ItemDTO import Item from app.library.router import route @@ -17,7 +18,6 @@ from app.library.Utils import ( REMOVE_KEYS, archive_read, arg_converter, - fetch_info, get_archive_id, validate_url, ) @@ -152,24 +152,16 @@ async def get_info(request: Request, cache: Cache, config: Config) -> Response: } return web.json_response(body=json.dumps(data, indent=4, default=str), status=web.HTTPOk.status_code) - logs: list = [] + ytdlp_opts: dict = opts.get_all() - ytdlp_opts: dict = { - **opts.get_all(), - "callback": { - "func": lambda _, msg: logs.append(msg), - "level": logging.WARNING, - "name": "callback-logger", - }, - } - - data: dict | None = await fetch_info( + (data, logs) = await fetch_info( config=ytdlp_opts, url=url, debug=False, no_archive=True, follow_redirect=True, sanitize_info=True, + capture_logs=logging.WARNING, ) if not data or not isinstance(data, dict): diff --git a/app/tests/test_utils.py b/app/tests/test_utils.py index 45d1c9d1..de54d178 100644 --- a/app/tests/test_utils.py +++ b/app/tests/test_utils.py @@ -25,7 +25,6 @@ from app.library.Utils import ( delete_dir, dt_delta, encrypt_data, - extract_info, extract_ytdlp_logs, get, get_archive_id, @@ -54,6 +53,7 @@ from app.library.Utils import ( validate_uuid, ytdlp_reject, ) +from app.library.downloads.extractor import extract_info_sync class TestStreamingError: @@ -1353,7 +1353,7 @@ class TestArchiveFunctions: class TestExtractInfo: """Test the extract_info function.""" - @patch("app.library.Utils.YTDLP") + @patch("app.library.downloads.extractor.YTDLP") def test_extract_info_basic(self, mock_ytdlp_class): """Test basic extract_info functionality.""" mock_ytdlp = MagicMock() @@ -1363,11 +1363,12 @@ class TestExtractInfo: config = {"quiet": True} url = "https://example.com/video" - result = extract_info(config, url) - assert isinstance(result, dict) + (result, logs) = extract_info_sync(config, url) + assert isinstance(result, dict), "Result should be a dictionary" + assert isinstance(logs, list), "Logs should be a list" mock_ytdlp.extract_info.assert_called_once() - @patch("app.library.Utils.YTDLP") + @patch("app.library.downloads.extractor.YTDLP") def test_extract_info_with_debug(self, mock_ytdlp_class): """Test extract_info with debug enabled.""" mock_ytdlp = MagicMock() @@ -1377,8 +1378,9 @@ class TestExtractInfo: config = {} url = "https://example.com/video" - result = extract_info(config, url, debug=True) - assert isinstance(result, dict) + (result, logs) = extract_info_sync(config, url, debug=True) + assert isinstance(result, dict), "Result should be a dictionary" + assert isinstance(logs, list), "Logs should be a list" class TestCheckId: From da9f08297866279cc0c70f31c5e9d104430501ac Mon Sep 17 00:00:00 2001 From: arabcoders Date: Tue, 27 Jan 2026 10:17:22 +0300 Subject: [PATCH 4/6] fix: unintended return type change. --- app/library/downloads/core.py | 2 +- app/library/downloads/extractor.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/app/library/downloads/core.py b/app/library/downloads/core.py index be7213da..4b512e3e 100644 --- a/app/library/downloads/core.py +++ b/app/library/downloads/core.py @@ -167,7 +167,7 @@ class Download: follow_redirect=True, capture_logs=logging.WARNING, ) - self.logs = [log["msg"] for log in logs] + self.logs = logs if info: self.info_dict = info diff --git a/app/library/downloads/extractor.py b/app/library/downloads/extractor.py index 3e11a7a6..88cdf093 100644 --- a/app/library/downloads/extractor.py +++ b/app/library/downloads/extractor.py @@ -257,10 +257,10 @@ def extract_info_sync( level=logging.DEBUG if debug else logging.WARNING, ) - captured_logs: list[dict[str, Any]] = kwargs.get("captured_logs", []) + captured_logs: list[str] = kwargs.get("captured_logs", []) if capture_logs is not None: log_wrapper.add_target( - target=lambda level, msg: captured_logs.append({"level": level, "msg": msg}), + target=lambda _, msg: captured_logs.append(msg), level=capture_logs, name="log-capture", ) From dae141a6c22e11d225172c5edce1ce21c862b2d2 Mon Sep 17 00:00:00 2001 From: arabcoders Date: Tue, 27 Jan 2026 10:41:17 +0300 Subject: [PATCH 5/6] revert: changes to playlist processor --- app/library/config.py | 4 + app/library/downloads/playlist_processor.py | 98 ++++++++++++++------- app/library/downloads/queue_manager.py | 3 + 3 files changed, 74 insertions(+), 31 deletions(-) diff --git a/app/library/config.py b/app/library/config.py index 3074e410..c139ed42 100644 --- a/app/library/config.py +++ b/app/library/config.py @@ -207,6 +207,9 @@ class Config(metaclass=Singleton): live_premiere_buffer: int = 5 """The buffer time in minutes to add to video duration to wait before starting premiere download.""" + playlist_items_concurrency: int = 4 + """The number of concurrent playlist items to be processed at same time.""" + auto_clear_history_days: int = 0 """Number of days after which completed download history is automatically cleared. 0 to disable.""" @@ -265,6 +268,7 @@ class Config(metaclass=Singleton): "max_workers_per_extractor", "extract_info_timeout", "debugpy_port", + "playlist_items_concurrency", "download_path_depth", "download_info_expires", "auto_clear_history_days", diff --git a/app/library/downloads/playlist_processor.py b/app/library/downloads/playlist_processor.py index 0f3b638f..2a7006d3 100644 --- a/app/library/downloads/playlist_processor.py +++ b/app/library/downloads/playlist_processor.py @@ -1,10 +1,13 @@ """Playlist processing.""" +import asyncio import logging from typing import TYPE_CHECKING, Any from app.library.Utils import merge_dict, ytdlp_reject +from .utils import handle_task_exception + if TYPE_CHECKING: from app.library.ItemDTO import Item @@ -56,54 +59,87 @@ async def process_playlist( "n_entries": len(entries), } - async def process_item(i: int, etr: dict) -> dict[str, str]: - """Process a single playlist item.""" - item_name: str = ( - f"'{entry.get('title')}: {i}/{playlist_keys['n_entries']}' - '{etr.get('id')}: {etr.get('title')}'" - ) - LOG.info(f"Processing '{item_name}'.") + async def playlist_processor(i: int, etr: dict): + acquired = False + try: + item_name: str = ( + f"'{entry.get('title')}: {i}/{playlist_keys['n_entries']}' - '{etr.get('id')}: {etr.get('title')}'" + ) + LOG.debug(f"Waiting to acquire lock for {item_name}") + await queue.processors.acquire() + acquired = True + LOG.debug(f"Acquired lock for {item_name}") - _status, _msg = ytdlp_reject(entry=etr, yt_params=yt_params) - if not _status: - return {"status": "error", "msg": _msg} + LOG.info(f"Processing '{item_name}'.") - extras: dict[str, Any] = { - **playlist_keys, - "playlist_index": i, - "playlist_index_number": i, - "playlist_autonumber": i, - } + _status, _msg = ytdlp_reject(entry=etr, yt_params=yt_params) + if not _status: + return {"status": "error", "msg": _msg} - for property in ("id", "title", "uploader", "uploader_id"): - if property in entry: - extras[f"playlist_{property}"] = entry.get(property) + extras: dict[str, Any] = { + **playlist_keys, + "playlist_index": i, + "playlist_index_number": i, + "playlist_autonumber": i, + } - extractor_key = entry.get("ie_key") or entry.get("extractor_key") or entry.get("extractor") or "" - if "thumbnail" not in etr and "youtube" in str(extractor_key).lower(): - extras["thumbnail"] = "https://img.youtube.com/vi/{id}/maxresdefault.jpg".format(**etr) + for property in ("id", "title", "uploader", "uploader_id"): + if property in entry: + extras[f"playlist_{property}"] = entry.get(property) - newItem: Item = item.new_with(url=etr.get("url") or etr.get("webpage_url"), extras=extras) + extractor_key = entry.get("ie_key") or entry.get("extractor_key") or entry.get("extractor") or "" + if "thumbnail" not in etr and "youtube" in str(extractor_key).lower(): + extras["thumbnail"] = "https://img.youtube.com/vi/{id}/maxresdefault.jpg".format(**etr) - if ("video" == etr.get("_type") and etr.get("url")) or ( - "formats" in etr and isinstance(etr["formats"], list) and len(etr["formats"]) > 0 - ): - dct = merge_dict(merge_dict({"_type": "video"}, etr), entry) - dct.pop("entries", None) - return await queue.add(item=newItem, entry=dct, already=already) + newItem: Item = item.new_with(url=etr.get("url") or etr.get("webpage_url"), extras=extras) - return await queue.add(item=newItem, already=already) + if ("video" == etr.get("_type") and etr.get("url")) or ( + "formats" in etr and isinstance(etr["formats"], list) and len(etr["formats"]) > 0 + ): + dct = merge_dict(merge_dict({"_type": "video"}, etr), entry) + dct.pop("entries", None) + return await queue.add(item=newItem, entry=dct, already=already) + + return await queue.add(item=newItem, already=already) + finally: + if acquired: + queue.processors.release() max_downloads: int = -1 ytdlp_opts: dict[str, Any] = item.get_ytdlp_opts().get_all() if ytdlp_opts.get("max_downloads") and isinstance(ytdlp_opts.get("max_downloads"), int): max_downloads: int = ytdlp_opts.get("max_downloads") - results: list[dict[str, str]] = [] + batch_size: int = max(50, int(queue.config.playlist_items_concurrency) * 10) + + async def run_batch(batch: list[tuple[int, dict]]) -> list[dict]: + tasks: list[asyncio.Task] = [] + for i, etr in batch: + task = asyncio.create_task( + playlist_processor(i, etr), + name=f"playlist_processor_{etr.get('id')}_{i}", + ) + task.add_done_callback(lambda t: handle_task_exception(t, LOG)) + tasks.append(task) + + batch_results: list[dict] = await asyncio.gather(*tasks) + await asyncio.sleep(0) + return batch_results + + results: list[dict] = [] + batch: list[tuple[int, dict]] = [] + for i, etr in enumerate(entries, start=1): if max_downloads > 0 and i > max_downloads: break - results.append(await process_item(i, etr)) + batch.append((i, etr)) + if len(batch) >= batch_size: + results.extend(await run_batch(batch)) + batch.clear() + + if batch: + results.extend(await run_batch(batch)) log_msg: str = f"Playlist '{playlist_name}' processing completed with '{len(results)}' entries." if max_downloads > 0 and len(entries) > max_downloads: diff --git a/app/library/downloads/queue_manager.py b/app/library/downloads/queue_manager.py index 1230d208..9f75c173 100644 --- a/app/library/downloads/queue_manager.py +++ b/app/library/downloads/queue_manager.py @@ -1,3 +1,4 @@ +import asyncio import functools import glob import logging @@ -41,6 +42,8 @@ class DownloadQueue(metaclass=Singleton): "DataStore for the completed downloads." self.queue = DataStore(type=StoreType.QUEUE, connection=SqliteStore.get_instance()) "DataStore for the download queue." + self.processors = asyncio.Semaphore(self.config.playlist_items_concurrency) + "Semaphore to limit the number of concurrent processors." self.pool = PoolManager(queue=self, config=self.config) "Pool manager for coordinating download execution." From 240a920b6cd60026ec7f6b853c182e1d86816347 Mon Sep 17 00:00:00 2001 From: arabcoders Date: Tue, 27 Jan 2026 11:46:00 +0300 Subject: [PATCH 6/6] fix; proper fix for dispatching many websocket events --- .../tasks/definitions/handlers/generic.py | 2 +- .../tasks/definitions/handlers/rss.py | 2 +- .../tasks/definitions/handlers/tver.py | 2 +- .../tasks/definitions/handlers/twitch.py | 2 +- .../tasks/definitions/handlers/youtube.py | 2 +- app/features/tasks/definitions/service.py | 2 +- app/features/tasks/service.py | 2 +- app/library/HttpSocket.py | 13 ++- app/library/config.py | 4 - app/library/downloads/playlist_processor.py | 98 ++++++------------- app/library/downloads/queue_manager.py | 3 - 11 files changed, 50 insertions(+), 82 deletions(-) diff --git a/app/features/tasks/definitions/handlers/generic.py b/app/features/tasks/definitions/handlers/generic.py index 6f4cfac3..9d94c41a 100644 --- a/app/features/tasks/definitions/handlers/generic.py +++ b/app/features/tasks/definitions/handlers/generic.py @@ -34,7 +34,7 @@ if TYPE_CHECKING: import httpx from parsel.selector import SelectorList -LOG: logging.Logger = logging.getLogger(__name__) +LOG: logging.Logger = logging.getLogger("handlers.generic") CACHE: Cache = Cache() diff --git a/app/features/tasks/definitions/handlers/rss.py b/app/features/tasks/definitions/handlers/rss.py index b7849b8d..f5fc6431 100644 --- a/app/features/tasks/definitions/handlers/rss.py +++ b/app/features/tasks/definitions/handlers/rss.py @@ -14,7 +14,7 @@ from ._base_handler import BaseHandler if TYPE_CHECKING: from xml.etree.ElementTree import Element -LOG: logging.Logger = logging.getLogger(__name__) +LOG: logging.Logger = logging.getLogger("handlers.rss") CACHE: Cache = Cache() diff --git a/app/features/tasks/definitions/handlers/tver.py b/app/features/tasks/definitions/handlers/tver.py index 81c2a4e8..29aa3695 100644 --- a/app/features/tasks/definitions/handlers/tver.py +++ b/app/features/tasks/definitions/handlers/tver.py @@ -6,7 +6,7 @@ from app.library.Utils import get_archive_id from ._base_handler import BaseHandler -LOG: logging.Logger = logging.getLogger(__name__) +LOG: logging.Logger = logging.getLogger("handlers.tver") class TverHandler(BaseHandler): diff --git a/app/features/tasks/definitions/handlers/twitch.py b/app/features/tasks/definitions/handlers/twitch.py index 25714bc2..2da417a0 100644 --- a/app/features/tasks/definitions/handlers/twitch.py +++ b/app/features/tasks/definitions/handlers/twitch.py @@ -11,7 +11,7 @@ from ._base_handler import BaseHandler if TYPE_CHECKING: from xml.etree.ElementTree import Element -LOG: logging.Logger = logging.getLogger(__name__) +LOG: logging.Logger = logging.getLogger("handlers.twitch") class TwitchHandler(BaseHandler): diff --git a/app/features/tasks/definitions/handlers/youtube.py b/app/features/tasks/definitions/handlers/youtube.py index 887ba78a..a963e14a 100644 --- a/app/features/tasks/definitions/handlers/youtube.py +++ b/app/features/tasks/definitions/handlers/youtube.py @@ -11,7 +11,7 @@ from ._base_handler import BaseHandler if TYPE_CHECKING: from xml.etree.ElementTree import Element -LOG: logging.Logger = logging.getLogger(__name__) +LOG: logging.Logger = logging.getLogger("handlers.youtube") class YoutubeHandler(BaseHandler): diff --git a/app/features/tasks/definitions/service.py b/app/features/tasks/definitions/service.py index f7702feb..36968c2d 100644 --- a/app/features/tasks/definitions/service.py +++ b/app/features/tasks/definitions/service.py @@ -22,7 +22,7 @@ if TYPE_CHECKING: from app.library.config import Config from app.library.Scheduler import Scheduler -LOG: logging.Logger = logging.getLogger("tasks.definitions.service") +LOG: logging.Logger = logging.getLogger("definitions.service") class TaskHandle: diff --git a/app/features/tasks/service.py b/app/features/tasks/service.py index 2e4900d7..82d5cd50 100644 --- a/app/features/tasks/service.py +++ b/app/features/tasks/service.py @@ -14,7 +14,7 @@ from app.library.Singleton import Singleton if TYPE_CHECKING: from aiohttp import web -LOG: logging.Logger = logging.getLogger(__name__) +LOG: logging.Logger = logging.getLogger("tasks.service") class Tasks(metaclass=Singleton): diff --git a/app/library/HttpSocket.py b/app/library/HttpSocket.py index 33163277..ceb8264c 100644 --- a/app/library/HttpSocket.py +++ b/app/library/HttpSocket.py @@ -1,3 +1,4 @@ +import asyncio import functools import json import logging @@ -186,10 +187,20 @@ class HttpSocket: await handle_connect(sid) + async def handle_message_safe(sid: str, message: str) -> None: + """Wrapper to handle message with error logging""" + try: + await handle_message(sid, message) + except Exception as e: + LOG.exception(e) + LOG.error(f"Error handling WebSocket message from client '{sid}': {e}") + try: + i: int = 0 async for msg in ws: if msg.type == web.WSMsgType.TEXT: - await handle_message(sid, msg.data) + i = i + 1 + asyncio.create_task(handle_message_safe(sid, msg.data), name=f"ws_msg_{sid}_{i}") elif msg.type == web.WSMsgType.ERROR: LOG.error(f"WebSocket connection closed with exception {ws.exception()}") finally: diff --git a/app/library/config.py b/app/library/config.py index c139ed42..3074e410 100644 --- a/app/library/config.py +++ b/app/library/config.py @@ -207,9 +207,6 @@ class Config(metaclass=Singleton): live_premiere_buffer: int = 5 """The buffer time in minutes to add to video duration to wait before starting premiere download.""" - playlist_items_concurrency: int = 4 - """The number of concurrent playlist items to be processed at same time.""" - auto_clear_history_days: int = 0 """Number of days after which completed download history is automatically cleared. 0 to disable.""" @@ -268,7 +265,6 @@ class Config(metaclass=Singleton): "max_workers_per_extractor", "extract_info_timeout", "debugpy_port", - "playlist_items_concurrency", "download_path_depth", "download_info_expires", "auto_clear_history_days", diff --git a/app/library/downloads/playlist_processor.py b/app/library/downloads/playlist_processor.py index 2a7006d3..0f3b638f 100644 --- a/app/library/downloads/playlist_processor.py +++ b/app/library/downloads/playlist_processor.py @@ -1,13 +1,10 @@ """Playlist processing.""" -import asyncio import logging from typing import TYPE_CHECKING, Any from app.library.Utils import merge_dict, ytdlp_reject -from .utils import handle_task_exception - if TYPE_CHECKING: from app.library.ItemDTO import Item @@ -59,87 +56,54 @@ async def process_playlist( "n_entries": len(entries), } - async def playlist_processor(i: int, etr: dict): - acquired = False - try: - item_name: str = ( - f"'{entry.get('title')}: {i}/{playlist_keys['n_entries']}' - '{etr.get('id')}: {etr.get('title')}'" - ) - LOG.debug(f"Waiting to acquire lock for {item_name}") - await queue.processors.acquire() - acquired = True - LOG.debug(f"Acquired lock for {item_name}") + async def process_item(i: int, etr: dict) -> dict[str, str]: + """Process a single playlist item.""" + item_name: str = ( + f"'{entry.get('title')}: {i}/{playlist_keys['n_entries']}' - '{etr.get('id')}: {etr.get('title')}'" + ) + LOG.info(f"Processing '{item_name}'.") - LOG.info(f"Processing '{item_name}'.") + _status, _msg = ytdlp_reject(entry=etr, yt_params=yt_params) + if not _status: + return {"status": "error", "msg": _msg} - _status, _msg = ytdlp_reject(entry=etr, yt_params=yt_params) - if not _status: - return {"status": "error", "msg": _msg} + extras: dict[str, Any] = { + **playlist_keys, + "playlist_index": i, + "playlist_index_number": i, + "playlist_autonumber": i, + } - extras: dict[str, Any] = { - **playlist_keys, - "playlist_index": i, - "playlist_index_number": i, - "playlist_autonumber": i, - } + for property in ("id", "title", "uploader", "uploader_id"): + if property in entry: + extras[f"playlist_{property}"] = entry.get(property) - for property in ("id", "title", "uploader", "uploader_id"): - if property in entry: - extras[f"playlist_{property}"] = entry.get(property) + extractor_key = entry.get("ie_key") or entry.get("extractor_key") or entry.get("extractor") or "" + if "thumbnail" not in etr and "youtube" in str(extractor_key).lower(): + extras["thumbnail"] = "https://img.youtube.com/vi/{id}/maxresdefault.jpg".format(**etr) - extractor_key = entry.get("ie_key") or entry.get("extractor_key") or entry.get("extractor") or "" - if "thumbnail" not in etr and "youtube" in str(extractor_key).lower(): - extras["thumbnail"] = "https://img.youtube.com/vi/{id}/maxresdefault.jpg".format(**etr) + newItem: Item = item.new_with(url=etr.get("url") or etr.get("webpage_url"), extras=extras) - newItem: Item = item.new_with(url=etr.get("url") or etr.get("webpage_url"), extras=extras) + if ("video" == etr.get("_type") and etr.get("url")) or ( + "formats" in etr and isinstance(etr["formats"], list) and len(etr["formats"]) > 0 + ): + dct = merge_dict(merge_dict({"_type": "video"}, etr), entry) + dct.pop("entries", None) + return await queue.add(item=newItem, entry=dct, already=already) - if ("video" == etr.get("_type") and etr.get("url")) or ( - "formats" in etr and isinstance(etr["formats"], list) and len(etr["formats"]) > 0 - ): - dct = merge_dict(merge_dict({"_type": "video"}, etr), entry) - dct.pop("entries", None) - return await queue.add(item=newItem, entry=dct, already=already) - - return await queue.add(item=newItem, already=already) - finally: - if acquired: - queue.processors.release() + return await queue.add(item=newItem, already=already) max_downloads: int = -1 ytdlp_opts: dict[str, Any] = item.get_ytdlp_opts().get_all() if ytdlp_opts.get("max_downloads") and isinstance(ytdlp_opts.get("max_downloads"), int): max_downloads: int = ytdlp_opts.get("max_downloads") - batch_size: int = max(50, int(queue.config.playlist_items_concurrency) * 10) - - async def run_batch(batch: list[tuple[int, dict]]) -> list[dict]: - tasks: list[asyncio.Task] = [] - for i, etr in batch: - task = asyncio.create_task( - playlist_processor(i, etr), - name=f"playlist_processor_{etr.get('id')}_{i}", - ) - task.add_done_callback(lambda t: handle_task_exception(t, LOG)) - tasks.append(task) - - batch_results: list[dict] = await asyncio.gather(*tasks) - await asyncio.sleep(0) - return batch_results - - results: list[dict] = [] - batch: list[tuple[int, dict]] = [] - + results: list[dict[str, str]] = [] for i, etr in enumerate(entries, start=1): if max_downloads > 0 and i > max_downloads: break - batch.append((i, etr)) - if len(batch) >= batch_size: - results.extend(await run_batch(batch)) - batch.clear() - - if batch: - results.extend(await run_batch(batch)) + results.append(await process_item(i, etr)) log_msg: str = f"Playlist '{playlist_name}' processing completed with '{len(results)}' entries." if max_downloads > 0 and len(entries) > max_downloads: diff --git a/app/library/downloads/queue_manager.py b/app/library/downloads/queue_manager.py index 9f75c173..1230d208 100644 --- a/app/library/downloads/queue_manager.py +++ b/app/library/downloads/queue_manager.py @@ -1,4 +1,3 @@ -import asyncio import functools import glob import logging @@ -42,8 +41,6 @@ class DownloadQueue(metaclass=Singleton): "DataStore for the completed downloads." self.queue = DataStore(type=StoreType.QUEUE, connection=SqliteStore.get_instance()) "DataStore for the download queue." - self.processors = asyncio.Semaphore(self.config.playlist_items_concurrency) - "Semaphore to limit the number of concurrent processors." self.pool = PoolManager(queue=self, config=self.config) "Pool manager for coordinating download execution."