Remove the redundant usage of AsyncPool and replace it with asyncio semaphore
This commit is contained in:
arabcoders
parent
a81ee127e9
commit
bd695e7067
4 changed files with 38 additions and 560 deletions
|
|
@ -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)
|
||||
|
|
@ -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] "]
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in a new issue