245 lines
9.6 KiB
Python
245 lines
9.6 KiB
Python
import time
|
|
import uuid
|
|
from asyncio import Event, Lock, Queue
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from api.domain.captcha_state import CaptchaState
|
|
from api.domain.task_status import TaskStatus
|
|
|
|
|
|
@dataclass
|
|
class TaskRecord:
|
|
task_id: str
|
|
thread_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
|
|
history: list[dict[str, Any]] = field(default_factory=list)
|
|
cancel_requested: bool = False
|
|
done_event: Event = field(default_factory=Event)
|
|
|
|
captcha_state: CaptchaState = CaptchaState.none
|
|
captcha_kind: str | None = None
|
|
captcha_reason: str | None = None
|
|
captcha_view_url: str | None = None
|
|
captcha_notified_at: float | None = None
|
|
captcha_solved_at: float | None = None
|
|
captcha_deadline: float | None = None
|
|
captcha_extra_seconds: int = 0
|
|
captcha_event: Event = field(default_factory=Event)
|
|
|
|
@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] = {}
|
|
self._thread_index: dict[str, list[str]] = {}
|
|
self._subscribers: dict[str, set[Queue[dict[str, Any]]]] = {}
|
|
|
|
async def create(self, task: str, timeout: int, metadata: dict[str, Any] | None, thread_id: str = "default") -> TaskRecord:
|
|
task_id = uuid.uuid4().hex
|
|
rec = TaskRecord(task_id=task_id, thread_id=thread_id, task=task, timeout=timeout, metadata=metadata)
|
|
async with self._lock:
|
|
self._tasks[task_id] = rec
|
|
self._thread_index.setdefault(thread_id, []).append(task_id)
|
|
self._subscribers.setdefault(task_id, set())
|
|
return rec
|
|
|
|
async def list_by_thread(self, thread_id: str) -> list[TaskRecord]:
|
|
async with self._lock:
|
|
ids = list(self._thread_index.get(thread_id, []))
|
|
return [self._tasks[item] for item in ids if item in self._tasks]
|
|
|
|
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
|
|
if rec.status == TaskStatus.cancelled:
|
|
return rec
|
|
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, history: list[dict[str, Any]] | 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.history = list(history or [])
|
|
rec.status = TaskStatus.succeeded if success else TaskStatus.failed
|
|
rec.done_event.set()
|
|
return rec
|
|
|
|
async def set_cancel_requested(self, task_id: str) -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
rec.cancel_requested = True
|
|
if rec.status == TaskStatus.queued:
|
|
rec.status = TaskStatus.cancelled
|
|
rec.finished_at = time.time()
|
|
rec.error = "Cancelled by user"
|
|
rec.done_event.set()
|
|
return rec
|
|
|
|
async def set_cancelled(self, task_id: str, error: str = "Cancelled by user") -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
if rec.status in (TaskStatus.succeeded, TaskStatus.failed, TaskStatus.cancelled):
|
|
return rec
|
|
rec.status = TaskStatus.cancelled
|
|
rec.finished_at = time.time()
|
|
rec.error = error
|
|
rec.done_event.set()
|
|
return rec
|
|
|
|
async def delete_if_finished(self, task_id: str) -> tuple[bool, bool]:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return False, False
|
|
if rec.status in (TaskStatus.queued, TaskStatus.running):
|
|
return True, False
|
|
del self._tasks[task_id]
|
|
thread_list = self._thread_index.get(rec.thread_id, [])
|
|
if task_id in thread_list:
|
|
thread_list.remove(task_id)
|
|
self._subscribers.pop(task_id, None)
|
|
return True, True
|
|
|
|
async def set_captcha_awaiting(self, task_id: str, kind: str | None, reason: str | None, view_url: str | None, timeout_seconds: int) -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
now = time.time()
|
|
rec.captcha_state = CaptchaState.awaiting
|
|
rec.captcha_kind = kind
|
|
rec.captcha_reason = reason
|
|
rec.captcha_view_url = view_url
|
|
rec.captcha_notified_at = now
|
|
rec.captcha_solved_at = None
|
|
rec.captcha_extra_seconds = 0
|
|
rec.captcha_deadline = now + max(1, int(timeout_seconds))
|
|
rec.captcha_event.clear()
|
|
return rec
|
|
|
|
async def set_captcha_solved(self, task_id: str) -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
if rec.captcha_state in (CaptchaState.none, CaptchaState.solved):
|
|
return rec
|
|
rec.captcha_state = CaptchaState.solved
|
|
rec.captcha_solved_at = time.time()
|
|
rec.captcha_event.set()
|
|
return rec
|
|
|
|
async def set_captcha_timeout_prompt(self, task_id: str) -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
if rec.captcha_state != CaptchaState.awaiting:
|
|
return rec
|
|
rec.captcha_state = CaptchaState.timeout_prompt
|
|
rec.captcha_event.set()
|
|
return rec
|
|
|
|
async def set_captcha_extended(self, task_id: str, extra_seconds: int) -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
if rec.captcha_state not in (CaptchaState.timeout_prompt, CaptchaState.awaiting):
|
|
return rec
|
|
extra = max(1, int(extra_seconds))
|
|
rec.captcha_extra_seconds += extra
|
|
base = rec.captcha_deadline or time.time()
|
|
rec.captcha_deadline = max(base, time.time()) + extra
|
|
rec.captcha_state = CaptchaState.extended
|
|
rec.captcha_event.set()
|
|
return rec
|
|
|
|
async def set_captcha_aborted(self, task_id: str, reason: str | None = None) -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
if rec.captcha_state in (CaptchaState.none, CaptchaState.aborted, CaptchaState.failed):
|
|
return rec
|
|
rec.captcha_state = CaptchaState.aborted
|
|
rec.captcha_reason = reason or rec.captcha_reason
|
|
rec.captcha_event.set()
|
|
return rec
|
|
|
|
async def set_captcha_failed(self, task_id: str, error: str | None = None) -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
rec.captcha_state = CaptchaState.failed
|
|
if error:
|
|
rec.captcha_reason = error
|
|
rec.captcha_event.set()
|
|
return rec
|
|
|
|
async def reset_captcha(self, task_id: str) -> TaskRecord | None:
|
|
async with self._lock:
|
|
rec = self._tasks.get(task_id)
|
|
if rec is None:
|
|
return None
|
|
rec.captcha_state = CaptchaState.none
|
|
rec.captcha_event.clear()
|
|
return rec
|
|
|
|
async def subscribe(self, task_id: str) -> Queue[dict[str, Any]] | None:
|
|
queue: Queue[dict[str, Any]] = Queue()
|
|
async with self._lock:
|
|
if task_id not in self._tasks:
|
|
return None
|
|
self._subscribers.setdefault(task_id, set()).add(queue)
|
|
return queue
|
|
|
|
async def unsubscribe(self, task_id: str, queue: Queue[dict[str, Any]]) -> None:
|
|
async with self._lock:
|
|
subscribers = self._subscribers.get(task_id)
|
|
if subscribers is not None:
|
|
subscribers.discard(queue)
|
|
|
|
async def publish(self, task_id: str, event: dict[str, Any]) -> None:
|
|
async with self._lock:
|
|
subscribers = list(self._subscribers.get(task_id, set()))
|
|
for queue in subscribers:
|
|
try:
|
|
queue.put_nowait(event)
|
|
except Exception:
|
|
continue
|