76 lines
2.5 KiB
Python
76 lines
2.5 KiB
Python
import time
|
|
import uuid
|
|
from asyncio import Lock
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from api.domain.task_status import TaskStatus
|
|
|
|
|
|
@dataclass
|
|
class TaskRecord:
|
|
task_id: str
|
|
task: str
|
|
timeout: int
|
|
metadata: dict[str, Any] | None
|
|
status: TaskStatus = TaskStatus.queued
|
|
create_at: float = field(default_factory=time.time)
|
|
started_at: float | None = None
|
|
finished_at: float | None = None
|
|
result: str | None = None
|
|
error: str | None = None
|
|
raw_response: dict[str, Any] | None = None
|
|
|
|
@property
|
|
def execution_time(self) -> float:
|
|
if self.started_at is None:
|
|
return 0
|
|
end = self.finished_at if self.finished_at is not None else time.time()
|
|
return max(0, end - self.started_at)
|
|
|
|
|
|
class TaskStore:
|
|
def __init__(self) -> None:
|
|
self._lock = Lock()
|
|
self._tasks: dict[str, TaskRecord] = {}
|
|
|
|
async def create(self, task: str, timeout: int, metadata: dict[str, Any] | None) -> TaskRecord:
|
|
task_id = uuid.uuid4().hex
|
|
rec = TaskRecord(task_id=task_id, task=task, timeout=timeout, metadata=metadata)
|
|
async with self._lock:
|
|
self._tasks[task_id] = rec
|
|
return rec
|
|
|
|
async def get(self, task_id: str) -> TaskRecord | None:
|
|
async with self._lock:
|
|
return self._tasks.get(task_id)
|
|
|
|
async def set_running(self, task_id: str) -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
rec.status = TaskStatus.running
|
|
rec.started_at = time.time()
|
|
return rec
|
|
|
|
async def set_done(
|
|
self,
|
|
task_id: str,
|
|
success: bool,
|
|
raw_response: dict[str, Any] | None,
|
|
error: str | None,
|
|
result: str | None = None,
|
|
) -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
rec.finished_at = time.time()
|
|
rec.raw_response = raw_response
|
|
rec.error = error if error is not None else (
|
|
raw_response.get("error") if isinstance(raw_response, dict) else None)
|
|
rec.result = result if result is not None else (
|
|
raw_response.get("result") if isinstance(raw_response, dict) else None)
|
|
rec.status = TaskStatus.succeeded if success else TaskStatus.failed
|
|
return rec
|