Remove the redundant usage of AsyncPool and replace it with asyncio semaphore

app/library/AsyncPool.py

@@ -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)

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] "]
 

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.")
+                LOG.info("Waiting for item to download.")
-                if self.event:
+                await self.event.wait()
-                    await self.event.wait()
+                self.event.clear()
-                    self.event.clear()
-                    LOG.debug("Cleared wait event.")
 
-            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()
+            for _id, entry in list(self.queue.items()):
-            await asyncio.sleep(0.2)
+                if entry.started() or entry.is_cancelled():
+                    continue
 
-            if entry is None:
+                # Bypass semaphore for temporary workers
-                continue
+                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:
+    async def _download(self, _id: str, entry: Download) -> None:
-                await self.pool.push(is_temp=entry.is_live, id=entry.info._id, entry=entry)
+        async with self.workers:
-                LOG.debug(f"Pushed {entry=} to executor.")
+            await self._download_file(_id, entry)
-                await asyncio.sleep(1)
+
+    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.

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"<<non-serializable: {type(o).__qualname__}>>"),
-    )
-
-
-@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)