diff --git a/app/AsyncPool.py b/app/AsyncPool.py index 7a76d1fc..3ed26a78 100644 --- a/app/AsyncPool.py +++ b/app/AsyncPool.py @@ -1,5 +1,6 @@ import asyncio from datetime import datetime, timezone +import logging class Terminator: @@ -7,9 +8,13 @@ class Terminator: class AsyncPool: - def __init__(self, loop, num_workers: int, name: str, logger, worker_co, load_factor: int = 1, - job_accept_duration: int = None, max_task_time: int = None, return_futures: bool = False, - raise_on_join: bool = False, log_every_n: int = None, expected_total=None): + def __init__(self, + num_workers: int, worker_co, + name: str, logger: logging.Logger, + loop: asyncio.AbstractEventLoop = None, + load_factor: int = 1, max_task_time: int = 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 @@ -23,37 +28,27 @@ class AsyncPool: :param logger: logger to use :param worker_co: async coroutine to call when an item is retrieved from the queue :param load_factor: multiplier used for number of items in queue - :param job_accept_duration: maximum number of seconds from first push to last push before a TimeoutError will be thrown. - Set to None for no limit. Note this does not get reset on aenter/aexit. :param max_task_time: maximum time allowed for each task before a CancelledError is raised in the task. Set to None for no limit. :param return_futures: set to reture to return a future for each `push` (imposes CPU overhead) :param raise_on_join: raise on join if any exceptions have occurred, default is False - :param log_every_n: (optional) set to number of `push`s each time a log statement should be printed (default does not print every-n pushes) - :param expected_total: (optional) expected total number of jobs (used for `log_event_n` logging) :return: instance of AsyncWorkerPool """ - loop = loop if loop else asyncio.get_event_loop() - self._loop = loop + self._loop = loop if loop else asyncio.get_event_loop() self._num_workers = num_workers - self._logger = logger + self._logger = logger if logger else logging.getLogger(__name__) self._queue = asyncio.Queue(num_workers * load_factor) self._workers = None self._exceptions = False - self._job_accept_duration = job_accept_duration - self._first_push_dt = None 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._total_queued = 0 - self._log_every_n = log_every_n - self._expected_total = expected_total - self._active_workers = 0 + self._active_worker: dict[str, dict | None] = {} - async def _worker_loop(self): + async def _worker_loop(self, worker_id: str): while True: got_obj = False future = None @@ -65,11 +60,16 @@ class AsyncPool: if item.__class__ is Terminator: break - self._active_workers += 1 future, args, kwargs = item + + self._active_worker[worker_id] = { + 'started': self._time().isoformat(), + 'data': kwargs.get('entry', {'info': {}}).info.__dict__, + } + # the wait_for will cancel the task (task sees CancelledError) and raises a TimeoutError from here # so be wary of catching TimeoutErrors in this loop - result = await asyncio.wait_for(self._worker_co(*args, **kwargs), self._max_task_time) + result = await asyncio.wait_for(self._worker_co(*args, **kwargs), timeout=self._max_task_time) if future: future.set_result(result) @@ -85,9 +85,10 @@ class AsyncPool: # don't log the failure when the client is receiving the future future.set_exception(e) else: - self._logger.exception('Worker call failed') + self._logger.exception(f'Worker call failed. {str(e)}') finally: - self._active_workers -= 1 + self._active_worker[worker_id] = None + if got_obj: self._queue.task_done() @@ -95,21 +96,23 @@ class AsyncPool: def exceptions(self): return self._exceptions - @property - def total_queued(self): - return self._total_queued - def has_open_workers(self) -> bool: """ :return: True if there are open workers. """ - return self._active_workers < self._num_workers + return self.get_available_workers() > 0 def get_available_workers(self) -> int: """ :return: number of available workers. """ - return self._num_workers - self._active_workers + return sum(1 for worker_status in self._active_worker.values() if worker_status is None) + + def get_workers_status(self) -> dict[str, datetime | None]: + """ + :return: dictionary of worker status. + """ + return self._active_worker async def __aenter__(self): self.start() @@ -127,18 +130,8 @@ class AsyncPool: :return: future of result. """ - if self._first_push_dt is None: - self._first_push_dt = self._time() - - if self._job_accept_duration is not None and (self._time() - self._first_push_dt) > self._job_accept_duration: - raise TimeoutError(f"Max life time of {self._job_accept_duration}s exceeded for {self._name} pool.") - future = asyncio.futures.Future(loop=self._loop) if self._return_futures else None await self._queue.put((future, args, kwargs)) - self._total_queued += 1 - - if self._log_every_n is not None and (self._total_queued % self._log_every_n) == 0: - self._logger.info(f"pushed {self._total_queued}/{self._expected_total} items to {self._name} pool.") self._logger.debug(f"'{self._name}' pool has received a new job. {args} {kwargs}") @@ -150,13 +143,15 @@ class AsyncPool: self._exceptions = False self._workers = [] - for _ in range(self._num_workers): + for worker_number in range(self._num_workers): + workerName = f"worker_{worker_number+1}" self._workers.append( asyncio.ensure_future( - coro_or_future=self._worker_loop(), + coro_or_future=self._worker_loop(worker_id=workerName), loop=self._loop ) ) + self._active_worker[workerName] = None async def join(self): # no-op if workers aren't running @@ -166,7 +161,7 @@ class AsyncPool: 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 + # each one to exit. for _ in range(self._num_workers): await self._queue.put(Terminator()) @@ -183,5 +178,4 @@ class AsyncPool: raise Exception(f"Exception occurred in {self._name} pool") def _time(self): - # utcnow returns a naive datetime, so we have to set the timezone manually - return datetime.utcnow().replace(tzinfo=timezone.utc) + return datetime.now(tz=timezone.utc) diff --git a/app/DownloadQueue.py b/app/DownloadQueue.py index a6f7c207..087937a1 100644 --- a/app/DownloadQueue.py +++ b/app/DownloadQueue.py @@ -25,6 +25,7 @@ class DownloadQueue: queue: DataStore = None done: DataStore = None event: asyncio.Event = None + pool: AsyncPool = None def __init__(self, notifier: Notifier, connection: Connection, serializer: ObjectSerializer): self.config = Config.get_instance() @@ -122,6 +123,9 @@ class DownloadQueue: is_manifestless = entry.get('is_manifestless', False) options.update({'is_manifestless': is_manifestless}) + live_status: list = ['is_live', 'is_upcoming'] + is_live = entry.get('is_live', None) or live_in or entry.get('live_status', None) in live_status + dl = ItemDTO( id=entry.get('id'), title=entry.get('title'), @@ -134,7 +138,7 @@ class DownloadQueue: output_template=output_template if output_template else self.config.output_template, datetime=formatdate(time.time()), error=error, - is_live=entry.get('is_live', None) or entry.get('live_status', None) in ['is_live', 'is_upcoming'] or live_in, + is_live=is_live, live_in=live_in, options=options ) @@ -161,7 +165,7 @@ class DownloadQueue: debug=bool(self.config.ytdl_debug) ) - if dlInfo.info.live_in or 'is_upcoming' == entry.get('live_status',None): + if dlInfo.info.live_in or 'is_upcoming' == entry.get('live_status', None): dlInfo.info.status = 'not_live' itemDownload = self.done.put(dlInfo) NotifyEvent = 'completed' @@ -328,42 +332,45 @@ class DownloadQueue: return items async def __download_pool(self): - async with AsyncPool( + self.pool = AsyncPool( loop=asyncio.get_running_loop(), num_workers=self.config.max_workers, worker_co=self.__downloadFile, name='download_pool', logger=logging.getLogger('WorkerPool'), - ) as executor: - lastLog = time.time() + ) + self.pool.start() + + lastLog = time.time() + + while True: while True: - while True: - if executor.has_open_workers() is True: - break - if time.time() - lastLog > 120: - lastLog = time.time() - LOG.info(f'Waiting for worker to be free.') - await asyncio.sleep(1) + if self.pool.has_open_workers() is True: + break + if time.time() - lastLog > 120: + lastLog = time.time() + LOG.info(f'Waiting for worker to be free. {self.pool.get_workers_status()}') + await asyncio.sleep(1) - while not self.queue.hasDownloads(): - LOG.info(f"Waiting for item to download. '{executor.get_available_workers()}' free workers.") - await self.event.wait() - self.event.clear() - LOG.debug(f"Cleared wait event.") + while not self.queue.hasDownloads(): + LOG.info(f"Waiting for item to download. '{self.pool.get_available_workers()}' free workers.") + await self.event.wait() + self.event.clear() + LOG.debug(f"Cleared wait event.") - entry = self.queue.getNextDownload() - await asyncio.sleep(0.2) + entry = self.queue.getNextDownload() + await asyncio.sleep(0.2) - if entry is None: - continue + if entry is None: + continue - LOG.debug(f"Pushing {entry=} to executor.") + LOG.debug(f"Pushing {entry=} to executor.") - if entry.started() is False and entry.is_canceled() is False: - await executor.push(id=entry.info._id, entry=entry) - LOG.debug(f"Pushed {entry=} to executor.") - await asyncio.sleep(1) + if entry.started() is False and entry.is_canceled() is False: + await self.pool.push(id=entry.info._id, entry=entry) + LOG.debug(f"Pushed {entry=} to executor.") + await asyncio.sleep(1) async def __download(self): while True: diff --git a/app/Utils.py b/app/Utils.py index 2794ec30..dfaf2862 100644 --- a/app/Utils.py +++ b/app/Utils.py @@ -293,7 +293,9 @@ class ObjectSerializer(json.JSONEncoder): """ def default(self, obj): - return obj.__dict__ if isinstance(obj, object) else json.JSONEncoder.default(self, obj) + if isinstance(obj, object) and hasattr(obj, '__dict__'): + return obj.__dict__ + return json.JSONEncoder.default(self, obj) class Notifier: diff --git a/app/main.py b/app/main.py index a42cb92f..7c595f77 100644 --- a/app/main.py +++ b/app/main.py @@ -11,6 +11,7 @@ from Utils import ObjectSerializer, Notifier from aiohttp import web, client from aiohttp.web import Request, Response from Webhooks import Webhooks +from Download import Download from player.M3u8 import M3u8 from player.Segments import Segments import socketio @@ -342,6 +343,31 @@ class Main: return web.Response(text=self.serializer.encode(history)) + @self.routes.get(self.config.url_prefix + 'workers') + async def workers(_) -> Response: + if self.queue.pool is None: + return web.json_response({"error": "Worker pool not initialized."}, status=404) + + status = self.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, + }) + + def safe_serialize(obj): + def default(o): return f"<>" + return json.dumps(obj, default=default) + + return web.Response(text=safe_serialize(data), headers={ + 'Content-Type': 'application/json', + }) + @self.routes.get(self.config.url_prefix + 'm3u8/{file:.*}') async def m3u8(request: Request) -> Response: file: str = request.match_info.get('file')