107 lines
3.5 KiB
Python
107 lines
3.5 KiB
Python
# src/dataloader/workers/manager.py
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from dataloader.context import APP_CTX
|
|
from dataloader.config import APP_CONFIG
|
|
from dataloader.storage.db import get_sessionmaker
|
|
from dataloader.storage.repositories import QueueRepository
|
|
from dataloader.workers.base import PGWorker, WorkerConfig
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class WorkerSpec:
|
|
"""
|
|
Конфигурация набора воркеров для очереди.
|
|
"""
|
|
queue: str
|
|
concurrency: int
|
|
|
|
|
|
class WorkerManager:
|
|
"""
|
|
Управляет жизненным циклом асинхронных воркеров.
|
|
"""
|
|
def __init__(self, specs: list[WorkerSpec]) -> None:
|
|
self._log = APP_CTX.get_logger()
|
|
self._specs = specs
|
|
self._stop = asyncio.Event()
|
|
self._tasks: list[asyncio.Task] = []
|
|
self._reaper_task: asyncio.Task | None = None
|
|
|
|
async def start(self) -> None:
|
|
"""
|
|
Стартует воркеры и фоновую задачу реапера.
|
|
"""
|
|
hb = int(APP_CONFIG.dl.dl_heartbeat_sec)
|
|
backoff = int(APP_CONFIG.dl.dl_claim_backoff_sec)
|
|
|
|
for spec in self._specs:
|
|
for i in range(max(1, spec.concurrency)):
|
|
cfg = WorkerConfig(queue=spec.queue, heartbeat_sec=hb, claim_backoff_sec=backoff)
|
|
t = asyncio.create_task(PGWorker(cfg, self._stop).run(), name=f"worker:{spec.queue}:{i}")
|
|
self._tasks.append(t)
|
|
|
|
self._reaper_task = asyncio.create_task(self._reaper_loop(), name="reaper")
|
|
|
|
self._log.info(
|
|
"worker_manager.started",
|
|
extra={"specs": [spec.__dict__ for spec in self._specs], "total_tasks": len(self._tasks)},
|
|
)
|
|
|
|
async def stop(self) -> None:
|
|
"""
|
|
Останавливает воркеры и реапер.
|
|
"""
|
|
self._stop.set()
|
|
|
|
for t in self._tasks:
|
|
t.cancel()
|
|
await asyncio.gather(*self._tasks, return_exceptions=True)
|
|
self._tasks.clear()
|
|
|
|
if self._reaper_task:
|
|
self._reaper_task.cancel()
|
|
with contextlib.suppress(Exception):
|
|
await self._reaper_task
|
|
self._reaper_task = None
|
|
|
|
self._log.info("worker_manager.stopped")
|
|
|
|
async def _reaper_loop(self) -> None:
|
|
"""
|
|
Фоновый цикл возврата потерянных задач в очередь.
|
|
"""
|
|
period = int(APP_CONFIG.dl.dl_reaper_period_sec)
|
|
sm = get_sessionmaker()
|
|
while not self._stop.is_set():
|
|
try:
|
|
async with sm() as s:
|
|
repo = QueueRepository(s)
|
|
ids = await repo.requeue_lost()
|
|
if ids:
|
|
APP_CTX.get_logger().info("reaper.requeued", extra={"count": len(ids)})
|
|
except Exception as e:
|
|
APP_CTX.get_logger().error("reaper.error", extra={"error": str(e)})
|
|
try:
|
|
await asyncio.wait_for(self._stop.wait(), timeout=period)
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
|
|
|
|
def build_manager_from_env() -> WorkerManager:
|
|
"""
|
|
Собирает WorkerManager из WORKERS_JSON.
|
|
"""
|
|
specs: list[WorkerSpec] = []
|
|
for item in APP_CONFIG.dl.parsed_workers():
|
|
q = str(item.get("queue", "")).strip()
|
|
c = int(item.get("concurrency", 1))
|
|
if q:
|
|
specs.append(WorkerSpec(queue=q, concurrency=max(1, c)))
|
|
return WorkerManager(specs)
|