import asyncio import functools import logging import multiprocessing import pickle import sys from concurrent.futures import ProcessPoolExecutor from pathlib import Path from typing import Any from aiohttp import web from app.features.ytdlp.utils import _DATA, LogWrapper, get_archive_id from app.features.ytdlp.ytdlp import YTDLP from app.library.log import get_logger from app.library.Services import Services from app.library.Singleton import Singleton LOG = get_logger() LIVE_REEXTRACT_STATUSES: set[str] = {"is_live", "post_live"} REEXTRACT_INFO_KEY = "_ytptube_reextract" def _ytdlp_logger(target: logging.Logger): return _YTDLPLogger(target) class _YTDLPLogger: def __init__(self, target: logging.Logger) -> None: self.target = target def __call__(self, level: int, msg: str, *args: Any, **kwargs: Any) -> None: kwargs.setdefault("stacklevel", 4) if level <= logging.DEBUG and isinstance(msg, str) and msg.startswith("[debug] "): self.target.debug(msg.removeprefix("[debug] "), *args, **kwargs) return if level <= logging.DEBUG: self.target.info(msg, *args, **kwargs) return self.target.log(level, msg, *args, **kwargs) class _LogCapture: def __init__(self, logs: list[str]) -> None: self.logs = logs def __call__(self, _: int, msg: str, *args: Any, **__: Any) -> None: self.logs.append(msg % args if args else msg) def _get_process_pool_kwargs() -> dict[str, Any]: """Use a fork-based pool for frozen Linux builds.""" if sys.platform == "linux" and getattr(sys, "frozen", False): return {"mp_context": multiprocessing.get_context("fork")} return {} 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 def _sleep_timeout(config: dict[str, Any], timeout: float, budget_sleep: bool) -> float: if not budget_sleep: return timeout sleep_requests = config.get("sleep_interval_requests") if not isinstance(sleep_requests, int | float) or sleep_requests <= 0: return timeout return timeout + min(float(sleep_requests) * 20, 300.0) 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 attach(self, app: web.Application) -> None: """Attach the extractor pool to the application (no-op for now).""" app.on_shutdown.append(self.shutdown) Services.get_instance().add(__class__.__name__, self) 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, **_get_process_pool_kwargs()) 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.exception( "Failed to shut down the extractor process pool.", extra={"exception_type": type(exc).__name__}, ) 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 needs_reextract(info: dict[str, Any]) -> bool: return bool(info.get("is_live") or info.get("live_status") in LIVE_REEXTRACT_STATUSES) def _process_safe_info(info: dict[str, Any]) -> dict[str, Any]: if needs_reextract(info): info = dict(info) info[REEXTRACT_INFO_KEY] = True for key in ("formats", "requested_formats", "requested_downloads", "fragments"): info.pop(key, None) if _is_picklable(info): return info return YTDLP.sanitize_info(info, remove_private_keys=False) 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, process_safe: bool = False, **kwargs, ) -> tuple[dict[str, Any] | None, list[str]]: """ 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. process_safe: Strip non-pickleable data for safe inter-process communication. **kwargs: Additional arguments Returns: tuple[dict | None, list[str]]: Extracted information and captured logs. """ params: dict[str, Any] = {**config, **_DATA.YTDLP_PARAMS, "simulate": True} if debug: params["verbose"] = True params.pop("quiet", None) suppress: tuple[str, ...] = () patterns = kwargs.get("suppress_logs") if isinstance(patterns, list | tuple): suppress = tuple(value for value in patterns if isinstance(value, str) and value) log_wrapper = LogWrapper(suppress=suppress) id_dict: dict[str, str | None] = get_archive_id(url=url) archive_id: str | None = f".{id_dict['id']}" if id_dict.get("id") else None logger_name: str = f"yt-dlp{archive_id or '.extract_info'}" try: log_wrapper.add_target( target=_ytdlp_logger(logging.getLogger(logger_name)), level=logging.DEBUG, name=logger_name, ) captured_logs: list[str] = kwargs.get("captured_logs", []) if capture_logs is not None: log_wrapper.add_target( target=_LogCapture(captured_logs), 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, process_safe=process_safe, captured_logs=captured_logs, **kwargs, ) if not data: return (data, captured_logs) try: from app.features.ytdlp.mini_filter import match_str 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") except ImportError: pass result = YTDLP.sanitize_info(data, remove_private_keys=True) if sanitize_info else data if process_safe and isinstance(result, dict): result = _process_safe_info(result) return (result, captured_logs) finally: logging.Logger.manager.loggerDict.pop(logger_name, None) 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, budget_sleep: bool = False, **kwargs, ) -> tuple[dict[str, Any] | None, list[str]]: """ 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 budget_sleep: Whether to add extra timeout budget for request-sleep-heavy extraction **kwargs: Additional arguments Returns: tuple[dict | None, list[str]]: 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) safe_kwargs = _sanitize_picklable(kwargs) safe_kwargs.pop("process_safe", None) timeout = _sleep_timeout(safe_config, extractor_config.timeout, budget_sleep) 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, process_safe=True, **safe_kwargs, ), ), timeout=timeout, ) except TimeoutError: raise except Exception as exc: LOG.warning( "yt-dlp extraction for '%s' fell back to the thread pool after the process pool failed.", url, extra={ "url": url, "timeout": timeout, "concurrency": extractor_config.concurrency, "pool": "process", "fallback": "thread", "follow_redirect": follow_redirect, "no_archive": no_archive, "sanitize_info": sanitize_info, "capture_logs_level": capture_logs, "exception_type": type(exc).__name__, }, exc_info=True, ) 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=timeout, ) finally: semaphore.release()