104 lines
5.5 KiB
Python
104 lines
5.5 KiB
Python
"""Durable project calculation queue. Every worker update is generation fenced."""
|
|
from uuid import uuid4
|
|
|
|
from fastapi import HTTPException
|
|
from psycopg.types.json import Jsonb
|
|
|
|
|
|
class CalculationJobRepository:
|
|
def __init__(self, connection, project_id):
|
|
self.connection = connection
|
|
self.project_id = str(project_id)
|
|
|
|
async def enqueue(self, *, owner_user_id, request_key, tool, parameters, binding, timeout_seconds):
|
|
cursor = await self.connection.execute(
|
|
"""INSERT INTO workbench.calculation_jobs
|
|
(job_id,project_id,owner_user_id,request_key,tool,parameters,binding,timeout_seconds)
|
|
VALUES (%s,%s,%s,%s,%s,%s,%s,%s)
|
|
ON CONFLICT(project_id,owner_user_id,request_key) DO NOTHING RETURNING *""",
|
|
(uuid4(), self.project_id, owner_user_id, request_key, tool,
|
|
Jsonb(parameters), Jsonb(binding), timeout_seconds),
|
|
)
|
|
row = await cursor.fetchone()
|
|
if row:
|
|
return dict(row)
|
|
cursor = await self.connection.execute(
|
|
"SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND owner_user_id=%s AND request_key=%s",
|
|
(self.project_id, owner_user_id, request_key),
|
|
)
|
|
row = dict(await cursor.fetchone())
|
|
if row['tool'] != tool or row['parameters'] != parameters or row['binding'] != binding:
|
|
raise HTTPException(409, 'Calculation request_key was used with different inputs')
|
|
return row
|
|
|
|
async def get(self, job_id, owner_user_id=None):
|
|
cursor = await self.connection.execute(
|
|
"""SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND job_id=%s
|
|
AND (%s::uuid IS NULL OR owner_user_id=%s)""",
|
|
(self.project_id, job_id, owner_user_id, owner_user_id),
|
|
)
|
|
row = await cursor.fetchone()
|
|
return dict(row) if row else None
|
|
|
|
async def list(self, *, owner_user_id, binding, limit, offset, status=None):
|
|
cursor = await self.connection.execute(
|
|
"""SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND owner_user_id=%s
|
|
AND binding @> %s AND (%s::text IS NULL OR status=%s) ORDER BY created_at DESC,job_id LIMIT %s OFFSET %s""",
|
|
(self.project_id, owner_user_id, Jsonb(binding), status, status, limit, offset),
|
|
)
|
|
return [dict(row) for row in await cursor.fetchall()]
|
|
|
|
async def cancel(self, job_id, owner_user_id):
|
|
await self.connection.execute(
|
|
"""UPDATE workbench.calculation_jobs SET status=CASE WHEN status='queued' THEN 'cancelled'
|
|
ELSE 'cancel_requested' END,updated_at=clock_timestamp()
|
|
WHERE project_id=%s AND job_id=%s AND owner_user_id=%s AND status IN ('queued','running')""",
|
|
(self.project_id, job_id, owner_user_id),
|
|
)
|
|
return await self.get(job_id, owner_user_id)
|
|
|
|
async def recover(self):
|
|
# A crashed write is never automatically replayed.
|
|
await self.connection.execute(
|
|
"""UPDATE workbench.calculation_jobs SET status='interrupted',
|
|
error='{"code":"WORKER_INTERRUPTED"}'::jsonb,updated_at=clock_timestamp()
|
|
WHERE project_id=%s AND status IN ('running','cancel_requested')
|
|
AND lease_expires_at < clock_timestamp()""", (self.project_id,),
|
|
)
|
|
|
|
async def claim(self, owner, lease_seconds):
|
|
cursor = await self.connection.execute(
|
|
"""UPDATE workbench.calculation_jobs j SET status='running',lease_owner=%s,
|
|
lease_generation=j.lease_generation+1,
|
|
lease_expires_at=clock_timestamp()+%s*interval '1 second',updated_at=clock_timestamp()
|
|
WHERE j.job_id=(SELECT job_id FROM workbench.calculation_jobs
|
|
WHERE project_id=%s AND status='queued' ORDER BY created_at,job_id
|
|
FOR UPDATE SKIP LOCKED LIMIT 1) RETURNING j.*""",
|
|
(owner, lease_seconds, self.project_id),
|
|
)
|
|
row = await cursor.fetchone()
|
|
return dict(row) if row else None
|
|
|
|
async def heartbeat(self, row, owner, lease_seconds):
|
|
cursor = await self.connection.execute(
|
|
"""UPDATE workbench.calculation_jobs SET lease_expires_at=clock_timestamp()+%s*interval '1 second',
|
|
updated_at=clock_timestamp() WHERE project_id=%s AND job_id=%s AND lease_owner=%s
|
|
AND lease_generation=%s AND lease_expires_at>clock_timestamp()
|
|
AND status IN ('running','cancel_requested') RETURNING status""",
|
|
(lease_seconds, self.project_id, row['job_id'], owner, row['lease_generation']),
|
|
)
|
|
result = await cursor.fetchone()
|
|
return result['status'] if result else None
|
|
|
|
async def finish(self, row, owner, status, *, result=None, error=None):
|
|
cursor = await self.connection.execute(
|
|
"""UPDATE workbench.calculation_jobs SET status=%s,result=%s,error=%s,updated_at=clock_timestamp()
|
|
WHERE project_id=%s AND job_id=%s AND lease_owner=%s AND lease_generation=%s
|
|
AND lease_expires_at>clock_timestamp() AND status IN ('running','cancel_requested')
|
|
AND (status!='cancel_requested' OR %s IN ('cancelled','interrupted')
|
|
OR (%s='failed' AND %s='CLEANUP_FAILED')) RETURNING job_id""",
|
|
(status, Jsonb(result) if result is not None else None, Jsonb(error) if error else None,
|
|
self.project_id, row['job_id'], owner, row['lease_generation'], status, status, (error or {}).get('code')),
|
|
)
|
|
return await cursor.fetchone() is not None
|