502 lines
16 KiB
Python
502 lines
16 KiB
Python
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.Services import Services
|
|
from app.library.Singleton import Singleton
|
|
|
|
LOG: logging.Logger = logging.getLogger("downloads.extractor")
|
|
|
|
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()
|