Sergey Antropoff
100 lines
3.9 KiB
Python
100 lines
3.9 KiB
Python
"""Bounded durable task worker with cooperative cancellation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from collections.abc import Awaitable, Callable
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from app.tasks.repository import Task, TaskRepository
|
|
|
|
TaskHandler = Callable[[Task], Awaitable[dict[str, Any] | None]]
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class TaskWorker:
|
|
repository: TaskRepository
|
|
worker_id: str
|
|
handlers: dict[str, TaskHandler]
|
|
concurrency: int = 2
|
|
lease_seconds: float = 30.0
|
|
poll_seconds: float = 0.1
|
|
_running: set[asyncio.Task[None]] = field(default_factory=set, init=False)
|
|
_stopping: asyncio.Event = field(default_factory=asyncio.Event, init=False)
|
|
|
|
async def run(self) -> None:
|
|
self._stopping.clear()
|
|
try:
|
|
while not self._stopping.is_set():
|
|
self._reap()
|
|
if len(self._running) >= self.concurrency:
|
|
await asyncio.sleep(self.poll_seconds)
|
|
continue
|
|
try:
|
|
task = await self.repository.claim(self.worker_id, self.lease_seconds)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception:
|
|
logger.exception("task claim failed; polling will retry")
|
|
await asyncio.sleep(self.poll_seconds)
|
|
continue
|
|
if task is None:
|
|
await asyncio.sleep(self.poll_seconds)
|
|
continue
|
|
execution = asyncio.create_task(self._execute(task))
|
|
self._running.add(execution)
|
|
finally:
|
|
if self._running:
|
|
await asyncio.gather(*self._running, return_exceptions=True)
|
|
self._running.clear()
|
|
|
|
def stop(self) -> None:
|
|
self._stopping.set()
|
|
|
|
def _reap(self) -> None:
|
|
self._running = {task for task in self._running if not task.done()}
|
|
|
|
async def _execute(self, task: Task) -> None:
|
|
handler = self.handlers.get(task.task_type)
|
|
if handler is None:
|
|
await self.repository.finish(
|
|
task.id, self.worker_id, status="error", error="unsupported task type"
|
|
)
|
|
return
|
|
try:
|
|
current = await self.repository.get(task.id)
|
|
if current is not None and current.cancel_requested:
|
|
await self.repository.finish(task.id, self.worker_id, status="cancelled")
|
|
return
|
|
execution: asyncio.Future[dict[str, Any] | None] = asyncio.ensure_future(handler(task))
|
|
heartbeat = asyncio.create_task(self._heartbeat(task))
|
|
try:
|
|
while not execution.done():
|
|
await asyncio.sleep(self.poll_seconds)
|
|
current = await self.repository.get(task.id)
|
|
if current is not None and current.cancel_requested:
|
|
execution.cancel()
|
|
await asyncio.gather(execution, return_exceptions=True)
|
|
await self.repository.finish(task.id, self.worker_id, status="cancelled")
|
|
return
|
|
result = await execution
|
|
finally:
|
|
heartbeat.cancel()
|
|
await asyncio.gather(heartbeat, return_exceptions=True)
|
|
await self.repository.finish(task.id, self.worker_id, status="success", result=result)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as error: # task failures are persisted, not leaked
|
|
await self.repository.finish(
|
|
task.id, self.worker_id, status="error", error=type(error).__name__
|
|
)
|
|
|
|
async def _heartbeat(self, task: Task) -> None:
|
|
interval = max(self.lease_seconds / 3, 0.01)
|
|
while True:
|
|
await asyncio.sleep(interval)
|
|
await self.repository.heartbeat(task.id, self.worker_id, self.lease_seconds)
|