From 2493d41271455cf697881578971c040ca13b35c2 Mon Sep 17 00:00:00 2001 From: ArabCoders Date: Wed, 20 Mar 2024 23:27:09 +0300 Subject: [PATCH 1/2] update the format for the workers output. --- app/main.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/app/main.py b/app/main.py index 7c595f77..9c8066e5 100644 --- a/app/main.py +++ b/app/main.py @@ -364,7 +364,11 @@ class Main: def default(o): return f"<>" return json.dumps(obj, default=default) - return web.Response(text=safe_serialize(data), headers={ + return web.Response(text=safe_serialize({ + 'open': self.queue.pool.has_open_workers(), + 'count': self.queue.pool.get_available_workers(), + 'workers': data, + }), headers={ 'Content-Type': 'application/json', }) From a1aa7b02ba569912ca208b120bb603d89a22f1d2 Mon Sep 17 00:00:00 2001 From: ArabCoders Date: Sat, 23 Mar 2024 01:02:53 +0300 Subject: [PATCH 2/2] wrap notifier in timeout call, and expose pool operations via API --- app/AsyncPool.py | 93 ++++++++++++++++++++++++++++++++++++------------ app/Utils.py | 11 +++--- app/main.py | 35 ++++++++++++++++++ 3 files changed, 113 insertions(+), 26 deletions(-) diff --git a/app/AsyncPool.py b/app/AsyncPool.py index 3ed26a78..852659b3 100644 --- a/app/AsyncPool.py +++ b/app/AsyncPool.py @@ -39,20 +39,19 @@ class AsyncPool: self._num_workers = num_workers self._logger = logger if logger else logging.getLogger(__name__) self._queue = asyncio.Queue(num_workers * load_factor) - self._workers = None + 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._active_worker: dict[str, dict | None] = {} + self._status: dict[str, dict | None] = {} async def _worker_loop(self, worker_id: str): while True: got_obj = False future = None - try: item = await self._queue.get() got_obj = True @@ -62,7 +61,7 @@ class AsyncPool: future, args, kwargs = item - self._active_worker[worker_id] = { + self._status[worker_id] = { 'started': self._time().isoformat(), 'data': kwargs.get('entry', {'info': {}}).info.__dict__, } @@ -79,6 +78,9 @@ class AsyncPool: self._exceptions = True raise except BaseException as e: + if isinstance(e, asyncio.exceptions.CancelledError): + raise e + self._exceptions = True if future: @@ -87,7 +89,7 @@ class AsyncPool: else: self._logger.exception(f'Worker call failed. {str(e)}') finally: - self._active_worker[worker_id] = None + self._status[worker_id] = None if got_obj: self._queue.task_done() @@ -106,13 +108,13 @@ class AsyncPool: """ :return: number of available workers. """ - return sum(1 for worker_status in self._active_worker.values() if worker_status is None) + return sum(1 for worker_status in self._status.values() if worker_status is None) def get_workers_status(self) -> dict[str, datetime | None]: """ :return: dictionary of worker status. """ - return self._active_worker + return self._status async def __aenter__(self): self.start() @@ -139,23 +141,55 @@ class AsyncPool: def start(self): """ Will start up worker pool and reset exception state """ - assert self._workers is None self._exceptions = False - self._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(worker_id=workerName), - loop=self._loop - ) - ) - self._active_worker[workerName] = None + worker_id = f"worker_{worker_number+1}" + self._createWorker(worker_id) - async def join(self): + async def restart(self, worker_id: str, msg: str = None) -> bool: + """ Will restart the worker pool """ + 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. {str(e)}') + if worker_id in self._status: + self._status.pop(worker_id) + + if worker_id in self._workers: + self._workers.pop(worker_id) + + self._createWorker(worker_id) + + return True + + async def stop(self, worker_id: str, msg: str = None) -> bool: + """ Will stop the worker """ + 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. {str(e)}') + 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 not self._workers: + if len(self._workers) < 1: return self._logger.info(f'Joining {self._name}') @@ -166,8 +200,8 @@ class AsyncPool: await self._queue.put(Terminator()) try: - await asyncio.gather(*self._workers) - self._workers = None + await asyncio.gather(*list(self._workers)) + self._workers = {} except: self._logger.exception(f'Exception joining {self._name}') raise @@ -177,5 +211,20 @@ class AsyncPool: if self._exceptions and self._raise_on_join: raise Exception(f"Exception occurred in {self._name} pool") - def _time(self): + def _createWorker(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=timezone.utc) diff --git a/app/Utils.py b/app/Utils.py index dfaf2862..5e8f9901 100644 --- a/app/Utils.py +++ b/app/Utils.py @@ -11,10 +11,10 @@ from Webhooks import Webhooks LOG = logging.getLogger('Utils') -AUDIO_FORMATS: tuple = ('m4a', 'mp3', 'opus', 'wav') -IGNORED_KEYS: tuple = ('cookiefile', 'paths', 'outtmpl', 'progress_hooks', 'postprocessor_hooks',) +AUDIO_FORMATS: tuple[str] = ('m4a', 'mp3', 'opus', 'wav',) +IGNORED_KEYS: tuple[str] = ('cookiefile', 'paths', 'outtmpl', 'progress_hooks', 'postprocessor_hooks',) -YTDLP_INFO_CLS = None +YTDLP_INFO_CLS: yt_dlp.YoutubeDL = None def get_format(format: str, quality: str) -> str: @@ -347,4 +347,7 @@ class Notifier: if self.webhooks and event in self.webhooks_allowed_events: tasks.append(self.webhooks.send(event, data)) - await asyncio.gather(*tasks) + try: + await asyncio.wait_for(asyncio.gather(*tasks), timeout=60) + except asyncio.TimeoutError: + LOG.error(f"Timed out sending event {event} to webhooks.") diff --git a/app/main.py b/app/main.py index 9c8066e5..19a87f75 100644 --- a/app/main.py +++ b/app/main.py @@ -372,6 +372,41 @@ class Main: 'Content-Type': 'application/json', }) + @self.routes.post(self.config.url_prefix + 'workers') + async def restart_pool(_) -> Response: + if self.queue.pool is None: + return web.json_response({"error": "Worker pool not initialized."}, status=404) + + status = self.queue.pool.start() + + return web.Response() + + @self.routes.patch(self.config.url_prefix + 'workers/{id}') + async def restart_worker(request: Request) -> Response: + id: str = request.match_info.get('id') + if not id: + raise web.HTTPBadRequest(reason='worker id is required.') + + if self.queue.pool is None: + return web.json_response({"error": "Worker pool not initialized."}, status=404) + + status = await self.queue.pool.restart(id, 'requested by user.') + + return web.json_response({"status": "restarted" if status else "in_error_state"}) + + @self.routes.delete(self.config.url_prefix + 'workers/{id}') + async def stop_worker(request: Request) -> Response: + id: str = request.match_info.get('id') + if not id: + raise web.HTTPBadRequest(reason='worker id is required.') + + if self.queue.pool is None: + return web.json_response({"error": "Worker pool not initialized."}, status=404) + + status = await self.queue.pool.stop(id, 'requested by user.') + + return web.json_response({"status": "stopped" if status else "in_error_state"}) + @self.routes.get(self.config.url_prefix + 'm3u8/{file:.*}') async def m3u8(request: Request) -> Response: file: str = request.match_info.get('file')