120 lines
4.3 KiB
Python
120 lines
4.3 KiB
Python
"""Воркеры, разбирающие очередь загрузок.
|
|
|
|
Каждый воркер — отдельная asyncio-таска: берёт id из очереди, грузит запись из
|
|
Storage, качает через выбранный downloader, по ходу обновляет Storage и публикует
|
|
события в EventBus. Менеджер WorkerPool поднимает N таких воркеров и гасит их.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
|
|
from app.core.events import EventBus
|
|
from app.core.queue import DownloadQueue
|
|
from app.core.storage import Storage
|
|
from app.models import Download, DownloadEvent, DownloadStatus, _now
|
|
from app.services.downloader import Progress, pick_downloader
|
|
|
|
logger = logging.getLogger("pvideodl.worker")
|
|
|
|
|
|
class WorkerPool:
|
|
def __init__(
|
|
self,
|
|
queue: DownloadQueue,
|
|
storage: Storage,
|
|
events: EventBus,
|
|
worker_count: int,
|
|
) -> None:
|
|
self._queue = queue
|
|
self._storage = storage
|
|
self._events = events
|
|
self._count = worker_count
|
|
self._tasks: list[asyncio.Task[None]] = []
|
|
|
|
async def start(self) -> None:
|
|
self._tasks = [
|
|
asyncio.create_task(self._run(i), name=f"worker-{i}")
|
|
for i in range(self._count)
|
|
]
|
|
logger.info("Запущено воркеров: %d", self._count)
|
|
|
|
async def stop(self) -> None:
|
|
for task in self._tasks:
|
|
task.cancel()
|
|
await asyncio.gather(*self._tasks, return_exceptions=True)
|
|
self._tasks.clear()
|
|
|
|
async def _run(self, index: int) -> None:
|
|
while True:
|
|
download_id = await self._queue.get()
|
|
try:
|
|
await self._process(download_id)
|
|
except asyncio.CancelledError:
|
|
self._queue.task_done()
|
|
raise
|
|
except Exception: # noqa: BLE001 — воркер не должен падать целиком
|
|
logger.exception("Воркер %d упал на задаче %s", index, download_id)
|
|
finally:
|
|
self._queue.task_done()
|
|
|
|
async def _process(self, download_id: str) -> None:
|
|
download = await self._storage.get(download_id)
|
|
if download is None or download.status != DownloadStatus.PENDING:
|
|
return
|
|
|
|
updated = await self._storage.update(
|
|
download_id, status=DownloadStatus.DOWNLOADING, error=None
|
|
)
|
|
await self._publish("progress", updated)
|
|
|
|
async def on_progress(p: Progress) -> None:
|
|
fields: dict = {
|
|
"downloaded_bytes": p.downloaded_bytes,
|
|
"size_bytes": p.total_bytes,
|
|
"speed": p.speed,
|
|
"eta": p.eta,
|
|
}
|
|
if p.total_bytes:
|
|
fields["progress"] = min(100.0, p.downloaded_bytes / p.total_bytes * 100)
|
|
if p.filename:
|
|
fields["filename"] = p.filename
|
|
snapshot = await self._storage.update(download_id, **fields)
|
|
await self._publish("progress", snapshot)
|
|
|
|
try:
|
|
downloader = pick_downloader(download.url)
|
|
result = await downloader.download(download.url, on_progress)
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.warning("Ошибка скачивания %s: %s", download.url, exc)
|
|
failed = await self._storage.update(
|
|
download_id,
|
|
status=DownloadStatus.FAILED,
|
|
error=str(exc),
|
|
speed=None,
|
|
eta=None,
|
|
finished_at=_now(),
|
|
)
|
|
await self._publish("failed", failed)
|
|
return
|
|
|
|
done = await self._storage.update(
|
|
download_id,
|
|
status=DownloadStatus.DONE,
|
|
filename=result.filename,
|
|
size_bytes=result.size_bytes,
|
|
downloaded_bytes=result.size_bytes or 0,
|
|
progress=100.0,
|
|
speed=None,
|
|
eta=0,
|
|
error=None,
|
|
finished_at=_now(),
|
|
)
|
|
await self._publish("done", done)
|
|
|
|
async def _publish(self, event_type: str, download: Download | None) -> None:
|
|
if download is None:
|
|
return
|
|
await self._events.publish(DownloadEvent(type=event_type, download=download))
|