feat: add durable leased task engine
This commit is contained in:
Sergey Antropoff
@@ -0,0 +1,185 @@
|
||||
"""PostgreSQL repository for durable leased tasks."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import asyncpg # type: ignore[import-untyped] # noqa: F401
|
||||
from asyncpg import Pool, Record
|
||||
|
||||
from app.db.primitives import ConflictError, require_affected, transaction
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Task:
|
||||
id: uuid.UUID
|
||||
upid: str
|
||||
task_type: str
|
||||
status: str
|
||||
payload: dict[str, Any]
|
||||
progress: int
|
||||
cancel_requested: bool
|
||||
attempt: int
|
||||
|
||||
|
||||
def _task(row: Record) -> Task:
|
||||
return Task(
|
||||
id=row["id"],
|
||||
upid=str(row["upid"]),
|
||||
task_type=str(row["task_type"]),
|
||||
status=str(row["status"]),
|
||||
payload=json.loads(row["payload"])
|
||||
if isinstance(row["payload"], str)
|
||||
else dict(row["payload"]),
|
||||
progress=int(row["progress"]),
|
||||
cancel_requested=bool(row["cancel_requested"]),
|
||||
attempt=int(row["attempt"]),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TaskRepository:
|
||||
pool: Pool
|
||||
|
||||
async def create(
|
||||
self,
|
||||
*,
|
||||
upid: str,
|
||||
task_type: str,
|
||||
payload: dict[str, Any],
|
||||
resource_key: str | None = None,
|
||||
idempotency_key: str | None = None,
|
||||
) -> Task:
|
||||
task_id = uuid.uuid4()
|
||||
async with transaction(self.pool) as connection:
|
||||
if idempotency_key is not None:
|
||||
existing = await connection.fetchrow(
|
||||
"SELECT * FROM tasks WHERE idempotency_key=$1", idempotency_key
|
||||
)
|
||||
if existing is not None:
|
||||
return _task(existing)
|
||||
row = await connection.fetchrow(
|
||||
"""INSERT INTO tasks(id, upid, task_type, status, payload, idempotency_key)
|
||||
VALUES($1,$2,$3,'queued',$4::jsonb,$5) RETURNING *""",
|
||||
task_id,
|
||||
upid,
|
||||
task_type,
|
||||
json.dumps(payload),
|
||||
idempotency_key,
|
||||
)
|
||||
if resource_key is not None:
|
||||
try:
|
||||
await connection.execute(
|
||||
"INSERT INTO resource_locks(resource_key, task_id) VALUES($1,$2)",
|
||||
resource_key,
|
||||
task_id,
|
||||
)
|
||||
except Exception as error:
|
||||
raise ConflictError(f"resource is locked: {resource_key}") from error
|
||||
await connection.execute(
|
||||
"INSERT INTO task_events(task_id, kind) VALUES($1,'created')", task_id
|
||||
)
|
||||
if row is None:
|
||||
raise RuntimeError("task insert returned no row")
|
||||
return _task(row)
|
||||
|
||||
async def claim(self, worker_id: str, lease_seconds: float) -> Task | None:
|
||||
async with transaction(self.pool) as connection:
|
||||
row = await connection.fetchrow(
|
||||
"""WITH candidate AS (
|
||||
SELECT id FROM tasks
|
||||
WHERE status='queued' OR (status='running' AND lease_expires_at < now())
|
||||
ORDER BY created_at FOR UPDATE SKIP LOCKED LIMIT 1
|
||||
) UPDATE tasks SET status='running', worker_id=$1,
|
||||
lease_expires_at=now() + $2 * interval '1 second', attempt=attempt+1,
|
||||
updated_at=now()
|
||||
WHERE id=(SELECT id FROM candidate) RETURNING *""",
|
||||
worker_id,
|
||||
lease_seconds,
|
||||
)
|
||||
if row is None:
|
||||
return None
|
||||
await connection.execute(
|
||||
"INSERT INTO task_events(task_id, kind, data) VALUES($1,'claimed',$2::jsonb)",
|
||||
row["id"],
|
||||
json.dumps({"worker": worker_id}),
|
||||
)
|
||||
return _task(row)
|
||||
|
||||
async def heartbeat(self, task_id: uuid.UUID, worker_id: str, lease_seconds: float) -> None:
|
||||
status = await self.pool.execute(
|
||||
"""UPDATE tasks SET lease_expires_at=now()+$3*interval '1 second', updated_at=now()
|
||||
WHERE id=$1 AND worker_id=$2 AND status='running'""",
|
||||
task_id,
|
||||
worker_id,
|
||||
lease_seconds,
|
||||
)
|
||||
require_affected(status)
|
||||
|
||||
async def progress(self, task_id: uuid.UUID, worker_id: str, value: int) -> None:
|
||||
status = await self.pool.execute(
|
||||
"""UPDATE tasks SET progress=$3, updated_at=now()
|
||||
WHERE id=$1 AND worker_id=$2 AND status='running'""",
|
||||
task_id,
|
||||
worker_id,
|
||||
value,
|
||||
)
|
||||
require_affected(status)
|
||||
|
||||
async def append_log(self, task_id: uuid.UUID, message: str) -> None:
|
||||
await self.pool.execute(
|
||||
"INSERT INTO task_logs(task_id, message) VALUES($1,$2)", task_id, message
|
||||
)
|
||||
|
||||
async def request_cancel(self, task_id: uuid.UUID) -> None:
|
||||
status = await self.pool.execute(
|
||||
"""UPDATE tasks SET cancel_requested=true, updated_at=now()
|
||||
WHERE id=$1 AND status IN ('queued','running')""",
|
||||
task_id,
|
||||
)
|
||||
require_affected(status)
|
||||
|
||||
async def finish(
|
||||
self,
|
||||
task_id: uuid.UUID,
|
||||
worker_id: str,
|
||||
*,
|
||||
status: str,
|
||||
result: dict[str, Any] | None = None,
|
||||
error: str | None = None,
|
||||
) -> None:
|
||||
if status not in {"success", "error", "cancelled"}:
|
||||
raise ValueError("invalid terminal task status")
|
||||
async with transaction(self.pool) as connection:
|
||||
command = await connection.execute(
|
||||
"""UPDATE tasks SET status=$3, result=$4::jsonb, error=$5,
|
||||
progress=CASE WHEN $3='success' THEN 100 ELSE progress END,
|
||||
lease_expires_at=NULL, updated_at=now()
|
||||
WHERE id=$1 AND worker_id=$2 AND status='running'""",
|
||||
task_id,
|
||||
worker_id,
|
||||
status,
|
||||
json.dumps(result) if result is not None else None,
|
||||
error,
|
||||
)
|
||||
require_affected(command)
|
||||
await connection.execute("DELETE FROM resource_locks WHERE task_id=$1", task_id)
|
||||
await connection.execute(
|
||||
"INSERT INTO task_events(task_id, kind, data) VALUES($1,$2,$3::jsonb)",
|
||||
task_id,
|
||||
status,
|
||||
json.dumps({"error": error} if error else {}),
|
||||
)
|
||||
|
||||
async def get(self, task_id: uuid.UUID) -> Task | None:
|
||||
row = await self.pool.fetchrow("SELECT * FROM tasks WHERE id=$1", task_id)
|
||||
return _task(row) if row is not None else None
|
||||
|
||||
async def logs(self, task_id: uuid.UUID) -> tuple[str, ...]:
|
||||
rows = await self.pool.fetch(
|
||||
"SELECT message FROM task_logs WHERE task_id=$1 ORDER BY sequence", task_id
|
||||
)
|
||||
return tuple(str(row["message"]) for row in rows)
|
||||
Reference in New Issue
Block a user