409 lines
14 KiB
Python
409 lines
14 KiB
Python
import asyncio
|
|
import json
|
|
import uuid
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Any, Dict, Optional, Sequence
|
|
|
|
from sqlalchemy import and_, or_, select, update, text
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from .. import database
|
|
from ..config import read_int_env
|
|
from ..models import SystemJobORM
|
|
|
|
|
|
JOB_STATUS_READY = "READY"
|
|
JOB_STATUS_RUNNING = "RUNNING"
|
|
JOB_STATUS_COMPLETED = "COMPLETED"
|
|
JOB_STATUS_FAILED = "FAILED"
|
|
JOB_STATUS_RETRY = "RETRY"
|
|
|
|
JOB_QUEUE_MAX_PAYLOAD_BYTES = read_int_env(
|
|
"JOB_QUEUE_MAX_PAYLOAD_BYTES",
|
|
512 * 1024,
|
|
minimum=1024,
|
|
maximum=10 * 1024 * 1024,
|
|
)
|
|
JOB_QUEUE_MAX_ATTEMPTS = read_int_env(
|
|
"JOB_QUEUE_MAX_ATTEMPTS",
|
|
10,
|
|
minimum=1,
|
|
maximum=100,
|
|
)
|
|
JOB_QUEUE_MAX_PRIORITY_ABS = read_int_env(
|
|
"JOB_QUEUE_MAX_PRIORITY_ABS",
|
|
1000,
|
|
minimum=1,
|
|
maximum=100000,
|
|
)
|
|
JOB_QUEUE_MAX_ID_LENGTH = read_int_env(
|
|
"JOB_QUEUE_MAX_ID_LENGTH",
|
|
128,
|
|
minimum=16,
|
|
maximum=512,
|
|
)
|
|
JOB_QUEUE_MAX_TYPE_LENGTH = read_int_env(
|
|
"JOB_QUEUE_MAX_TYPE_LENGTH",
|
|
64,
|
|
minimum=8,
|
|
maximum=256,
|
|
)
|
|
JOB_QUEUE_STALE_RUNNING_SECONDS = read_int_env(
|
|
"JOB_QUEUE_STALE_RUNNING_SECONDS",
|
|
300,
|
|
minimum=30,
|
|
maximum=7 * 24 * 60 * 60,
|
|
)
|
|
|
|
|
|
def _new_session() -> AsyncSession:
|
|
if database.AsyncSessionLocal is None:
|
|
database.init_db()
|
|
if database.AsyncSessionLocal is None:
|
|
raise RuntimeError("Database session factory is not initialized.")
|
|
return database.AsyncSessionLocal()
|
|
|
|
|
|
def _normalize_job_type(job_type: str) -> str:
|
|
normalized = str(job_type or "").strip().upper()
|
|
if not normalized:
|
|
raise ValueError("job_type must not be empty.")
|
|
if len(normalized) > JOB_QUEUE_MAX_TYPE_LENGTH:
|
|
raise ValueError(
|
|
f"job_type is too long (max {JOB_QUEUE_MAX_TYPE_LENGTH} chars)."
|
|
)
|
|
return normalized
|
|
|
|
|
|
def _normalize_optional_text(value: Optional[str], field_name: str) -> Optional[str]:
|
|
if value is None:
|
|
return None
|
|
normalized = str(value).strip()
|
|
if not normalized:
|
|
return None
|
|
if len(normalized) > JOB_QUEUE_MAX_ID_LENGTH:
|
|
raise ValueError(f"{field_name} is too long (max {JOB_QUEUE_MAX_ID_LENGTH} chars).")
|
|
return normalized
|
|
|
|
|
|
def _normalize_priority(value: int) -> int:
|
|
try:
|
|
parsed = int(value)
|
|
except (TypeError, ValueError):
|
|
parsed = 0
|
|
return min(JOB_QUEUE_MAX_PRIORITY_ABS, max(-JOB_QUEUE_MAX_PRIORITY_ABS, parsed))
|
|
|
|
|
|
def _normalize_max_attempts(value: int) -> int:
|
|
try:
|
|
parsed = int(value)
|
|
except (TypeError, ValueError):
|
|
parsed = 3
|
|
return min(JOB_QUEUE_MAX_ATTEMPTS, max(1, parsed))
|
|
|
|
|
|
def _normalize_payload(payload: Optional[Dict[str, Any]]) -> Dict[str, Any]:
|
|
if payload is None:
|
|
return {}
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("Job payload must be a JSON object (dict).")
|
|
|
|
try:
|
|
serialized = json.dumps(
|
|
payload,
|
|
ensure_ascii=False,
|
|
separators=(",", ":"),
|
|
allow_nan=False,
|
|
)
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError("Job payload must be JSON-serializable.") from exc
|
|
|
|
payload_bytes = len(serialized.encode("utf-8"))
|
|
if payload_bytes > JOB_QUEUE_MAX_PAYLOAD_BYTES:
|
|
raise ValueError(
|
|
f"Job payload too large ({payload_bytes} bytes), "
|
|
f"max allowed is {JOB_QUEUE_MAX_PAYLOAD_BYTES} bytes."
|
|
)
|
|
|
|
return payload
|
|
|
|
|
|
class JobQueueService:
|
|
"""
|
|
DB-backed job queue service (PostgreSQL).
|
|
"""
|
|
|
|
async def create_job(
|
|
self,
|
|
job_type: str,
|
|
payload: Optional[Dict[str, Any]] = None,
|
|
priority: int = 0,
|
|
max_attempts: int = 3,
|
|
workflow_run_id: Optional[str] = None,
|
|
workflow_step_id: Optional[str] = None,
|
|
task_id: Optional[str] = None,
|
|
db: Optional[AsyncSession] = None,
|
|
) -> str:
|
|
gen_db = db is None
|
|
if gen_db:
|
|
db = _new_session()
|
|
|
|
try:
|
|
normalized_job_type = _normalize_job_type(job_type)
|
|
normalized_payload = _normalize_payload(payload)
|
|
normalized_priority = _normalize_priority(priority)
|
|
normalized_max_attempts = _normalize_max_attempts(max_attempts)
|
|
normalized_workflow_run_id = _normalize_optional_text(workflow_run_id, "workflow_run_id")
|
|
normalized_workflow_step_id = _normalize_optional_text(workflow_step_id, "workflow_step_id")
|
|
normalized_task_id = _normalize_optional_text(task_id, "task_id")
|
|
|
|
job_id = str(uuid.uuid4())
|
|
job = SystemJobORM(
|
|
job_id=job_id,
|
|
job_type=normalized_job_type,
|
|
status=JOB_STATUS_READY,
|
|
priority=normalized_priority,
|
|
payload=normalized_payload,
|
|
max_attempts=normalized_max_attempts,
|
|
workflow_run_id=normalized_workflow_run_id,
|
|
workflow_step_id=normalized_workflow_step_id,
|
|
task_id=normalized_task_id,
|
|
)
|
|
db.add(job)
|
|
if gen_db:
|
|
await db.commit()
|
|
else:
|
|
await db.flush()
|
|
return job_id
|
|
finally:
|
|
if gen_db:
|
|
await db.close()
|
|
|
|
async def claim_next_job(
|
|
self,
|
|
worker_id: str,
|
|
lock_timeout_seconds: int = 1800,
|
|
allowed_job_types: Optional[Sequence[str]] = None,
|
|
db: Optional[AsyncSession] = None,
|
|
) -> Optional[SystemJobORM]:
|
|
gen_db = db is None
|
|
if gen_db:
|
|
db = _new_session()
|
|
|
|
try:
|
|
lock_timeout_seconds = int(lock_timeout_seconds)
|
|
normalized_allowed = [
|
|
_normalize_job_type(job_type)
|
|
for job_type in (allowed_job_types or [])
|
|
if str(job_type or "").strip()
|
|
]
|
|
normalized_allowed = list(dict.fromkeys(normalized_allowed))
|
|
allowed_clause = ""
|
|
params: Dict[str, Any] = {"lock_timeout_seconds": lock_timeout_seconds}
|
|
if normalized_allowed:
|
|
allowed_params = []
|
|
for index, job_type in enumerate(normalized_allowed):
|
|
key = f"allowed_job_type_{index}"
|
|
allowed_params.append(f":{key}")
|
|
params[key] = job_type
|
|
allowed_clause = f"AND job_type IN ({', '.join(allowed_params)})"
|
|
async with db.begin():
|
|
result = await db.execute(
|
|
text(
|
|
f"""
|
|
SELECT id
|
|
FROM system_jobs
|
|
WHERE status IN ('READY', 'RETRY')
|
|
AND (next_run_at IS NULL OR next_run_at <= NOW())
|
|
{allowed_clause}
|
|
AND (
|
|
locked_by IS NULL OR
|
|
locked_at IS NULL OR
|
|
locked_at < NOW() - make_interval(secs => :lock_timeout_seconds)
|
|
)
|
|
ORDER BY priority DESC, id ASC
|
|
FOR UPDATE SKIP LOCKED
|
|
LIMIT 1
|
|
"""
|
|
),
|
|
params,
|
|
)
|
|
row = result.first()
|
|
if not row:
|
|
return None
|
|
|
|
job_pk = row[0]
|
|
await db.execute(
|
|
update(SystemJobORM)
|
|
.where(SystemJobORM.id == job_pk)
|
|
.values(
|
|
status=JOB_STATUS_RUNNING,
|
|
locked_by=worker_id,
|
|
locked_at=datetime.now(timezone.utc).replace(tzinfo=None),
|
|
heartbeat_at=datetime.now(timezone.utc).replace(tzinfo=None),
|
|
started_at=datetime.now(timezone.utc).replace(tzinfo=None),
|
|
)
|
|
)
|
|
|
|
job_result = await db.execute(
|
|
select(SystemJobORM).where(SystemJobORM.id == job_pk)
|
|
)
|
|
return job_result.scalar_one_or_none()
|
|
finally:
|
|
if gen_db:
|
|
await db.close()
|
|
|
|
async def heartbeat(self, job_id: str, db: Optional[AsyncSession] = None) -> None:
|
|
gen_db = db is None
|
|
if gen_db:
|
|
db = _new_session()
|
|
try:
|
|
await db.execute(
|
|
update(SystemJobORM)
|
|
.where(SystemJobORM.job_id == job_id)
|
|
.values(heartbeat_at=datetime.now(timezone.utc).replace(tzinfo=None))
|
|
)
|
|
await db.commit()
|
|
finally:
|
|
if gen_db:
|
|
await db.close()
|
|
|
|
async def recover_stale_running_jobs(
|
|
self,
|
|
stale_seconds: int = JOB_QUEUE_STALE_RUNNING_SECONDS,
|
|
db: Optional[AsyncSession] = None,
|
|
) -> Dict[str, int]:
|
|
gen_db = db is None
|
|
if gen_db:
|
|
db = _new_session()
|
|
try:
|
|
safe_stale_seconds = max(30, int(stale_seconds))
|
|
now = datetime.now(timezone.utc).replace(tzinfo=None)
|
|
cutoff = now - timedelta(seconds=safe_stale_seconds)
|
|
|
|
result = await db.execute(
|
|
select(SystemJobORM).where(
|
|
SystemJobORM.status == JOB_STATUS_RUNNING,
|
|
or_(
|
|
SystemJobORM.heartbeat_at < cutoff,
|
|
and_(
|
|
SystemJobORM.heartbeat_at.is_(None),
|
|
SystemJobORM.locked_at < cutoff,
|
|
),
|
|
and_(
|
|
SystemJobORM.heartbeat_at.is_(None),
|
|
SystemJobORM.locked_at.is_(None),
|
|
),
|
|
),
|
|
)
|
|
)
|
|
stale_jobs = result.scalars().all()
|
|
if not stale_jobs:
|
|
return {"recovered": 0, "failed": 0, "failed_task_ids": []}
|
|
|
|
recovered = 0
|
|
failed = 0
|
|
failed_task_ids = []
|
|
for job in stale_jobs:
|
|
attempts = int(job.attempts or 0) + 1
|
|
if attempts < int(job.max_attempts or 1):
|
|
status = JOB_STATUS_RETRY
|
|
next_run_at = now + timedelta(seconds=5)
|
|
recovered += 1
|
|
finished_at = None
|
|
else:
|
|
status = JOB_STATUS_FAILED
|
|
next_run_at = None
|
|
failed += 1
|
|
finished_at = now
|
|
if job.task_id:
|
|
failed_task_ids.append(str(job.task_id))
|
|
|
|
await db.execute(
|
|
update(SystemJobORM)
|
|
.where(SystemJobORM.id == job.id)
|
|
.values(
|
|
status=status,
|
|
attempts=attempts,
|
|
next_run_at=next_run_at,
|
|
last_error=f"Recovered stale RUNNING job after {safe_stale_seconds}s timeout.",
|
|
finished_at=finished_at,
|
|
locked_by=None,
|
|
locked_at=None,
|
|
heartbeat_at=None,
|
|
)
|
|
)
|
|
await db.commit()
|
|
return {"recovered": recovered, "failed": failed, "failed_task_ids": failed_task_ids}
|
|
finally:
|
|
if gen_db:
|
|
await db.close()
|
|
|
|
async def mark_completed(
|
|
self,
|
|
job_id: str,
|
|
db: Optional[AsyncSession] = None,
|
|
) -> None:
|
|
gen_db = db is None
|
|
if gen_db:
|
|
db = _new_session()
|
|
try:
|
|
await db.execute(
|
|
update(SystemJobORM)
|
|
.where(SystemJobORM.job_id == job_id)
|
|
.values(
|
|
status=JOB_STATUS_COMPLETED,
|
|
finished_at=datetime.now(timezone.utc).replace(tzinfo=None),
|
|
last_error=None,
|
|
locked_by=None,
|
|
locked_at=None,
|
|
heartbeat_at=None,
|
|
)
|
|
)
|
|
await db.commit()
|
|
finally:
|
|
if gen_db:
|
|
await db.close()
|
|
|
|
async def mark_failed(
|
|
self,
|
|
job: SystemJobORM,
|
|
error_message: str,
|
|
retry_delay_seconds: int = 30,
|
|
db: Optional[AsyncSession] = None,
|
|
) -> str:
|
|
gen_db = db is None
|
|
if gen_db:
|
|
db = _new_session()
|
|
try:
|
|
attempts = (job.attempts or 0) + 1
|
|
status = JOB_STATUS_FAILED
|
|
next_run_at = None
|
|
if attempts < (job.max_attempts or 0):
|
|
status = JOB_STATUS_RETRY
|
|
backoff = retry_delay_seconds * (2 ** max(attempts - 1, 0))
|
|
next_run_at = datetime.now(timezone.utc).replace(tzinfo=None) + timedelta(seconds=backoff)
|
|
|
|
await db.execute(
|
|
update(SystemJobORM)
|
|
.where(SystemJobORM.job_id == job.job_id)
|
|
.values(
|
|
status=status,
|
|
attempts=attempts,
|
|
next_run_at=next_run_at,
|
|
last_error=error_message,
|
|
finished_at=datetime.now(timezone.utc).replace(tzinfo=None) if status == JOB_STATUS_FAILED else None,
|
|
locked_by=None,
|
|
locked_at=None,
|
|
heartbeat_at=None,
|
|
)
|
|
)
|
|
await db.commit()
|
|
return status
|
|
finally:
|
|
if gen_db:
|
|
await db.close()
|
|
|
|
|
|
job_queue_service = JobQueueService()
|