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: