diff --git a/app/library/AsyncPool.py b/app/library/AsyncPool.py deleted file mode 100644 index 2c0581ad..00000000 --- a/app/library/AsyncPool.py +++ /dev/null @@ -1,320 +0,0 @@ -import asyncio -import logging -import uuid -from datetime import UTC, datetime - - -class Terminator: - pass - - -class AsyncPool: - def __init__( - self, - num_workers: int, - worker_co, - name: str, - logger: logging.Logger, - loop: asyncio.AbstractEventLoop | None = None, - load_factor: int = 1, - max_task_time: int | None = None, - return_futures: bool = False, - raise_on_join: bool = False, - ): - """ - This class will create `num_workers` asyncio tasks to work against a queue of - `num_workers * load_factor` items of back-pressure (IOW we will block after such - number of items of work is in the queue). `worker_co` will be called - against each item retrieved from the queue. If any exceptions are raised out of - worker_co, self.exceptions will be set to True. - - Args: - num_workers (int): number of async tasks which will pull from the internal queue - worker_co (coroutine): async coroutine to call when an item is retrieved from the queue - name (str): name of the worker pool (used for logging) - logger (logging.Logger): logger to use - loop (asyncio.AbstractEventLoop|None): asyncio loop to use - load_factor (int): multiplier used for number of items in queue - max_task_time (int|None): maximum time allowed for each task before a CancelledError is raised in the task. - return_futures (bool): set to reture to return a future for each `push` (imposes CPU overhead) - raise_on_join (bool): raise on join if any exceptions have occurred, default is False - - Returns: - AsyncPool: instance of AsyncPool - - """ - self._loop = loop if loop else asyncio.get_event_loop() - self._num_workers = num_workers - self._logger = logger if logger else logging.getLogger(__name__) - self._queue = asyncio.Queue(num_workers * load_factor) - self._workers: dict[str, asyncio.Future] = {} - self._exceptions = False - self._max_task_time = max_task_time - self._return_futures = return_futures - self._raise_on_join = raise_on_join - self._name = name - self._worker_co = worker_co - self._status: dict[str, dict | None] = {} - - @property - def exceptions(self): - return self._exceptions - - async def _worker_loop(self, worker_id: str): - """ - The main persistent worker loop that continuously processes jobs from the shared queue. - - Args: - worker_id (str): the id of the worker processing this job. - - """ - while True: - item = await self._queue.get() - should_continue = await self._process_item(worker_id, item, from_queue=True) - if not should_continue: - break - - async def _process_item(self, worker_id: str, item: tuple | Terminator, from_queue: bool = True) -> bool: - """ - Processes a single job item. - - Args: - worker_id (str): the id of the worker processing this job. - item (tuple|Terminator): the job item (tuple of (future, args, kwargs)) or a Terminator. - from_queue (bool): indicates whether the job came from the shared queue; if True, task_done() is called. - - Returns: - bool: False if the item is a Terminator (indicating termination), True otherwise. - - """ - future = None - is_terminator = isinstance(item, Terminator) - - try: - if is_terminator: - return False - - future, args, kwargs = item - - self._status[worker_id] = { - "started": self._time().isoformat(), - "data": kwargs.get("entry", {"info": {}}).info.__dict__, - } - - result = await asyncio.wait_for(self._worker_co(*args, **kwargs), timeout=self._max_task_time) - - if future: - future.set_result(result) - except (KeyboardInterrupt, MemoryError, SystemExit) as e: - if future: - future.set_exception(e) - self._exceptions = True - raise - except Exception as e: - self._exceptions = True - if future: - future.set_exception(e) - else: - self._logger.exception(f"Worker call failed. {e!s}") - finally: - self._status[worker_id] = None - if not is_terminator and from_queue: - self._queue.task_done() - - return True - - def has_open_workers(self) -> bool: - """ - Check if there are open workers. - - Returns: - bool: True if there are open workers, False otherwise. - - """ - return self.get_available_workers() > 0 - - def get_available_workers(self) -> int: - """ - Get the number of available workers. - - Returns: - int: number of available workers. - - """ - return sum(1 for worker_status in self._status.values() if worker_status is None) - - def get_workers_status(self) -> dict[str, datetime | None]: - """ - Get the status of all workers. - - Returns: - dict: dictionary of worker status. - - """ - return self._status - - async def __aenter__(self): - self.start() - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - await self.join() - - async def push(self, *args, is_temp: bool = False, **kwargs) -> asyncio.Future: - """ - Push work to the worker_co. If is_temp is True, a temporary worker is spawned - - Args: - args: position arguments to be passed to `worker_co` - is_temp (bool): flag to indicate if the job is temporary, default is False - kwargs (dict): keyword arguments to be passed to `worker_co` - - Returns: - asyncio.Future: future of result. - - """ - future = asyncio.futures.Future(loop=self._loop) if self._return_futures else None - - if is_temp: - temp_worker_id = f"temp_worker_{uuid.uuid4()!s}" - self._logger.info(f"Creating temporary worker '{temp_worker_id}'.") - # Spawn a temporary worker that processes this one job. - task = asyncio.ensure_future( - self._process_item(temp_worker_id, (future, args, kwargs), from_queue=False), loop=self._loop - ) - self._workers[temp_worker_id] = task - - # Attach a callback to clean up temporary worker entries when done. - def cleanup_temp_worker(_): - self._workers.pop(temp_worker_id, None) - self._status.pop(temp_worker_id, None) - self._logger.info(f"Temporary worker '{temp_worker_id}' has terminated.") - - task.add_done_callback(cleanup_temp_worker) - - return future - - await self._queue.put((future, args, kwargs)) - self._logger.debug(f"'{self._name}' pool has received a new job. {args} {kwargs}") - return future - - def start(self): - """ - Will start up worker pool and reset exception state - """ - self._exceptions = False - - for worker_number in range(self._num_workers): - worker_id = f"worker_{worker_number+1}" - self._create_worker(worker_id) - - async def restart(self, worker_id: str, msg: str | None = None) -> bool: - """ - Will restart the worker pool - - Args: - worker_id (str): worker id to restart - msg (str|None): message to send to the worker, default is None - - Returns: - bool: True if worker was restarted. - - """ - if worker_id not in self._workers: - self._logger.warning(f"Worker {worker_id} does not exist.") - return False - - try: - self._workers[worker_id].cancel(msg) - await self._workers[worker_id] - except asyncio.exceptions.CancelledError as e: - self._logger.warning(f"Worker {worker_id} restarted. {e!s}") - if worker_id in self._status: - self._status.pop(worker_id) - - if worker_id in self._workers: - self._workers.pop(worker_id) - - self._create_worker(worker_id) - - return True - - async def on_shutdown(self, _) -> None: - """ - Will shutdown the worker pool - """ - try: - await asyncio.wait_for(asyncio.gather(*[self.stop(worker_id) for worker_id in self._workers]), timeout=10) - except Exception: - self._logger.error(f"Exception shutting down {self._name}") - - async def stop(self, worker_id: str, msg: str | None = None) -> bool: - """ - Will stop the worker - - Args: - worker_id (str): worker id to stop - msg (str|None): message to send to the worker, default is None - - Returns: - bool: True if worker was stopped. - - """ - if worker_id not in self._workers: - self._logger.warning(f"Worker {worker_id} does not exist.") - return False - - if self._workers[worker_id].cancel(msg): - try: - await self._workers[worker_id] - except asyncio.exceptions.CancelledError as e: - self._logger.warning(f"Worker {worker_id} stopped. {e!s}") - if worker_id in self._status: - self._status.pop(worker_id) - - if worker_id in self._workers: - self._workers.pop(worker_id) - - return True - - async def join(self) -> None: - # no-op if workers aren't running - if len(self._workers) < 1: - return - - self._logger.info(f"Joining {self._name}") - - # The Terminators will kick each worker from being blocked against the _queue.get() and allow - # each one to exit. - for _ in range(self._num_workers): - await self._queue.put(Terminator()) - - try: - await asyncio.gather(*list(self._workers)) - self._workers = {} - except: - self._logger.exception(f"Exception joining {self._name}") - raise - finally: - self._logger.info(f"Completed {self._name}") - - if self._exceptions and self._raise_on_join: - msg = f"Exception occurred in {self._name} pool" - raise Exception(msg) - - def _create_worker(self, worker_id: str) -> asyncio.Future: - if worker_id in self._workers: - self._logger.debug(f"Worker {worker_id} already exists.") - return self._workers[worker_id] - - self._status[worker_id] = None - self._workers[worker_id] = asyncio.ensure_future( - coro_or_future=self._worker_loop(worker_id=worker_id), loop=self._loop - ) - - self._logger.debug(f"Created {worker_id}") - - return self._workers[worker_id] - - def _time(self) -> datetime: - return datetime.now(tz=UTC) diff --git a/app/library/Download.py b/app/library/Download.py index 958de882..59731734 100644 --- a/app/library/Download.py +++ b/app/library/Download.py @@ -11,7 +11,6 @@ from pathlib import Path import yt_dlp -from .AsyncPool import Terminator from .config import Config from .Events import EventBus, Events from .ffprobe import ffprobe @@ -20,6 +19,10 @@ from .Utils import delete_dir, extract_info, extract_ytdlp_logs, load_cookies from .YTDLPOpts import YTDLPOpts +class Terminator: + pass + + class NestedLogger: debug_messages = ["[debug] ", "[download] "] diff --git a/app/library/DownloadQueue.py b/app/library/DownloadQueue.py index 35285245..89a6b535 100644 --- a/app/library/DownloadQueue.py +++ b/app/library/DownloadQueue.py @@ -15,7 +15,6 @@ from aiohttp import web from app.library.ag_utils import ag -from .AsyncPool import AsyncPool from .conditions import Conditions from .config import Config from .DataStore import DataStore @@ -62,9 +61,6 @@ class DownloadQueue(metaclass=Singleton): event: asyncio.Event """Event to signal the download queue to start downloading.""" - pool: AsyncPool | None = None - """Pool of workers to download the files.""" - _active: dict[str, Download] = {} """Dictionary of active downloads.""" @@ -77,6 +73,12 @@ class DownloadQueue(metaclass=Singleton): done: DataStore """DataStore for the completed downloads.""" + workers: asyncio.Semaphore + """Semaphore to limit the number of concurrent downloads.""" + + processors: asyncio.Semaphore + """Semaphore to limit the number of concurrent processors.""" + def __init__(self, connection: Connection, config: Config | None = None): DownloadQueue._instance = self @@ -89,6 +91,8 @@ class DownloadQueue(metaclass=Singleton): self.paused = asyncio.Event() self.paused.set() self.event = asyncio.Event() + self.workers = asyncio.Semaphore(self.config.max_workers) + self.processors = asyncio.Semaphore(self.config.playlist_items_concurrency) @staticmethod def get_instance(): @@ -127,14 +131,6 @@ class DownloadQueue(metaclass=Singleton): ) # app.on_shutdown.append(self.on_shutdown) - # async def close_pool(_: web.Application): - # try: - # await self.pool.on_shutdown(_) - # except Exception as e: - # LOG.error(f"Failed to cleanup download pool. {e!s}") - - # app.on_cleanup.append(close_pool) - async def test(self) -> bool: """ Test the datastore connection to the database. @@ -213,8 +209,6 @@ class DownloadQueue(metaclass=Singleton): playlistCount = int(entry.get("playlist_count", len(entries))) results = [] - semaphore = asyncio.Semaphore(self.config.playlist_items_concurrency) - async def process_entry(i, etr): extras = { "playlist": entry.get("id"), @@ -229,7 +223,7 @@ class DownloadQueue(metaclass=Singleton): if "thumbnail" not in etr and "youtube:" in entry.get("extractor", ""): extras["thumbnail"] = f"https://img.youtube.com/vi/{etr['id']}/maxresdefault.jpg" - async with semaphore: + async with self.processors: LOG.debug(f"Processing entry {i}/{playlistCount} - ID: {etr.get('id')} - Title: {etr.get('title')}") return await self.add( item=item.new_with(url=etr.get("url") or etr.get("webpage_url"), extras=extras), @@ -708,56 +702,41 @@ class DownloadQueue(metaclass=Singleton): async def _download_pool(self) -> None: """ Create a pool of workers to download the files. - - Returns: - None - """ - self.pool = AsyncPool( - loop=asyncio.get_running_loop(), - num_workers=self.config.max_workers, - worker_co=self._download_file, - name="download_pool", - logger=logging.getLogger("WorkerPool"), - ) - - self.pool.start() - - lastLog = time.time() - while True: - while True: - if self.pool.has_open_workers() is True: - break - if self.config.max_workers > 1 and time.time() - lastLog > 600: - lastLog = time.time() - LOG.info("Waiting for worker to be free.", extra={"workers": self.pool.get_available_workers()}) - await asyncio.sleep(1) - while not self.queue.has_downloads(): - LOG.info(f"Waiting for item to download. '{self.pool.get_available_workers()}' free workers.") - if self.event: - await self.event.wait() - self.event.clear() - LOG.debug("Cleared wait event.") + LOG.info("Waiting for item to download.") + await self.event.wait() + self.event.clear() - if self.paused and isinstance(self.paused, asyncio.Event) and self.is_paused(): + if self.is_paused(): LOG.info("Download pool is paused.") await self.paused.wait() LOG.info("Download pool resumed downloading.") - entry = self.queue.get_next_download() - await asyncio.sleep(0.2) + for _id, entry in list(self.queue.items()): + if entry.started() or entry.is_cancelled(): + continue - if entry is None: - continue + # Bypass semaphore for temporary workers + if entry.is_live: + asyncio.create_task(self._download_temp(_id, entry)) + else: + asyncio.create_task(self._download(_id, entry)) - LOG.debug(f"Pushing {entry=} to executor.") + await asyncio.sleep(0.5) - if entry.started() is False and entry.is_cancelled() is False: - await self.pool.push(is_temp=entry.is_live, id=entry.info._id, entry=entry) - LOG.debug(f"Pushed {entry=} to executor.") - await asyncio.sleep(1) + async def _download(self, _id: str, entry: Download) -> None: + async with self.workers: + await self._download_file(_id, entry) + + async def _download_temp(self, _id: str, entry: Download) -> None: + LOG.info(f"Creating temporary worker for entry '{entry.info.name()}'.") + + try: + await self._download_file(_id, entry) + finally: + LOG.info(f"Temporary worker for '{entry.info.name()}' completed.") async def _download_file(self, id: str, entry: Download) -> None: """ @@ -842,52 +821,6 @@ class DownloadQueue(metaclass=Singleton): LOG.error(f"Failed to cancel staled item '{item_ref}'. {e!s}") LOG.exception(e) - if self.pool: - time_now = datetime.now(tz=UTC) - workers = self.pool.get_workers_status() - if len(workers) > 0: - LOG.debug(f"Checking for stale workers. {len(workers)} workers found.") - - for worker_id in workers: - worker = workers.get(worker_id, {}) - if not worker: - continue - - started = worker.get("started", None) - if not started: - LOG.debug(f"Worker '{worker_id}' not working yet.") - continue - - started = datetime.fromisoformat(started) - - if time_now - started < timedelta(minutes=5): - LOG.debug(f"Worker '{worker_id}' is not consider stale yet.") - continue - - data = worker.get("data", {}) - - status = data.get("status", None) - if "preparing" != status: - LOG.debug(f"Worker '{worker_id}' not stuck. Status '{status}'.") - continue - - _id = data.get("data._id", None) - if not _id: - LOG.debug(f"Worker '{worker_id}' has no id.") - continue - - id = data.get("id", None) - title = data.get("title", None) - - item_ref = f"{_id=} {id=} {title=}" - - LOG.warning(f"Cancelling staled item '{item_ref}' from worker pool.") - try: - await self.cancel([_id]) - except Exception as e: - LOG.error(f"Failed to cancel staled item '{item_ref}' from worker pool. {e!s}") - LOG.exception(e) - async def _check_live(self): """ Monitor the queue for items marked as live events and queue them when time is reached. diff --git a/app/routes/api/dev.py b/app/routes/api/dev.py index 06dc4e81..c53fe397 100644 --- a/app/routes/api/dev.py +++ b/app/routes/api/dev.py @@ -1,11 +1,9 @@ import asyncio -import json import logging from aiohttp import web -from aiohttp.web import Request, Response +from aiohttp.web import Response -from app.library import DownloadQueue from app.library.config import Config from app.library.encoder import Encoder from app.library.router import route @@ -42,139 +40,3 @@ async def debug_asyncio(config: Config, encoder: Encoder) -> Response: dumps=encoder.encode, ) - -@route("GET", "api/dev/workers/", "pool_list") -async def pool_list(config: Config, queue: DownloadQueue) -> Response: - """ - Get the workers status. - - Args: - config (Config): The configuration object. - queue (DownloadQueue): The download queue object. - - Returns: - Response: The response object. - - """ - if not config.is_dev(): - return web.json_response( - {"error": "This endpoint is only available in development mode."}, - status=web.HTTPNotFound.status_code, - ) - - if queue.pool is None: - return web.json_response({"error": "Worker pool not initialized."}, status=web.HTTPNotFound.status_code) - - status = queue.pool.get_workers_status() - - data = [] - - for worker in status: - worker_status = status.get(worker) - data.append( - { - "id": worker, - "data": {"status": "Waiting for download."} if worker_status is None else worker_status, - } - ) - - return web.json_response( - data={ - "open": queue.pool.has_open_workers(), - "count": queue.pool.get_available_workers(), - "workers": data, - }, - status=web.HTTPOk.status_code, - dumps=lambda obj: json.dumps(obj, default=lambda o: f"<>"), - ) - - -@route("POST", "api/dev/workers/", "pool_start") -async def pool_restart(config: Config, queue: DownloadQueue) -> Response: - """ - Restart the workers pool. - - Args: - config (Config): The configuration object. - queue (DownloadQueue): The download queue object. - - Returns: - Response: The response object. - - """ - if not config.is_dev(): - return web.json_response( - {"error": "This endpoint is only available in development mode."}, - status=web.HTTPNotFound.status_code, - ) - - if queue.pool is None: - return web.json_response({"error": "Worker pool not initialized."}, status=web.HTTPNotFound.status_code) - - queue.pool.start() - - return web.json_response({"message": "Workers pool being restarted."}, status=web.HTTPOk.status_code) - - -@route("PATCH", "api/dev/workers/{id}", "worker_restart") -async def worker_restart(request: Request, config: Config, queue: DownloadQueue) -> Response: - """ - Restart a worker. - - Args: - request (Request): The request object. - config (Config): The configuration object. - queue (DownloadQueue): The download queue object. - - Returns: - Response: The response object - - """ - if not config.is_dev(): - return web.json_response( - {"error": "This endpoint is only available in development mode."}, - status=web.HTTPNotFound.status_code, - ) - - worker_id: str = request.match_info.get("id") - if not worker_id: - return web.json_response({"error": "worker id is required."}, status=web.HTTPBadRequest.status_code) - - if queue.pool is None: - return web.json_response({"error": "Worker pool not initialized."}, status=web.HTTPNotFound.status_code) - - status = await queue.pool.restart(worker_id, "requested by user.") - - return web.json_response({"status": "restarted" if status else "in_error_state"}, status=web.HTTPOk.status_code) - - -@route("DELETE", "api/dev/workers/{id}", "worker_stop") -async def worker_stop(request: Request, config: Config, queue: DownloadQueue) -> Response: - """ - Stop a worker. - - Args: - request (Request): The request object. - config (Config): The configuration object. - queue (DownloadQueue): The download queue object. - - Returns: - Response: The response object. - - """ - if not config.is_dev(): - return web.json_response( - {"error": "This endpoint is only available in development mode."}, - status=web.HTTPNotFound.status_code, - ) - - worker_id: str = request.match_info.get("id") - if not worker_id: - raise web.HTTPBadRequest(text="worker id is required.") - - if queue.pool is None: - return web.json_response({"error": "Worker pool not initialized."}, status=web.HTTPNotFound.status_code) - - status = await queue.pool.stop(worker_id, "requested by user.") - - return web.json_response({"status": "stopped" if status else "in_error_state"}, status=web.HTTPOk.status_code)