252 lines
9 KiB
Python
252 lines
9 KiB
Python
"""Download pool management - worker coordination and execution."""
|
|
|
|
import asyncio
|
|
import logging
|
|
from typing import TYPE_CHECKING
|
|
|
|
from app.library.Events import EventBus, Events
|
|
from app.library.Utils import calc_download_path
|
|
|
|
from .core import Download
|
|
from .utils import get_extractor_limit, handle_task_exception
|
|
|
|
if TYPE_CHECKING:
|
|
from app.library.config import Config
|
|
|
|
from .queue_manager import DownloadQueue
|
|
|
|
LOG: logging.Logger = logging.getLogger("downloads.pool")
|
|
|
|
|
|
class PoolManager:
|
|
"""
|
|
Manages download worker pool and execution.
|
|
|
|
Coordinates concurrent downloads with semaphore-based limits:
|
|
- Global worker limit (max_workers)
|
|
- Per-extractor limits (max_workers_per_extractor)
|
|
- Live streams bypass all limits
|
|
|
|
Integrates with DownloadQueue for queue access and configuration.
|
|
"""
|
|
|
|
def __init__(self, queue: "DownloadQueue", config: "Config"):
|
|
self.queue = queue
|
|
self.config = config
|
|
self._notify: EventBus = EventBus.get_instance()
|
|
|
|
self.workers = asyncio.Semaphore(config.max_workers)
|
|
"Semaphore to limit the number of concurrent downloads."
|
|
|
|
self.paused = asyncio.Event()
|
|
"Event to pause the download queue."
|
|
|
|
self.event = asyncio.Event()
|
|
"Event to signal the download queue to start downloading."
|
|
|
|
self._active: dict[str, Download] = {}
|
|
"Dictionary of active downloads."
|
|
|
|
self.paused.set()
|
|
|
|
def is_paused(self) -> bool:
|
|
return not self.paused.is_set()
|
|
|
|
def pause(self) -> None:
|
|
if self.paused.is_set():
|
|
self.paused.clear()
|
|
LOG.info("Download pool paused.")
|
|
|
|
def resume(self) -> None:
|
|
if not self.paused.is_set():
|
|
self.paused.set()
|
|
LOG.info("Download pool resumed.")
|
|
|
|
def trigger_download(self) -> None:
|
|
"""Trigger the download pool to check for new items."""
|
|
if self.event:
|
|
self.event.set()
|
|
|
|
def get_active_downloads(self) -> dict[str, Download]:
|
|
return self._active.copy()
|
|
|
|
def get_active_download(self, download_id: str) -> Download | None:
|
|
return self._active.get(download_id)
|
|
|
|
async def start_pool(self) -> None:
|
|
asyncio.create_task(self._download_pool(), name="download_pool")
|
|
|
|
async def shutdown(self) -> None:
|
|
if self._active:
|
|
LOG.info(
|
|
"Cancelling %s active download(s).",
|
|
len(self._active),
|
|
extra={"active_count": len(self._active), "download_ids": list(self._active)},
|
|
)
|
|
await self.queue.cancel(list(self._active.keys()))
|
|
|
|
async def _download_pool(self) -> None:
|
|
adaptive_sleep = 0.2
|
|
max_sleep = 5.0
|
|
|
|
while True:
|
|
while not self.queue.queue.has_downloads():
|
|
LOG.info("Waiting for item to download.")
|
|
await self.event.wait()
|
|
self.event.clear()
|
|
adaptive_sleep = 0.2
|
|
|
|
if self.is_paused():
|
|
LOG.warning("Download pool is paused.")
|
|
await self.paused.wait()
|
|
LOG.info("Download pool resumed downloading.")
|
|
adaptive_sleep = 0.2
|
|
|
|
items_processed = 0
|
|
|
|
for _id, entry in list(self.queue.queue.items()):
|
|
if entry.started() or entry.is_cancelled() or entry.info.auto_start is False:
|
|
continue
|
|
|
|
extractor: str = entry.info.get_extractor() or "unknown"
|
|
|
|
# Live downloads bypass all limits.
|
|
if entry.is_live:
|
|
task: asyncio.Task[None] = asyncio.create_task(
|
|
self._download_file(_id, entry), name=f"download_live_{extractor}_{_id}"
|
|
)
|
|
task.add_done_callback(lambda t: handle_task_exception(t, LOG))
|
|
items_processed += 1
|
|
else:
|
|
_limit: asyncio.Semaphore = get_extractor_limit(
|
|
extractor, self.config.max_workers, self.config.max_workers_per_extractor, LOG
|
|
)
|
|
|
|
# Skip this item in this iteration if no slots are available.
|
|
if self.workers.locked() or _limit.locked():
|
|
continue
|
|
|
|
await self.workers.acquire()
|
|
await _limit.acquire()
|
|
|
|
task: asyncio.Task[None] = asyncio.create_task(
|
|
self._download_file(_id, entry), name=f"download_file_{extractor}_{_id}"
|
|
)
|
|
|
|
def _release(t: asyncio.Task, sem=_limit) -> None:
|
|
sem.release()
|
|
self.workers.release()
|
|
handle_task_exception(t, LOG)
|
|
|
|
task.add_done_callback(_release)
|
|
items_processed += 1
|
|
|
|
# No items could be processed, back off a bit to avoid busy-waiting.
|
|
if 0 == items_processed:
|
|
adaptive_sleep: float = min(adaptive_sleep * 1.5, max_sleep)
|
|
LOG.debug(
|
|
"No download slots available; backing off for %.2f seconds.",
|
|
adaptive_sleep,
|
|
extra={"backoff_s": round(adaptive_sleep, 2), "active_count": len(self._active)},
|
|
)
|
|
else:
|
|
adaptive_sleep = 0.2
|
|
|
|
await asyncio.sleep(adaptive_sleep)
|
|
|
|
async def _download_file(self, id: str, entry: Download) -> None:
|
|
"""
|
|
Download the file.
|
|
|
|
Args:
|
|
id: The id of the download
|
|
entry: The download entry
|
|
|
|
"""
|
|
filePath: str = calc_download_path(base_path=self.config.download_path, folder=entry.info.folder)
|
|
LOG.info(
|
|
"Downloading '%s' from '%s' to '%s'.",
|
|
entry.info.title,
|
|
entry.info.url,
|
|
filePath,
|
|
extra={
|
|
"download": {
|
|
"download_id": entry.id,
|
|
"item_id": entry.info.id,
|
|
"title": entry.info.title,
|
|
"url": entry.info.url,
|
|
"download_dir": filePath,
|
|
"preset": entry.info.preset,
|
|
"is_live": entry.is_live,
|
|
}
|
|
},
|
|
)
|
|
|
|
try:
|
|
self._active[entry.info._id] = entry
|
|
await entry.start()
|
|
|
|
if entry.info.status not in ("finished", "skip", "cancelled"):
|
|
if not entry.info.error:
|
|
entry.info.error = f"Download ended with unexpected status '{entry.info.status}'."
|
|
entry.info.status = "error"
|
|
except Exception as e:
|
|
entry.info.status = "error"
|
|
entry.info.error = str(e)
|
|
finally:
|
|
if entry.info._id in self._active:
|
|
self._active.pop(entry.info._id, None)
|
|
|
|
await entry.close()
|
|
|
|
if await self.queue.queue.exists(key=id):
|
|
LOG.debug(
|
|
"Removing completed download '%s' from queue.",
|
|
entry.info.title,
|
|
extra={
|
|
"download": {
|
|
"download_id": id,
|
|
"item_id": entry.info.id,
|
|
"title": entry.info.title,
|
|
"status": entry.info.status,
|
|
}
|
|
},
|
|
)
|
|
await self.queue.queue.delete(key=id)
|
|
|
|
nTitle: str | None = None
|
|
nMessage: str | None = None
|
|
finalized_live = entry.is_live and entry.info.status == "finished" and bool(entry.info.filename)
|
|
|
|
if finalized_live:
|
|
nTitle = "Download Completed"
|
|
nMessage = f"Completed '{entry.info.title}' download."
|
|
if entry.info.is_archivable and not entry.info.is_archived:
|
|
entry.info.is_archived = True
|
|
|
|
self._notify.emit(Events.ITEM_COMPLETED, data=entry.info, title=nTitle, message=nMessage)
|
|
elif entry.is_cancelled() is True:
|
|
nTitle = "Download Cancelled"
|
|
nMessage = f"Cancelled '{entry.info.title}' download."
|
|
self._notify.emit(Events.ITEM_CANCELLED, data=entry.info, title=nTitle, message=nMessage)
|
|
entry.info.status = "cancelled"
|
|
elif entry.info.status == "finished" and entry.info.filename:
|
|
nTitle = "Download Completed"
|
|
nMessage = f"Completed '{entry.info.title}' download."
|
|
if entry.info.is_archivable and not entry.info.is_archived:
|
|
entry.info.is_archived = True
|
|
|
|
self._notify.emit(Events.ITEM_COMPLETED, data=entry.info, title=nTitle, message=nMessage)
|
|
|
|
await self.queue.done.put(entry)
|
|
self._notify.emit(
|
|
Events.ITEM_MOVED,
|
|
data={"to": "history", "preset": entry.info.preset, "item": entry.info},
|
|
title=nTitle,
|
|
message=nMessage,
|
|
)
|
|
else:
|
|
LOG.warning("Completed download '%s' was not found in queue.", id, extra={"download_id": id})
|
|
|
|
if self.event:
|
|
self.event.set()
|