diff --git a/services/ai/db/__init__.py b/services/ai/db/__init__.py index bec5c797..4d2d3056 100644 --- a/services/ai/db/__init__.py +++ b/services/ai/db/__init__.py @@ -10,11 +10,20 @@ from .model_providers import ModelProviderRecord, ModelProvidersRepository, ModelsRepository from .models import Chat, ChatMessage, ChatSearchHit, ModelRecord, Source, User from .skills import Skill, SkillsRepository +from .task_queue import ( + ClaimOptions, + EnqueueTaskRequest, + Task, + TaskClaim, + TaskQueueRepository, + TaskStats, + TaskStatus, +) from .tool_approvals import ( ToolApproval, + ToolApprovalsRepository, ToolApprovalStatus, ToolApprovalType, - ToolApprovalsRepository, ) from .usage import UsageRepository, UsageSummary from .users import UsersRepository @@ -62,4 +71,11 @@ "ToolApprovalsRepository", "SkillsRepository", "Skill", + "ClaimOptions", + "EnqueueTaskRequest", + "Task", + "TaskClaim", + "TaskQueueRepository", + "TaskStats", + "TaskStatus", ] diff --git a/services/ai/db/task_queue.py b/services/ai/db/task_queue.py new file mode 100644 index 00000000..efe7cc6d --- /dev/null +++ b/services/ai/db/task_queue.py @@ -0,0 +1,346 @@ +"""Repository for the generic PostgreSQL task queue. + +This is the Python counterpart of the Rust `shared::task_queue` module. Both +languages are thin adapters over the canonical `task_*` PostgreSQL functions +created by migration 112, so the lifecycle state machine lives in exactly one +place (the database) and the two facades cannot drift apart. + +Contract notes: + +- Delivery is at-least-once: a task can be claimed again after its lease + expires, so consumers must be idempotent around database writes and + external effects. +- A consumer that loses its lease (``heartbeat`` returns False, or a terminal + write affects no rows) must stop processing the task. +- Claim result order is not guaranteed by ``UPDATE ... RETURNING``; consumers + that need order must sort the returned tasks. +- Enqueue idempotency is per task id: producers retrying must reuse the same + ``EnqueueTaskRequest.id``. There is no generic deduplication key; logical work + coalescing is workload policy (enforced with task-specific indexes in the + workload migrations). +- A non-null ``concurrency_key`` only serializes execution: multiple tasks may + queue for the same key, but they run one at a time in oldest-task order. +""" + +import json +import logging +from dataclasses import dataclass +from datetime import UTC, datetime +from enum import StrEnum + +from asyncpg import Connection, Pool +from ulid import ULID + +from .connection import get_db_pool + +logger = logging.getLogger(__name__) + +JsonValue = str | int | float | bool | None | list["JsonValue"] | dict[str, "JsonValue"] + + +class TaskStatus(StrEnum): + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + DEAD_LETTER = "dead_letter" + + +@dataclass +class Task: + """A row from the tasks table.""" + + id: str + task_type: str + payload: dict[str, JsonValue] + payload_version: int + status: TaskStatus + priority: int + available_at: datetime + weight: int + concurrency_key: str | None + attempt_count: int + max_attempts: int + last_error: str | None + claim_token: str | None + claimed_by: str | None + lease_expires_at: datetime | None + created_at: datetime + updated_at: datetime + last_started_at: datetime | None + completed_at: datetime | None + + @classmethod + def from_row(cls, row) -> "Task": + data = dict(row) + data["status"] = TaskStatus(data["status"]) + if isinstance(data["payload"], str): + data["payload"] = json.loads(data["payload"]) + return cls(**data) + + +@dataclass +class EnqueueTaskRequest: + """A task to enqueue. ``id`` defaults to a fresh ULID; producers that + need idempotent retries must set it explicitly and reuse it.""" + + task_type: str + payload: dict[str, JsonValue] + id: str | None = None + payload_version: int = 1 + priority: int = 0 + available_at: datetime | None = None + weight: int = 1 + concurrency_key: str | None = None + max_attempts: int = 3 + + +@dataclass +class ClaimOptions: + """Claim selection and policy options.""" + + candidate_ids: list[str] | None = None + limit: int = 1 + max_weight: int | None = None + max_concurrency: int | None = None + lease_seconds: int = 300 + + +@dataclass +class TaskClaim: + """A successful claim: a fresh fencing token plus the leased tasks. + Terminal writes (complete/fail) must be fenced with ``claim_token``.""" + + claim_token: str + tasks: list[Task] + + +@dataclass +class TaskStats: + """Status statistics grouped by (task_type, status).""" + + task_type: str + status: TaskStatus + count: int + + +def _to_epoch_ms(dt: datetime | None) -> int: + if dt is None: + return int(datetime.now(UTC).timestamp() * 1000) + if dt.tzinfo is None: + dt = dt.replace(tzinfo=UTC) + return int(dt.timestamp() * 1000) + + +class TaskQueueRepository: + """Thin facade over the canonical task queue PostgreSQL functions.""" + + def __init__(self, pool: Pool | None = None): + self.pool = pool + + async def _get_pool(self) -> Pool: + if self.pool: + return self.pool + return await get_db_pool() + + async def get(self, task_id: str) -> Task | None: + """Fetch one task by id.""" + pool = await self._get_pool() + row = await pool.fetchrow("SELECT * FROM tasks WHERE id = $1", task_id) + return Task.from_row(row) if row else None + + async def enqueue(self, task: EnqueueTaskRequest) -> Task: + """Enqueue one task. Re-enqueueing an existing id is idempotent and + returns the already-stored task.""" + created = await self.enqueue_bulk([task]) + if created: + return created[0] + if task.id is None: + raise RuntimeError("task was not enqueued") + existing = await self.get(task.id) + if existing is None: + raise RuntimeError(f"task {task.id} was not enqueued") + return existing + + async def enqueue_bulk(self, tasks: list[EnqueueTaskRequest]) -> list[Task]: + """Enqueue tasks, returning only the rows actually inserted.""" + if not tasks: + return [] + + for task in tasks: + self._validate_enqueue_task_request(task) + + ids = [task.id or str(ULID()) for task in tasks] + task_types = [task.task_type for task in tasks] + payloads = [json.dumps(task.payload) for task in tasks] + payload_versions = [task.payload_version for task in tasks] + priorities = [task.priority for task in tasks] + available_at_ms = [_to_epoch_ms(task.available_at) for task in tasks] + weights = [task.weight for task in tasks] + concurrency_keys = [task.concurrency_key for task in tasks] + max_attempts = [task.max_attempts for task in tasks] + + pool = await self._get_pool() + rows = await pool.fetch( + "SELECT * FROM task_enqueue_bulk($1, $2, $3, $4, $5, $6, $7, $8, $9)", + ids, + task_types, + payloads, + payload_versions, + priorities, + available_at_ms, + weights, + concurrency_keys, + max_attempts, + ) + return [Task.from_row(row) for row in rows] + + async def claim( + self, + task_type: str, + claimed_by: str, + options: ClaimOptions, + *, + connection: Connection | None = None, + ) -> TaskClaim: + """Atomically claim a batch of tasks. + + Pass ``connection`` to claim inside a transaction that already applied + workload-specific candidate selection; the transaction makes the + selection and the claim atomic. + """ + if not task_type.strip(): + raise ValueError("task_type must not be empty") + if not claimed_by.strip(): + raise ValueError("claimed_by must not be empty") + if options.limit < 1: + raise ValueError("claim limit must be >= 1") + if options.max_weight is not None and options.max_weight < 0: + raise ValueError("claim max_weight must be >= 0") + if options.max_concurrency is not None and options.max_concurrency < 1: + raise ValueError("claim max_concurrency must be >= 1") + if options.lease_seconds < 1: + raise ValueError("claim lease_seconds must be >= 1") + if options.candidate_ids: + for candidate_id in options.candidate_ids: + if len(candidate_id) != 26: + raise ValueError( + f"candidate task id must be a 26-char ULID, got {candidate_id!r}" + ) + + pool = connection or await self._get_pool() + claim_token = str(ULID()) + rows = await pool.fetch( + "SELECT * FROM task_claim_bulk($1, $2, $3, $4, $5, $6, $7, $8)", + task_type, + options.candidate_ids, + options.limit, + options.max_weight, + options.max_concurrency, + options.lease_seconds, + claim_token, + claimed_by, + ) + return TaskClaim(claim_token=claim_token, tasks=[Task.from_row(r) for r in rows]) + + async def heartbeat(self, task_id: str, claim_token: str, lease_seconds: int) -> bool: + """Renew a running task's lease. Returns False when the task is no + longer running under this token or its lease has expired; the worker + must then stop processing it.""" + if len(task_id) != 26: + raise ValueError(f"task id must be a 26-char ULID, got {task_id!r}") + if len(claim_token) != 26: + raise ValueError(f"claim_token must be a 26-char ULID, got {claim_token!r}") + if lease_seconds < 1: + raise ValueError("heartbeat lease_seconds must be >= 1") + pool = await self._get_pool() + renewed = await pool.fetchval( + "SELECT task_heartbeat($1, $2, $3)", task_id, claim_token, lease_seconds + ) + return bool(renewed) + + async def complete(self, task_ids: list[str], claim_token: str) -> int: + """Mark claimed tasks completed. Returns the number completed.""" + if len(claim_token) != 26: + raise ValueError(f"claim_token must be a 26-char ULID, got {claim_token!r}") + if not task_ids: + return 0 + pool = await self._get_pool() + completed = await pool.fetchval( + "SELECT task_complete_bulk($1, $2)", task_ids, claim_token + ) + return int(completed) + + async def fail( + self, + task_ids: list[str], + claim_token: str, + error: str, + retryable: bool, + retry_delay_seconds: int = 0, + ) -> dict[str, TaskStatus]: + """Fail claimed tasks. + + Retryable tasks with attempt budget left return to pending and become + claimable again after ``retry_delay_seconds``; the rest become + dead_letter. Returns the resulting status per task id. + """ + if len(claim_token) != 26: + raise ValueError(f"claim_token must be a 26-char ULID, got {claim_token!r}") + if retry_delay_seconds < 0: + raise ValueError("fail retry_delay_seconds must be >= 0") + if not task_ids: + return {} + pool = await self._get_pool() + rows = await pool.fetch( + "SELECT * FROM task_fail_bulk($1, $2, $3, $4, $5)", + task_ids, + claim_token, + error, + retryable, + retry_delay_seconds, + ) + return {row["task_id"]: TaskStatus(row["result_status"]) for row in rows} + + async def recover_expired(self) -> list[Task]: + """Recover tasks whose lease expired while running: retryable tasks + are requeued as pending, exhausted tasks become dead_letter. Returns + every recovered row so adapters that mirror queue state into domain + tables can synchronize.""" + pool = await self._get_pool() + rows = await pool.fetch("SELECT * FROM task_recover_expired()") + return [Task.from_row(row) for row in rows] + + async def stats(self, task_type: str | None = None) -> list[TaskStats]: + """Status statistics grouped by (task_type, status).""" + pool = await self._get_pool() + rows = await pool.fetch("SELECT * FROM task_stats($1)", task_type) + return [ + TaskStats( + task_type=row["task_type"], + status=TaskStatus(row["status"]), + count=int(row["count"]), + ) + for row in rows + ] + + async def cleanup(self, before: datetime) -> int: + """Delete terminal (completed/dead_letter) tasks completed before the + cutoff. Returns the number of deleted rows.""" + pool = await self._get_pool() + deleted = await pool.fetchval("SELECT task_cleanup($1)", before) + return int(deleted) + + @staticmethod + def _validate_enqueue_task_request(task: EnqueueTaskRequest) -> None: + if task.id is not None and len(task.id) != 26: + raise ValueError(f"task id must be a 26-char ULID, got {task.id!r}") + if not task.task_type.strip(): + raise ValueError("task_type must not be empty") + if not isinstance(task.payload, dict): + raise ValueError("task payload must be a JSON object") + if task.weight < 0: + raise ValueError("task weight must be >= 0") + if task.max_attempts < 1: + raise ValueError("max_attempts must be >= 1") + if task.payload_version < 1: + raise ValueError("payload_version must be >= 1") diff --git a/services/ai/tests/integration/test_task_queue.py b/services/ai/tests/integration/test_task_queue.py new file mode 100644 index 00000000..a6de45b2 --- /dev/null +++ b/services/ai/tests/integration/test_task_queue.py @@ -0,0 +1,223 @@ +"""PostgreSQL-backed integration tests for the Python task queue facade.""" + +from datetime import UTC, datetime, timedelta + +import pytest +from ulid import ULID + +from db.task_queue import ( + ClaimOptions, + EnqueueTaskRequest, + TaskQueueRepository, + TaskStatus, +) + +pytestmark = pytest.mark.integration + + +def _unique_task_type(prefix: str) -> str: + return f"{prefix}_{ULID()}" + + +@pytest.mark.asyncio +async def test_enqueue_round_trip_generated_id_and_idempotent_retry(db_pool): + repo = TaskQueueRepository(pool=db_pool) + explicit_id = str(ULID()) + explicit_type = _unique_task_type("enqueue_explicit") + generated_type = _unique_task_type("enqueue_generated") + available_at = datetime.now(UTC) + timedelta(seconds=5) + + created = await repo.enqueue_bulk( + [ + EnqueueTaskRequest( + id=explicit_id, + task_type=explicit_type, + payload={"nested": {"enabled": True}, "items": [1, "two"]}, + payload_version=2, + priority=4, + available_at=available_at, + weight=7, + concurrency_key="partition-1", + max_attempts=5, + ), + EnqueueTaskRequest(task_type=generated_type, payload={"generated": True}), + ] + ) + + assert len(created) == 2 + by_type = {task.task_type: task for task in created} + explicit = by_type[explicit_type] + generated = by_type[generated_type] + + assert explicit.id == explicit_id + assert explicit.payload == {"nested": {"enabled": True}, "items": [1, "two"]} + assert explicit.payload_version == 2 + assert explicit.priority == 4 + assert explicit.weight == 7 + assert explicit.concurrency_key == "partition-1" + assert explicit.max_attempts == 5 + assert abs((explicit.available_at - available_at).total_seconds()) < 0.01 + assert len(generated.id) == 26 + assert generated.payload == {"generated": True} + + retried = await repo.enqueue( + EnqueueTaskRequest( + id=explicit_id, task_type=explicit_type, payload={"replacement": True} + ) + ) + assert retried.id == explicit_id + assert retried.payload == explicit.payload + assert retried.payload_version == explicit.payload_version + + +@pytest.mark.asyncio +async def test_claim_uses_provided_transaction(db_pool): + repo = TaskQueueRepository(pool=db_pool) + task_type = _unique_task_type("claim_transaction") + task = await repo.enqueue( + EnqueueTaskRequest(task_type=task_type, payload={"n": 1}) + ) + + async with db_pool.acquire() as connection: + transaction = connection.transaction() + await transaction.start() + try: + claim = await repo.claim( + task_type, + "worker-transaction", + ClaimOptions(limit=1, lease_seconds=60), + connection=connection, + ) + assert len(claim.tasks) == 1 + assert claim.tasks[0].id == task.id + assert claim.tasks[0].status == TaskStatus.RUNNING + assert claim.tasks[0].attempt_count == 1 + finally: + await transaction.rollback() + + rolled_back = await repo.get(task.id) + assert rolled_back is not None + assert rolled_back.status == TaskStatus.PENDING + assert rolled_back.attempt_count == 0 + + +@pytest.mark.asyncio +async def test_claim_heartbeat_and_complete_lifecycle(db_pool): + repo = TaskQueueRepository(pool=db_pool) + task_type = _unique_task_type("complete_lifecycle") + task = await repo.enqueue( + EnqueueTaskRequest(task_type=task_type, payload={"n": 1}) + ) + + claim = await repo.claim( + task_type, + "worker-complete", + ClaimOptions(limit=1, lease_seconds=60), + ) + assert len(claim.claim_token) == 26 + assert [claimed.id for claimed in claim.tasks] == [task.id] + assert claim.tasks[0].status == TaskStatus.RUNNING + + assert await repo.heartbeat(task.id, claim.claim_token, 120) is True + assert await repo.complete([task.id], claim.claim_token) == 1 + assert await repo.complete([task.id], claim.claim_token) == 0 + + completed = await repo.get(task.id) + assert completed is not None + assert completed.status == TaskStatus.COMPLETED + assert completed.completed_at is not None + assert completed.claim_token is None + assert completed.claimed_by is None + assert completed.lease_expires_at is None + + +@pytest.mark.asyncio +async def test_retry_and_dead_letter_lifecycle(db_pool): + repo = TaskQueueRepository(pool=db_pool) + task_type = _unique_task_type("fail_lifecycle") + task = await repo.enqueue( + EnqueueTaskRequest(task_type=task_type, payload={"n": 1}, max_attempts=2) + ) + + first_claim = await repo.claim(task_type, "worker-fail", ClaimOptions()) + first_result = await repo.fail( + [task.id], first_claim.claim_token, "temporary", retryable=True + ) + assert first_result == {task.id: TaskStatus.PENDING} + + pending = await repo.get(task.id) + assert pending is not None + assert pending.status == TaskStatus.PENDING + assert pending.attempt_count == 1 + assert pending.last_error == "temporary" + + second_claim = await repo.claim(task_type, "worker-fail", ClaimOptions()) + second_result = await repo.fail( + [task.id], second_claim.claim_token, "still failing", retryable=True + ) + assert second_result == {task.id: TaskStatus.DEAD_LETTER} + + dead_letter = await repo.get(task.id) + assert dead_letter is not None + assert dead_letter.status == TaskStatus.DEAD_LETTER + assert dead_letter.attempt_count == 2 + assert dead_letter.completed_at is not None + assert dead_letter.last_error == "still failing" + + +@pytest.mark.asyncio +async def test_recover_expired_maps_real_rows(db_pool): + repo = TaskQueueRepository(pool=db_pool) + task_type = _unique_task_type("recover_expired") + task = await repo.enqueue( + EnqueueTaskRequest(task_type=task_type, payload={"n": 1}, max_attempts=2) + ) + claim = await repo.claim(task_type, "worker-recover", ClaimOptions()) + assert len(claim.tasks) == 1 + + await db_pool.execute( + "UPDATE tasks SET lease_expires_at = clock_timestamp() - INTERVAL '1 second' WHERE id = $1", + task.id, + ) + recovered = await repo.recover_expired() + own_task = next(item for item in recovered if item.id == task.id) + + assert own_task.status == TaskStatus.PENDING + assert own_task.attempt_count == 1 + assert own_task.claim_token is None + assert own_task.claimed_by is None + assert own_task.lease_expires_at is None + assert own_task.last_error == "task lease expired; requeued" + + +@pytest.mark.asyncio +async def test_stats_and_cleanup(db_pool): + repo = TaskQueueRepository(pool=db_pool) + task_type = _unique_task_type("stats_cleanup") + first = await repo.enqueue( + EnqueueTaskRequest(task_type=task_type, payload={"n": 1}) + ) + second = await repo.enqueue( + EnqueueTaskRequest(task_type=task_type, payload={"n": 2}) + ) + + claim = await repo.claim(task_type, "worker-stats", ClaimOptions(limit=1)) + completed_id = claim.tasks[0].id + pending_id = second.id if completed_id == first.id else first.id + assert await repo.complete([completed_id], claim.claim_token) == 1 + + stats = await repo.stats(task_type) + counts = {row.status: row.count for row in stats} + assert counts == {TaskStatus.COMPLETED: 1, TaskStatus.PENDING: 1} + + await db_pool.execute( + "UPDATE tasks SET completed_at = clock_timestamp() - INTERVAL '7 days' WHERE id = $1", + completed_id, + ) + deleted = await repo.cleanup(datetime.now(UTC) - timedelta(days=1)) + assert deleted >= 1 + assert await repo.get(completed_id) is None + + pending = await repo.get(pending_id) + assert pending is not None + assert pending.status == TaskStatus.PENDING diff --git a/services/ai/tests/unit/test_task_queue.py b/services/ai/tests/unit/test_task_queue.py new file mode 100644 index 00000000..26c76a7b --- /dev/null +++ b/services/ai/tests/unit/test_task_queue.py @@ -0,0 +1,134 @@ +"""Pure unit tests for task queue validation and row conversion.""" + +from datetime import UTC, datetime + +import pytest + +from db.task_queue import ( + ClaimOptions, + EnqueueTaskRequest, + Task, + TaskQueueRepository, + TaskStatus, + _to_epoch_ms, +) + +pytestmark = pytest.mark.unit + +NOW = datetime(2024, 1, 1, tzinfo=UTC) +TASK_ID = "01J00000000000000000000000" +CLAIM_TOKEN = "01J00000000000000000000001" + + +def _task_row(**overrides) -> dict: + row = { + "id": TASK_ID, + "task_type": "test", + "payload": {"n": 1}, + "payload_version": 1, + "status": "pending", + "priority": 0, + "available_at": NOW, + "weight": 1, + "concurrency_key": None, + "attempt_count": 0, + "max_attempts": 3, + "last_error": None, + "claim_token": None, + "claimed_by": None, + "lease_expires_at": None, + "created_at": NOW, + "updated_at": NOW, + "last_started_at": None, + "completed_at": None, + } + row.update(overrides) + return row + + +def test_to_epoch_ms_handles_aware_and_naive_datetimes(): + expected = int(NOW.timestamp() * 1000) + + assert _to_epoch_ms(NOW) == expected + assert _to_epoch_ms(NOW.replace(tzinfo=None)) == expected + + +def test_from_row_parses_string_payload_and_status(): + task = Task.from_row(_task_row(payload='{"n": 1}', status="running")) + + assert task.payload == {"n": 1} + assert task.status == TaskStatus.RUNNING + + +@pytest.mark.asyncio +async def test_enqueue_validation_raises_before_database_access(): + repo = TaskQueueRepository() + + with pytest.raises(ValueError, match="26-char"): + await repo.enqueue_bulk( + [EnqueueTaskRequest(task_type="test", payload={}, id="short")] + ) + with pytest.raises(ValueError, match="task_type"): + await repo.enqueue_bulk([EnqueueTaskRequest(task_type=" ", payload={})]) + with pytest.raises(ValueError, match="payload"): + await repo.enqueue_bulk( + [EnqueueTaskRequest(task_type="test", payload="nope")] # type: ignore + ) + with pytest.raises(ValueError, match="weight"): + await repo.enqueue_bulk( + [EnqueueTaskRequest(task_type="test", payload={}, weight=-1)] + ) + with pytest.raises(ValueError, match="max_attempts"): + await repo.enqueue_bulk( + [EnqueueTaskRequest(task_type="test", payload={}, max_attempts=0)] + ) + with pytest.raises(ValueError, match="payload_version"): + await repo.enqueue_bulk( + [EnqueueTaskRequest(task_type="test", payload={}, payload_version=0)] + ) + + +@pytest.mark.asyncio +async def test_claim_validation_raises_before_database_access(): + repo = TaskQueueRepository() + + with pytest.raises(ValueError, match="task_type"): + await repo.claim("", "worker", ClaimOptions()) + with pytest.raises(ValueError, match="claimed_by"): + await repo.claim("test", "", ClaimOptions()) + with pytest.raises(ValueError, match="limit"): + await repo.claim("test", "worker", ClaimOptions(limit=0)) + with pytest.raises(ValueError, match="max_weight"): + await repo.claim("test", "worker", ClaimOptions(max_weight=-1)) + with pytest.raises(ValueError, match="max_concurrency"): + await repo.claim("test", "worker", ClaimOptions(max_concurrency=0)) + with pytest.raises(ValueError, match="lease_seconds"): + await repo.claim("test", "worker", ClaimOptions(lease_seconds=0)) + with pytest.raises(ValueError, match="26-char"): + await repo.claim("test", "worker", ClaimOptions(candidate_ids=["short"])) + + +@pytest.mark.asyncio +async def test_terminal_operation_validation_raises_before_database_access(): + repo = TaskQueueRepository() + + with pytest.raises(ValueError, match="task id"): + await repo.heartbeat("short", CLAIM_TOKEN, 60) + with pytest.raises(ValueError, match="claim_token"): + await repo.heartbeat(TASK_ID, "short", 60) + with pytest.raises(ValueError, match="lease"): + await repo.heartbeat(TASK_ID, CLAIM_TOKEN, 0) + with pytest.raises(ValueError, match="claim_token"): + await repo.complete([TASK_ID], "short") + with pytest.raises(ValueError, match="claim_token"): + await repo.fail([TASK_ID], "short", "boom", True) + with pytest.raises(ValueError, match="retry_delay"): + await repo.fail([TASK_ID], CLAIM_TOKEN, "boom", True, -1) + + +@pytest.mark.asyncio +async def test_empty_terminal_batches_short_circuit_without_database(): + repo = TaskQueueRepository() + + assert await repo.complete([], CLAIM_TOKEN) == 0 + assert await repo.fail([], CLAIM_TOKEN, "boom", True) == {} diff --git a/services/migrations/112_create_tasks_table.sql b/services/migrations/112_create_tasks_table.sql new file mode 100644 index 00000000..fc34436f --- /dev/null +++ b/services/migrations/112_create_tasks_table.sql @@ -0,0 +1,550 @@ +-- Generic PostgreSQL task queue. +-- +-- One table backs every kind of queued work (connector events, document +-- embeddings, agent runs, ...). The lifecycle state machine lives in the +-- task_* functions below; the shared Rust (`shared::task_queue`) and Python +-- (`db.task_queue`) facades only translate typed arguments to those +-- functions, so both languages share exactly one implementation. +-- +-- Workload-specific payloads live in `payload` and are interpreted by +-- (task_type, payload_version). Delivery is at-least-once: a task can be +-- claimed again after its lease expires, so consumers must be idempotent +-- around database writes and external effects. + +CREATE TABLE tasks ( + id TEXT PRIMARY KEY, + task_type TEXT NOT NULL, + payload JSONB NOT NULL, + payload_version INTEGER NOT NULL DEFAULT 1, + status TEXT NOT NULL DEFAULT 'pending', + priority INTEGER NOT NULL DEFAULT 0, + available_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + weight BIGINT NOT NULL DEFAULT 1, + concurrency_key TEXT, + attempt_count INTEGER NOT NULL DEFAULT 0, + max_attempts INTEGER NOT NULL DEFAULT 3, + last_error TEXT, + claim_token TEXT, + claimed_by TEXT, + lease_expires_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + last_started_at TIMESTAMPTZ, + completed_at TIMESTAMPTZ, + + CONSTRAINT tasks_id_check + CHECK (char_length(id) = 26), + CONSTRAINT tasks_task_type_check + CHECK (btrim(task_type) <> ''), + CONSTRAINT tasks_payload_check + CHECK (jsonb_typeof(payload) = 'object'), + CONSTRAINT tasks_payload_version_check + CHECK (payload_version >= 1), + CONSTRAINT tasks_status_check + CHECK (status IN ('pending', 'running', 'completed', 'dead_letter')), + CONSTRAINT tasks_weight_check + CHECK (weight >= 0), + CONSTRAINT tasks_attempt_count_check + CHECK (attempt_count >= 0), + CONSTRAINT tasks_max_attempts_check + CHECK (max_attempts >= 1), + CONSTRAINT tasks_attempt_budget_check + CHECK (attempt_count <= max_attempts), + CONSTRAINT tasks_claim_token_check + CHECK (claim_token IS NULL OR char_length(claim_token) = 26), + CONSTRAINT tasks_running_lease_check + CHECK ( + (status = 'running' + AND claim_token IS NOT NULL + AND claimed_by IS NOT NULL + AND lease_expires_at IS NOT NULL) + OR + (status <> 'running' + AND claim_token IS NULL + AND claimed_by IS NULL + AND lease_expires_at IS NULL) + ), + CONSTRAINT tasks_terminal_completed_check + CHECK ( + (status IN ('completed', 'dead_letter') AND completed_at IS NOT NULL) + OR + (status IN ('pending', 'running') AND completed_at IS NULL) + ) +); + +-- Fast eligibility scan for the common claim path (by task_type). +CREATE INDEX idx_tasks_pending_claim + ON tasks (task_type, priority DESC, available_at, id) + WHERE status = 'pending'; + +-- Stale-lease recovery scans running tasks whose lease has expired. +CREATE INDEX idx_tasks_running_lease + ON tasks (lease_expires_at) + WHERE status = 'running'; + +-- The oldest-unresolved lookup used to serialize tasks that share a +-- concurrency_key. +CREATE INDEX idx_tasks_unresolved_concurrency + ON tasks (task_type, concurrency_key, id) + WHERE status IN ('pending', 'running') AND concurrency_key IS NOT NULL; + +-- Hard safeguard: a concurrency_key can never have two running tasks. +-- The claim function already admits at most the oldest unresolved task per +-- key; this index turns any enforcement gap into a loud constraint failure. +CREATE UNIQUE INDEX idx_tasks_running_concurrency_key + ON tasks (task_type, concurrency_key) + WHERE status = 'running' AND concurrency_key IS NOT NULL; + +COMMENT ON TABLE tasks IS + 'Generic at-least-once task queue. Workload payloads are interpreted by task_type and payload_version.'; +COMMENT ON COLUMN tasks.id IS + 'Caller-generated ULID. Producers retrying idempotently reuse the same id; task_enqueue_bulk ignores duplicate ids.'; +COMMENT ON COLUMN tasks.task_type IS + 'Opaque consumer category. Planned values: connector_event, document_embedding, agent_run.'; +COMMENT ON COLUMN tasks.payload IS + 'Versioned workload input. Connector payloads include type and sync_run_id; embedding payloads include document_id; agent payloads include run_id and agent_id.'; +COMMENT ON COLUMN tasks.payload_version IS + 'Schema version of payload within task_type; starts at 1.'; +COMMENT ON COLUMN tasks.status IS + 'pending (eligible once available_at passes), running (leased by a worker), completed (successful terminal work), dead_letter (exhausted or non-retryable terminal work).'; +COMMENT ON COLUMN tasks.priority IS + 'Higher values claim first. Initial workloads use 0 to preserve FIFO.'; +COMMENT ON COLUMN tasks.available_at IS + 'Earliest claim time; also the scheduled retry time after a failure.'; +COMMENT ON COLUMN tasks.weight IS + 'Non-negative claim cost used by task_claim_bulk for cumulative batch limits. The first eligible task is admitted even if it exceeds p_max_weight. Planned values: connector tasks use JSON payload plus referenced content bytes; embeddings and agent runs use 1.'; +COMMENT ON COLUMN tasks.concurrency_key IS + 'Optional execution-serialization partition. For a given task_type, only the oldest unresolved task sharing a key is claimable and at most one matching task may be running; NULL disables per-key serialization. This does not deduplicate enqueue requests; reuse id for that.'; +COMMENT ON COLUMN tasks.attempt_count IS + 'Number of claims, including the initial attempt; incremented atomically when claimed.'; +COMMENT ON COLUMN tasks.max_attempts IS + 'Maximum number of claims before a task dead-letters.'; +COMMENT ON COLUMN tasks.last_error IS + 'Most recent failure or recovery reason; cleared on completion.'; +COMMENT ON COLUMN tasks.claim_token IS + 'Per-claim ULID fencing token required for lease renewal and terminal transitions.'; +COMMENT ON COLUMN tasks.claimed_by IS + 'Service/worker identifier holding the lease; diagnostics only.'; +COMMENT ON COLUMN tasks.lease_expires_at IS + 'Lease deadline while running; a worker that does not renew before this loses the task.'; +COMMENT ON COLUMN tasks.created_at IS + 'Enqueue time.'; +COMMENT ON COLUMN tasks.updated_at IS + 'Time of the last state transition.'; +COMMENT ON COLUMN tasks.last_started_at IS + 'Time of the most recent claim.'; +COMMENT ON COLUMN tasks.completed_at IS + 'Terminal time; set only for completed and dead_letter tasks.'; + +-- --------------------------------------------------------------------------- +-- Canonical queue functions +-- --------------------------------------------------------------------------- + +-- Bulk enqueue. Duplicate ids are ignored so a producer retry that reuses an +-- id is a no-op. Returns only the rows that were actually inserted. +-- available_at is expressed in unix epoch milliseconds. +CREATE OR REPLACE FUNCTION task_enqueue_bulk( + p_ids TEXT[], + p_task_types TEXT[], + p_payloads JSONB[], + p_payload_versions INTEGER[], + p_priorities INTEGER[], + p_available_at_ms BIGINT[], + p_weights BIGINT[], + p_concurrency_keys TEXT[], + p_max_attempts INTEGER[] +) RETURNS SETOF tasks AS $$ +BEGIN + RETURN QUERY + INSERT INTO tasks ( + id, task_type, payload, payload_version, + priority, available_at, weight, concurrency_key, max_attempts + ) + SELECT + ids.id, + ids.task_type, + ids.payload, + ids.payload_version, + ids.priority, + to_timestamp(ids.available_at_ms::double precision / 1000.0), + ids.weight, + ids.concurrency_key, + ids.max_attempts + FROM UNNEST( + p_ids, + p_task_types, + p_payloads, + p_payload_versions, + p_priorities, + p_available_at_ms, + p_weights, + p_concurrency_keys, + p_max_attempts + ) AS ids( + id, task_type, payload, payload_version, + priority, available_at_ms, weight, concurrency_key, max_attempts + ) + ON CONFLICT (id) DO NOTHING + RETURNING *; +END; +$$ LANGUAGE plpgsql; + +-- Atomic batch claim. Selects eligible pending tasks (by task_type, or by +-- caller-chosen ids when p_candidate_ids is set), transitions them to +-- running with a shared claim token and lease, increments attempt_count, and +-- returns the claimed rows. +-- +-- Ordering: priority DESC, then available_at, then id (ULID) so FIFO is +-- preserved within a priority class. UPDATE ... RETURNING does not preserve +-- this order, so callers that need it must sort by id. +-- +-- Weighted batches: p_max_weight caps the cumulative weight of the batch; +-- the first eligible task is always admitted even if it alone exceeds the +-- cap (mirrors the indexer byte budget). +-- +-- Concurrency: p_max_concurrency caps running tasks per task_type; the cap +-- is enforced atomically via an advisory lock so concurrent workers cannot +-- over-admit. The concurrency_key rule (oldest unresolved only) applies to +-- every claim, and the unique running index guards the invariant. +CREATE OR REPLACE FUNCTION task_claim_bulk( + p_task_type TEXT, + p_candidate_ids TEXT[], + p_limit INTEGER, + p_max_weight BIGINT, + p_max_concurrency INTEGER, + p_lease_seconds INTEGER, + p_claim_token TEXT, + p_claimed_by TEXT +) RETURNS SETOF tasks AS $$ +DECLARE + v_candidate_id TEXT; + v_now TIMESTAMPTZ; +BEGIN + IF p_task_type IS NULL OR btrim(p_task_type) = '' THEN + RAISE EXCEPTION 'task_claim_bulk: task_type must not be empty'; + END IF; + IF p_claimed_by IS NULL OR btrim(p_claimed_by) = '' THEN + RAISE EXCEPTION 'task_claim_bulk: claimed_by must not be empty'; + END IF; + IF p_claim_token IS NULL OR char_length(p_claim_token) <> 26 THEN + RAISE EXCEPTION 'task_claim_bulk: claim_token must be a 26-char ULID'; + END IF; + IF p_limit IS NULL OR p_limit < 1 THEN + RAISE EXCEPTION 'task_claim_bulk: limit must be >= 1'; + END IF; + IF p_lease_seconds IS NULL OR p_lease_seconds < 1 THEN + RAISE EXCEPTION 'task_claim_bulk: lease_seconds must be >= 1'; + END IF; + IF p_max_weight IS NOT NULL AND p_max_weight < 0 THEN + RAISE EXCEPTION 'task_claim_bulk: max_weight must be >= 0'; + END IF; + IF p_max_concurrency IS NOT NULL AND p_max_concurrency < 1 THEN + RAISE EXCEPTION 'task_claim_bulk: max_concurrency must be >= 1'; + END IF; + IF p_candidate_ids IS NOT NULL THEN + FOREACH v_candidate_id IN ARRAY p_candidate_ids LOOP + IF v_candidate_id IS NULL OR char_length(v_candidate_id) <> 26 THEN + RAISE EXCEPTION 'task_claim_bulk: candidate task id must be a 26-char ULID: %', v_candidate_id; + END IF; + END LOOP; + END IF; + + IF p_max_concurrency IS NOT NULL THEN + -- Serialize claims per task_type so the running-count check and the + -- state transition below are atomic across workers. + PERFORM pg_advisory_xact_lock(hashtext(p_task_type)::bigint); + END IF; + + -- clock_timestamp() keeps advancing while this statement runs, so a + -- lease created after a long advisory-lock wait is measured from the + -- moment the claim actually proceeds. + v_now := clock_timestamp(); + + RETURN QUERY + WITH candidates AS ( + SELECT c.id, c.priority, c.available_at, c.weight + FROM tasks c + WHERE c.status = 'pending' + AND c.available_at <= v_now + AND c.attempt_count < c.max_attempts + AND c.task_type = p_task_type + AND (p_candidate_ids IS NULL OR c.id = ANY(p_candidate_ids)) + AND ( + c.concurrency_key IS NULL + OR NOT EXISTS ( + SELECT 1 FROM tasks older + WHERE older.task_type = c.task_type + AND older.concurrency_key = c.concurrency_key + AND older.status IN ('pending', 'running') + AND older.id < c.id + ) + ) + ORDER BY c.priority DESC, c.available_at, c.id + LIMIT p_limit + FOR UPDATE SKIP LOCKED + ), + -- MATERIALIZED: without it the planner inlines `ranked` into `batch` and + -- drops the WHERE filter when the query has two window functions + -- (observed on PostgreSQL 16/17), returning every candidate regardless of + -- the weight/concurrency conditions. + ranked AS MATERIALIZED ( + SELECT id, + row_number() OVER ( + ORDER BY priority DESC, available_at, id + ) AS row_num, + SUM(weight) OVER ( + ORDER BY priority DESC, available_at, id + ROWS UNBOUNDED PRECEDING + ) AS running_weight + FROM candidates + ), + batch AS ( + SELECT id + FROM ranked + WHERE ( + p_max_weight IS NULL + OR row_num = 1 + OR running_weight <= p_max_weight + ) + AND ( + p_max_concurrency IS NULL + OR row_num <= p_max_concurrency - ( + SELECT COUNT(*)::INTEGER FROM tasks active + WHERE active.task_type = p_task_type + AND active.status = 'running' + AND active.lease_expires_at > v_now + ) + ) + ) + UPDATE tasks t + SET status = 'running', + claim_token = p_claim_token, + claimed_by = p_claimed_by, + lease_expires_at = v_now + make_interval(secs => p_lease_seconds), + last_started_at = v_now, + attempt_count = t.attempt_count + 1, + updated_at = v_now + FROM batch + WHERE t.id = batch.id + RETURNING t.*; +END; +$$ LANGUAGE plpgsql; + +-- Renew a running task's lease. Returns false when the task is not running +-- under this token or its lease has already expired (the worker has lost the +-- task and must stop processing it). +CREATE OR REPLACE FUNCTION task_heartbeat( + p_task_id TEXT, + p_claim_token TEXT, + p_lease_seconds INTEGER +) RETURNS BOOLEAN AS $$ +DECLARE + v_updated BIGINT; + v_now TIMESTAMPTZ; +BEGIN + IF p_task_id IS NULL OR char_length(p_task_id) <> 26 THEN + RAISE EXCEPTION 'task_heartbeat: task id must be a 26-char ULID'; + END IF; + IF p_claim_token IS NULL OR char_length(p_claim_token) <> 26 THEN + RAISE EXCEPTION 'task_heartbeat: claim_token must be a 26-char ULID'; + END IF; + IF p_lease_seconds IS NULL OR p_lease_seconds < 1 THEN + RAISE EXCEPTION 'task_heartbeat: lease_seconds must be >= 1'; + END IF; + + -- Lock the row before taking a timestamp so a concurrent updater cannot + -- push the lease check past its real wall-clock expiry. + PERFORM 1 FROM tasks WHERE id = p_task_id FOR UPDATE; + IF NOT FOUND THEN + RETURN false; + END IF; + + v_now := clock_timestamp(); + + UPDATE tasks + SET lease_expires_at = v_now + make_interval(secs => p_lease_seconds), + updated_at = v_now + WHERE id = p_task_id + AND claim_token = p_claim_token + AND status = 'running' + AND lease_expires_at > v_now; + GET DIAGNOSTICS v_updated = ROW_COUNT; + RETURN v_updated = 1; +END; +$$ LANGUAGE plpgsql; + +-- Mark claimed tasks completed. Fenced by the batch claim token and an +-- unexpired lease. Returns the number of tasks completed. +CREATE OR REPLACE FUNCTION task_complete_bulk( + p_task_ids TEXT[], + p_claim_token TEXT +) RETURNS BIGINT AS $$ +DECLARE + v_updated BIGINT; + v_now TIMESTAMPTZ; +BEGIN + IF p_task_ids IS NULL OR cardinality(p_task_ids) = 0 THEN + RAISE EXCEPTION 'task_complete_bulk: task_ids must not be empty'; + END IF; + IF p_claim_token IS NULL OR char_length(p_claim_token) <> 26 THEN + RAISE EXCEPTION 'task_complete_bulk: claim_token must be a 26-char ULID'; + END IF; + + -- Lock every target row in id order before taking a timestamp so a long + -- lock wait cannot let an expired lease pass the fence. + PERFORM 1 + FROM tasks + WHERE id = ANY(p_task_ids) + ORDER BY id + FOR UPDATE; + + v_now := clock_timestamp(); + + UPDATE tasks + SET status = 'completed', + completed_at = v_now, + last_error = NULL, + claim_token = NULL, + claimed_by = NULL, + lease_expires_at = NULL, + updated_at = v_now + WHERE id = ANY(p_task_ids) + AND claim_token = p_claim_token + AND status = 'running' + AND lease_expires_at > v_now; + GET DIAGNOSTICS v_updated = ROW_COUNT; + RETURN v_updated; +END; +$$ LANGUAGE plpgsql; + +-- Fail claimed tasks. Tasks that are retryable and still have attempt budget +-- return to pending with available_at = NOW() + retry_delay_seconds; the +-- rest become dead_letter. Returns the resulting status per task. +CREATE OR REPLACE FUNCTION task_fail_bulk( + p_task_ids TEXT[], + p_claim_token TEXT, + p_error TEXT, + p_retryable BOOLEAN, + p_retry_delay_seconds INTEGER +) RETURNS TABLE (task_id TEXT, result_status TEXT) AS $$ +DECLARE + v_now TIMESTAMPTZ; +BEGIN + IF p_task_ids IS NULL OR cardinality(p_task_ids) = 0 THEN + RAISE EXCEPTION 'task_fail_bulk: task_ids must not be empty'; + END IF; + IF p_claim_token IS NULL OR char_length(p_claim_token) <> 26 THEN + RAISE EXCEPTION 'task_fail_bulk: claim_token must be a 26-char ULID'; + END IF; + IF p_retryable IS NULL THEN + RAISE EXCEPTION 'task_fail_bulk: retryable must be true or false'; + END IF; + IF p_retry_delay_seconds IS NULL OR p_retry_delay_seconds < 0 THEN + RAISE EXCEPTION 'task_fail_bulk: retry_delay_seconds must be >= 0'; + END IF; + + -- Lock every target row in id order before taking a timestamp so a long + -- lock wait cannot let an expired lease pass the fence. + PERFORM 1 + FROM tasks + WHERE id = ANY(p_task_ids) + ORDER BY id + FOR UPDATE; + + v_now := clock_timestamp(); + + RETURN QUERY + UPDATE tasks t + SET status = CASE + WHEN p_retryable AND t.attempt_count < t.max_attempts THEN 'pending' + ELSE 'dead_letter' + END, + available_at = CASE + WHEN p_retryable AND t.attempt_count < t.max_attempts + THEN v_now + make_interval(secs => p_retry_delay_seconds) + ELSE t.available_at + END, + completed_at = CASE + WHEN p_retryable AND t.attempt_count < t.max_attempts THEN NULL + ELSE v_now + END, + last_error = p_error, + claim_token = NULL, + claimed_by = NULL, + lease_expires_at = NULL, + updated_at = v_now + WHERE t.id = ANY(p_task_ids) + AND t.claim_token = p_claim_token + AND t.status = 'running' + AND t.lease_expires_at > v_now + RETURNING t.id AS task_id, t.status AS result_status; +END; +$$ LANGUAGE plpgsql; + +-- Recover tasks whose lease expired while running. Tasks with remaining +-- attempt budget return to pending (immediately claimable); exhausted tasks +-- become dead_letter. Returns every recovered row so adapters that keep +-- domain state (e.g. agent_runs) can synchronize. +CREATE OR REPLACE FUNCTION task_recover_expired() RETURNS SETOF tasks AS $$ +DECLARE + v_now TIMESTAMPTZ; +BEGIN + v_now := statement_timestamp(); + + RETURN QUERY + UPDATE tasks t + SET status = CASE + WHEN t.attempt_count < t.max_attempts THEN 'pending' + ELSE 'dead_letter' + END, + available_at = CASE + WHEN t.attempt_count < t.max_attempts THEN v_now + ELSE t.available_at + END, + completed_at = CASE + WHEN t.attempt_count < t.max_attempts THEN NULL + ELSE v_now + END, + last_error = CASE + WHEN t.attempt_count < t.max_attempts THEN 'task lease expired; requeued' + ELSE 'task lease expired; retries exhausted' + END, + claim_token = NULL, + claimed_by = NULL, + lease_expires_at = NULL, + updated_at = v_now + WHERE t.status = 'running' + AND t.lease_expires_at <= v_now + RETURNING t.*; +END; +$$ LANGUAGE plpgsql; + +-- Status statistics grouped by (task_type, status). Pass NULL for all types. +CREATE OR REPLACE FUNCTION task_stats( + p_task_type TEXT +) RETURNS TABLE (task_type TEXT, status TEXT, count BIGINT) AS $$ +BEGIN + RETURN QUERY + SELECT t.task_type::TEXT, t.status::TEXT, COUNT(*)::BIGINT + FROM tasks t + WHERE p_task_type IS NULL OR t.task_type = p_task_type + GROUP BY t.task_type, t.status + ORDER BY t.task_type, t.status; +END; +$$ LANGUAGE plpgsql; + +-- Delete terminal (completed/dead_letter) tasks completed before the cutoff. +-- Returns the number of deleted rows. +CREATE OR REPLACE FUNCTION task_cleanup( + p_before TIMESTAMPTZ +) RETURNS BIGINT AS $$ +DECLARE + v_deleted BIGINT; +BEGIN + DELETE FROM tasks + WHERE status IN ('completed', 'dead_letter') + AND completed_at < p_before; + GET DIAGNOSTICS v_deleted = ROW_COUNT; + RETURN v_deleted; +END; +$$ LANGUAGE plpgsql; diff --git a/shared/src/lib.rs b/shared/src/lib.rs index 26679613..cdd18c64 100644 --- a/shared/src/lib.rs +++ b/shared/src/lib.rs @@ -12,6 +12,7 @@ pub mod queue; pub mod rate_limiter; pub mod service_auth; pub mod storage; +pub mod task_queue; pub mod telemetry; pub mod traits; pub mod utils; @@ -40,6 +41,9 @@ pub use storage::{ ContentMetadata as StorageContentMetadata, ObjectStorage, StorageError, factory::{StorageBackend, StorageFactory}, }; +pub use task_queue::{ + ClaimOptions, EnqueueTaskRequest, Task, TaskClaim, TaskQueue, TaskStats, TaskStatus, +}; pub use traits::Repository; pub fn init() { diff --git a/shared/src/task_queue.rs b/shared/src/task_queue.rs new file mode 100644 index 00000000..9e186501 --- /dev/null +++ b/shared/src/task_queue.rs @@ -0,0 +1,472 @@ +use anyhow::{Result, bail}; +use serde::{Deserialize, Serialize}; +use sqlx::{PgPool, Row}; +use time::OffsetDateTime; + +use crate::utils::generate_ulid; + +/// Lifecycle state of a task. There is no transient `failed` state: a +/// retryable failure returns the task to `pending` with a future +/// `available_at`, and exhausted or non-retryable failures become +/// `dead_letter`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TaskStatus { + Pending, + Running, + Completed, + DeadLetter, +} + +impl TaskStatus { + pub fn as_str(&self) -> &'static str { + match self { + TaskStatus::Pending => "pending", + TaskStatus::Running => "running", + TaskStatus::Completed => "completed", + TaskStatus::DeadLetter => "dead_letter", + } + } +} + +impl std::fmt::Display for TaskStatus { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +impl std::str::FromStr for TaskStatus { + type Err = anyhow::Error; + + fn from_str(s: &str) -> Result { + match s { + "pending" => Ok(TaskStatus::Pending), + "running" => Ok(TaskStatus::Running), + "completed" => Ok(TaskStatus::Completed), + "dead_letter" => Ok(TaskStatus::DeadLetter), + _ => Err(anyhow::anyhow!("invalid task status: {}", s)), + } + } +} + +/// A full row from the `tasks` table. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Task { + pub id: String, + pub task_type: String, + pub payload: serde_json::Value, + pub payload_version: i32, + pub status: TaskStatus, + pub priority: i32, + pub available_at: OffsetDateTime, + pub weight: i64, + pub concurrency_key: Option, + pub attempt_count: i32, + pub max_attempts: i32, + pub last_error: Option, + pub claim_token: Option, + pub claimed_by: Option, + pub lease_expires_at: Option, + pub created_at: OffsetDateTime, + pub updated_at: OffsetDateTime, + pub last_started_at: Option, + pub completed_at: Option, +} + +impl sqlx::FromRow<'_, sqlx::postgres::PgRow> for Task { + fn from_row(row: &sqlx::postgres::PgRow) -> Result { + use sqlx::Row; + + let status_str: String = row.try_get("status")?; + let status = status_str + .parse::() + .map_err(|e| sqlx::Error::ColumnDecode { + index: "status".to_string(), + source: e.into(), + })?; + + Ok(Task { + id: row.try_get("id")?, + task_type: row.try_get("task_type")?, + payload: row.try_get("payload")?, + payload_version: row.try_get("payload_version")?, + status, + priority: row.try_get("priority")?, + available_at: row.try_get("available_at")?, + weight: row.try_get("weight")?, + concurrency_key: row.try_get("concurrency_key")?, + attempt_count: row.try_get("attempt_count")?, + max_attempts: row.try_get("max_attempts")?, + last_error: row.try_get("last_error")?, + claim_token: row.try_get("claim_token")?, + claimed_by: row.try_get("claimed_by")?, + lease_expires_at: row.try_get("lease_expires_at")?, + created_at: row.try_get("created_at")?, + updated_at: row.try_get("updated_at")?, + last_started_at: row.try_get("last_started_at")?, + completed_at: row.try_get("completed_at")?, + }) + } +} + +/// A task to enqueue. `id` defaults to a fresh ULID; producers that need +/// idempotent retries must set it explicitly and reuse it. +#[derive(Debug, Clone)] +pub struct EnqueueTaskRequest { + pub id: String, + pub task_type: String, + pub payload: serde_json::Value, + pub payload_version: i32, + pub priority: i32, + pub available_at: OffsetDateTime, + pub weight: i64, + pub concurrency_key: Option, + pub max_attempts: i32, +} + +impl EnqueueTaskRequest { + pub fn new(task_type: impl Into, payload: serde_json::Value) -> Self { + Self { + id: generate_ulid(), + task_type: task_type.into(), + payload, + payload_version: 1, + priority: 0, + available_at: OffsetDateTime::now_utc(), + weight: 1, + concurrency_key: None, + max_attempts: 3, + } + } +} + +/// Claim selection and policy options. A claim selects eligible tasks either +/// by `task_type` or by explicit `candidate_ids` chosen by the caller (e.g. +/// in a transaction that already applied workload-specific filtering). +#[derive(Debug, Clone)] +pub struct ClaimOptions { + /// Claim only these task ids (already eligible per the caller's own + /// selection). When `None`, claim by `task_type`. + pub candidate_ids: Option>, + /// Maximum number of tasks to claim. + pub limit: i32, + /// Cap on cumulative batch weight; the first eligible task is always + /// admitted even if it alone exceeds the cap. + pub max_weight: Option, + /// Cap on concurrently running tasks of this type. Enforced atomically + /// per task type. + pub max_concurrency: Option, + /// Lease duration in seconds granted by this claim. + pub lease_seconds: i32, +} + +impl Default for ClaimOptions { + fn default() -> Self { + Self { + candidate_ids: None, + limit: 1, + max_weight: None, + max_concurrency: None, + lease_seconds: 300, + } + } +} + +/// A successful claim: a fresh fencing token plus the tasks now leased to the +/// caller. Terminal writes (complete/fail) must be fenced with this token. +#[derive(Debug, Clone)] +pub struct TaskClaim { + pub claim_token: String, + pub tasks: Vec, +} + +/// Status statistics grouped by (task_type, status). +#[derive(Debug, Clone, Serialize)] +pub struct TaskStats { + pub task_type: String, + pub status: TaskStatus, + pub count: i64, +} + +/// Thin, strongly typed facade over the canonical task queue PostgreSQL +/// functions created by migration 112. All queue semantics live in SQL so the +/// Rust and Python facades can never drift apart. +#[derive(Clone)] +pub struct TaskQueue { + pool: PgPool, +} + +impl TaskQueue { + pub fn new(pool: PgPool) -> Self { + Self { pool } + } + + /// Enqueue tasks, returning only the rows that were actually inserted. + /// Re-enqueueing an existing id is a no-op, which is what makes producer + /// retries idempotent. + pub async fn enqueue_bulk(&self, tasks: &[EnqueueTaskRequest]) -> Result> { + if tasks.is_empty() { + return Ok(Vec::new()); + } + for task in tasks { + validate_enqueue_task_request(task)?; + } + + let ids: Vec = tasks.iter().map(|t| t.id.clone()).collect(); + let task_types: Vec = tasks.iter().map(|t| t.task_type.clone()).collect(); + let payloads: Vec = tasks.iter().map(|t| t.payload.clone()).collect(); + let payload_versions: Vec = tasks.iter().map(|t| t.payload_version).collect(); + let priorities: Vec = tasks.iter().map(|t| t.priority).collect(); + let available_at_ms: Vec = tasks.iter().map(|t| to_epoch_ms(t.available_at)).collect(); + let weights: Vec = tasks.iter().map(|t| t.weight).collect(); + let concurrency_keys: Vec> = + tasks.iter().map(|t| t.concurrency_key.clone()).collect(); + let max_attempts: Vec = tasks.iter().map(|t| t.max_attempts).collect(); + + let rows = sqlx::query_as::<_, Task>( + "SELECT * FROM task_enqueue_bulk($1, $2, $3, $4, $5, $6, $7, $8, $9)", + ) + .bind(&ids) + .bind(&task_types) + .bind(&payloads) + .bind(&payload_versions) + .bind(&priorities) + .bind(&available_at_ms) + .bind(&weights) + .bind(&concurrency_keys) + .bind(&max_attempts) + .fetch_all(&self.pool) + .await?; + + Ok(rows) + } + + /// Convenience wrapper around [`TaskQueue::enqueue_bulk`]. Re-enqueueing + /// an existing id is idempotent and returns the already-stored task. + pub async fn enqueue(&self, task: EnqueueTaskRequest) -> Result { + let created = self.enqueue_bulk(std::slice::from_ref(&task)).await?; + if let Some(row) = created.into_iter().next() { + return Ok(row); + } + let existing = sqlx::query_as::<_, Task>("SELECT * FROM tasks WHERE id = $1") + .bind(&task.id) + .fetch_optional(&self.pool) + .await?; + existing.ok_or_else(|| anyhow::anyhow!("task {} was not enqueued", task.id)) + } + + /// Atomically claim a batch of tasks. Pass a pool, connection, or + /// transaction as `executor`; the transaction form lets callers select + /// candidate ids with workload-specific SQL and claim them in the same + /// transaction. Generates the batch claim token internally. + pub async fn claim_bulk<'e, E>( + &self, + executor: E, + task_type: &str, + claimed_by: &str, + options: &ClaimOptions, + ) -> Result + where + E: sqlx::Executor<'e, Database = sqlx::Postgres>, + { + if task_type.trim().is_empty() { + bail!("claim task_type must not be empty"); + } + if claimed_by.trim().is_empty() { + bail!("claim claimed_by must not be empty"); + } + validate_claim_options(options)?; + + let claim_token = generate_ulid(); + let rows = sqlx::query_as::<_, Task>( + "SELECT * FROM task_claim_bulk($1, $2, $3, $4, $5, $6, $7, $8)", + ) + .bind(task_type) + .bind(&options.candidate_ids) + .bind(options.limit) + .bind(options.max_weight) + .bind(options.max_concurrency) + .bind(options.lease_seconds) + .bind(&claim_token) + .bind(claimed_by) + .fetch_all(executor) + .await?; + + Ok(TaskClaim { + claim_token, + tasks: rows, + }) + } + + /// Renew a running task's lease. Returns `false` when the task is no + /// longer running under this token or its lease has already expired; the + /// worker must then stop processing the task. + pub async fn heartbeat( + &self, + task_id: &str, + claim_token: &str, + lease_seconds: i32, + ) -> Result { + if task_id.len() != 26 { + bail!("heartbeat task id must be a 26-char ULID"); + } + if claim_token.len() != 26 { + bail!("heartbeat claim_token must be a 26-char ULID"); + } + if lease_seconds < 1 { + bail!("heartbeat lease_seconds must be >= 1"); + } + let renewed = sqlx::query_scalar::<_, bool>("SELECT task_heartbeat($1, $2, $3)") + .bind(task_id) + .bind(claim_token) + .bind(lease_seconds) + .fetch_one(&self.pool) + .await?; + Ok(renewed) + } + + /// Mark claimed tasks completed. Fenced by the batch claim token and an + /// unexpired lease. Returns the number of tasks completed. + pub async fn complete_bulk(&self, task_ids: &[String], claim_token: &str) -> Result { + if claim_token.len() != 26 { + bail!("complete claim_token must be a 26-char ULID"); + } + if task_ids.is_empty() { + return Ok(0); + } + let completed = sqlx::query_scalar::<_, i64>("SELECT task_complete_bulk($1, $2)") + .bind(task_ids) + .bind(claim_token) + .fetch_one(&self.pool) + .await?; + Ok(completed) + } + + /// Fail claimed tasks. Retryable tasks with attempt budget left return to + /// `pending` and become claimable again after `retry_delay_seconds`; + /// everything else becomes `dead_letter`. Returns the resulting status + /// per task id. + pub async fn fail_bulk( + &self, + task_ids: &[String], + claim_token: &str, + error: &str, + retryable: bool, + retry_delay_seconds: i32, + ) -> Result> { + if claim_token.len() != 26 { + bail!("fail claim_token must be a 26-char ULID"); + } + if retry_delay_seconds < 0 { + bail!("fail retry_delay_seconds must be >= 0"); + } + if task_ids.is_empty() { + return Ok(Vec::new()); + } + let rows = sqlx::query("SELECT * FROM task_fail_bulk($1, $2, $3, $4, $5)") + .bind(task_ids) + .bind(claim_token) + .bind(error) + .bind(retryable) + .bind(retry_delay_seconds) + .fetch_all(&self.pool) + .await?; + + let mut result = Vec::with_capacity(rows.len()); + for row in rows { + let task_id: String = row.get("task_id"); + let status: TaskStatus = row.get::("result_status").parse()?; + result.push((task_id, status)); + } + Ok(result) + } + + /// Recover tasks whose lease expired while running: retryable tasks are + /// requeued as pending, exhausted tasks become dead_letter. Returns every + /// recovered row so adapters that mirror queue state into domain tables + /// can synchronize. + pub async fn recover_expired(&self) -> Result> { + let rows = sqlx::query_as::<_, Task>("SELECT * FROM task_recover_expired()") + .fetch_all(&self.pool) + .await?; + Ok(rows) + } + + /// Status statistics grouped by (task_type, status). Pass `None` for all + /// task types. + pub async fn stats(&self, task_type: Option<&str>) -> Result> { + let rows = sqlx::query("SELECT * FROM task_stats($1)") + .bind(task_type) + .fetch_all(&self.pool) + .await?; + + let mut stats = Vec::with_capacity(rows.len()); + for row in rows { + stats.push(TaskStats { + task_type: row.get("task_type"), + status: row.get::("status").parse()?, + count: row.get("count"), + }); + } + Ok(stats) + } + + /// Delete terminal (completed/dead_letter) tasks completed before the + /// cutoff. Returns the number of deleted rows. + pub async fn cleanup(&self, before: OffsetDateTime) -> Result { + let deleted = sqlx::query_scalar::<_, i64>("SELECT task_cleanup($1)") + .bind(before) + .fetch_one(&self.pool) + .await?; + Ok(deleted) + } +} + +fn to_epoch_ms(odt: OffsetDateTime) -> i64 { + odt.unix_timestamp() * 1000 + i64::from(odt.millisecond()) +} + +fn validate_enqueue_task_request(task: &EnqueueTaskRequest) -> Result<()> { + if task.id.len() != 26 { + bail!("task id must be a 26-char ULID, got {:?}", task.id); + } + if task.task_type.trim().is_empty() { + bail!("task_type must not be empty"); + } + if !task.payload.is_object() { + bail!("task payload must be a JSON object"); + } + if task.payload_version < 1 { + bail!("payload_version must be >= 1"); + } + if task.weight < 0 { + bail!("task weight must be >= 0"); + } + if task.max_attempts < 1 { + bail!("max_attempts must be >= 1"); + } + Ok(()) +} + +fn validate_claim_options(options: &ClaimOptions) -> Result<()> { + if options.limit < 1 { + bail!("claim limit must be >= 1"); + } + if options.max_weight.is_some_and(|w| w < 0) { + bail!("claim max_weight must be >= 0"); + } + if options.max_concurrency.is_some_and(|c| c < 1) { + bail!("claim max_concurrency must be >= 1"); + } + if options.lease_seconds < 1 { + bail!("claim lease_seconds must be >= 1"); + } + if let Some(ids) = &options.candidate_ids { + for id in ids { + if id.len() != 26 { + bail!("candidate task id must be a 26-char ULID, got {:?}", id); + } + } + } + Ok(()) +} diff --git a/shared/tests/task_queue_test.rs b/shared/tests/task_queue_test.rs new file mode 100644 index 00000000..1d32e50f --- /dev/null +++ b/shared/tests/task_queue_test.rs @@ -0,0 +1,1157 @@ +use futures_util::future::join_all; +use shared::task_queue::{ClaimOptions, EnqueueTaskRequest, TaskQueue, TaskStatus}; +use shared::test_environment::TestEnvironment; +use sqlx::PgPool; +use time::{Duration, OffsetDateTime}; + +async fn new_queue() -> (TestEnvironment, PgPool, TaskQueue) { + let env = TestEnvironment::new().await.unwrap(); + let pool = env.db_pool.pool().clone(); + let queue = TaskQueue::new(pool.clone()); + (env, pool, queue) +} + +fn enqueue_tasks(task_type: &str, count: usize) -> Vec { + (0..count) + .map(|i| EnqueueTaskRequest::new(task_type, serde_json::json!({ "n": i }))) + .collect() +} + +fn claim_opts(limit: i32) -> ClaimOptions { + ClaimOptions { + limit, + ..Default::default() + } +} + +#[tokio::test] +async fn test_enqueue_claim_complete_lifecycle() { + let (_env, pool, queue) = new_queue().await; + + let created = queue.enqueue_bulk(&enqueue_tasks("test", 1)).await.unwrap(); + assert_eq!(created.len(), 1); + let task = &created[0]; + assert_eq!(task.status, TaskStatus::Pending); + assert_eq!(task.attempt_count, 0); + assert!(task.claim_token.is_none()); + assert!(task.completed_at.is_none()); + + let claim = queue + .claim_bulk(&pool, "test", "worker-1", &claim_opts(1)) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1); + let claimed = &claim.tasks[0]; + assert_eq!(claimed.id, task.id); + assert_eq!(claimed.status, TaskStatus::Running); + assert_eq!(claimed.attempt_count, 1); + assert_eq!( + claimed.claim_token.as_deref(), + Some(claim.claim_token.as_str()) + ); + assert_eq!(claimed.claimed_by.as_deref(), Some("worker-1")); + assert!(claimed.lease_expires_at.is_some()); + assert!(claimed.last_started_at.is_some()); + + // A second claim must not see the running task. + let second = queue + .claim_bulk(&pool, "test", "worker-2", &claim_opts(1)) + .await + .unwrap(); + assert!(second.tasks.is_empty()); + + let completed = queue + .complete_bulk(std::slice::from_ref(&task.id), &claim.claim_token) + .await + .unwrap(); + assert_eq!(completed, 1); + + let stats = queue.stats(Some("test")).await.unwrap(); + assert_eq!( + stats + .iter() + .find(|s| s.status == TaskStatus::Completed) + .map(|s| s.count), + Some(1) + ); +} + +#[tokio::test] +async fn test_enqueue_bulk_and_caller_id_idempotency() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 3); + let ids: Vec = tasks.iter().map(|t| t.id.clone()).collect(); + + let first = queue.enqueue_bulk(&tasks).await.unwrap(); + assert_eq!(first.len(), 3); + + // Re-enqueueing the same ids is a no-op. + let second = queue.enqueue_bulk(&tasks).await.unwrap(); + assert!(second.is_empty()); + + // The tasks are still claimable exactly once each. + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(10)) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 3); + let mut claimed_ids: Vec = claim.tasks.iter().map(|t| t.id.clone()).collect(); + claimed_ids.sort(); + let mut ids = ids; + ids.sort(); + assert_eq!(claimed_ids, ids); + + // A partial batch with one duplicate and one new task only inserts the new one. + let mut new_task = EnqueueTaskRequest::new("test", serde_json::json!({ "n": 99 })); + new_task.id = "01J00000000000000000000000".to_string(); + let partial = queue + .enqueue_bulk(&[tasks[0].clone(), new_task]) + .await + .unwrap(); + assert_eq!(partial.len(), 1); + assert_eq!(partial[0].id, "01J00000000000000000000000"); +} + +#[tokio::test] +async fn test_enqueue_single_is_idempotent() { + let (_env, _pool, queue) = new_queue().await; + + let task = EnqueueTaskRequest::new("test", serde_json::json!({})); + let first = queue.enqueue(task.clone()).await.unwrap(); + let second = queue.enqueue(task.clone()).await.unwrap(); + + assert_eq!(first.id, second.id); + assert_eq!(second.status, TaskStatus::Pending); +} + +#[tokio::test] +async fn test_enqueue_validation() { + let (_env, _pool, queue) = new_queue().await; + + let mut bad_id = EnqueueTaskRequest::new("test", serde_json::json!({})); + bad_id.id = "short".to_string(); + assert!(queue.enqueue_bulk(&[bad_id]).await.is_err()); + + let bad_type = EnqueueTaskRequest::new(" ", serde_json::json!({})); + assert!(queue.enqueue_bulk(&[bad_type]).await.is_err()); + + let mut bad_payload = EnqueueTaskRequest::new("test", serde_json::json!("not an object")); + bad_payload.payload = serde_json::json!("nope"); + assert!(queue.enqueue_bulk(&[bad_payload]).await.is_err()); + + let mut bad_weight = EnqueueTaskRequest::new("test", serde_json::json!({})); + bad_weight.weight = -1; + assert!(queue.enqueue_bulk(&[bad_weight]).await.is_err()); + + let mut bad_attempts = EnqueueTaskRequest::new("test", serde_json::json!({})); + bad_attempts.max_attempts = 0; + assert!(queue.enqueue_bulk(&[bad_attempts]).await.is_err()); +} + +#[tokio::test] +async fn test_claim_fifo_by_ulid() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 3); + let ids: Vec = tasks.iter().map(|t| t.id.clone()).collect(); + queue.enqueue_bulk(&tasks).await.unwrap(); + + let mut claimed = Vec::new(); + for _ in 0..3 { + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(1)) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1); + claimed.push(claim.tasks[0].id.clone()); + queue + .complete_bulk( + std::slice::from_ref(claimed.last().unwrap()), + &claim.claim_token, + ) + .await + .unwrap(); + } + + // ULIDs are generated monotonically, so claim order matches enqueue order. + assert_eq!(claimed, ids); +} + +#[tokio::test] +async fn test_claim_orders_by_priority() { + let (_env, pool, queue) = new_queue().await; + + let mut low = EnqueueTaskRequest::new("test", serde_json::json!({ "p": 5 })); + low.priority = 5; + let mut high = EnqueueTaskRequest::new("test", serde_json::json!({ "p": 10 })); + high.priority = 10; + queue.enqueue_bulk(&[low, high]).await.unwrap(); + + let first = queue + .claim_bulk(&pool, "test", "w", &claim_opts(1)) + .await + .unwrap(); + assert_eq!(first.tasks[0].payload["p"], 10); + + let second = queue + .claim_bulk(&pool, "test", "w", &claim_opts(1)) + .await + .unwrap(); + assert_eq!(second.tasks[0].payload["p"], 5); +} + +#[tokio::test] +async fn test_delayed_availability_blocks_claim() { + let (_env, pool, queue) = new_queue().await; + + let mut later = EnqueueTaskRequest::new("test", serde_json::json!({})); + later.available_at = OffsetDateTime::now_utc() + Duration::hours(1); + queue.enqueue_bulk(&[later]).await.unwrap(); + + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(1)) + .await + .unwrap(); + assert!(claim.tasks.is_empty()); + + sqlx::query("UPDATE tasks SET available_at = NOW() - INTERVAL '1 second'") + .execute(&pool) + .await + .unwrap(); + + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(1)) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1); +} + +#[tokio::test] +async fn test_weighted_batch_admits_first_even_over_cap() { + let (_env, pool, queue) = new_queue().await; + + let mut tasks = enqueue_tasks("test", 3); + tasks[0].weight = 5; + tasks[1].weight = 5; + tasks[2].weight = 10; + queue.enqueue_bulk(&tasks).await.unwrap(); + + let claim = queue + .claim_bulk( + &pool, + "test", + "w", + &ClaimOptions { + limit: 3, + max_weight: Some(10), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 2); + let mut claimed_ids: Vec = claim.tasks.iter().map(|t| t.id.clone()).collect(); + claimed_ids.sort(); + let mut expected_ids = vec![tasks[0].id.clone(), tasks[1].id.clone()]; + expected_ids.sort(); + assert_eq!(claimed_ids, expected_ids); + + // A single task heavier than the cap is still admitted (row 1 rule). + let mut heavy = EnqueueTaskRequest::new("test", serde_json::json!({})); + heavy.weight = 100; + queue.enqueue_bulk(&[heavy]).await.unwrap(); + + let claim = queue + .claim_bulk( + &pool, + "test", + "w", + &ClaimOptions { + limit: 1, + max_weight: Some(10), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1); +} + +#[tokio::test] +async fn test_concurrent_consumers_never_receive_same_task() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 10); + queue.enqueue_bulk(&tasks).await.unwrap(); + + let workers: Vec<_> = (0..20) + .map(|i| { + let queue = queue.clone(); + let pool = pool.clone(); + async move { + let claim = queue + .claim_bulk(&pool, "test", &format!("worker-{}", i), &claim_opts(1)) + .await + .unwrap(); + claim.tasks.into_iter().map(|t| t.id).collect::>() + } + }) + .collect(); + + let results: Vec = join_all(workers).await.into_iter().flatten().collect(); + assert_eq!( + results.len(), + 10, + "each of the 10 tasks claimed exactly once" + ); + + let mut unique = results.clone(); + unique.sort(); + unique.dedup(); + assert_eq!(unique.len(), 10, "no task was claimed twice"); +} + +#[tokio::test] +async fn test_concurrency_key_serializes_oldest_first() { + let (_env, pool, queue) = new_queue().await; + + let mut tasks = enqueue_tasks("test", 3); + for task in &mut tasks { + task.concurrency_key = Some("identity-1".to_string()); + } + queue.enqueue_bulk(&tasks).await.unwrap(); + + // Only the oldest task for the key is claimable at a time. + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(3)) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1); + assert_eq!(claim.tasks[0].id, tasks[0].id); + + // Newer tasks stay blocked while the oldest is running. + let blocked = queue + .claim_bulk(&pool, "test", "w", &claim_opts(3)) + .await + .unwrap(); + assert!(blocked.tasks.is_empty()); + + queue + .complete_bulk(&[tasks[0].id.clone()], &claim.claim_token) + .await + .unwrap(); + + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(3)) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1); + assert_eq!(claim.tasks[0].id, tasks[1].id); +} + +#[tokio::test] +async fn test_max_concurrency_cap_is_atomic() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 5); + queue.enqueue_bulk(&tasks).await.unwrap(); + + let claim = queue + .claim_bulk( + &pool, + "test", + "w", + &ClaimOptions { + limit: 5, + max_concurrency: Some(2), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 2); + + // Already at the cap: no further claims. + let blocked = queue + .claim_bulk( + &pool, + "test", + "w", + &ClaimOptions { + limit: 5, + max_concurrency: Some(2), + ..Default::default() + }, + ) + .await + .unwrap(); + assert!(blocked.tasks.is_empty()); + + // Completing frees capacity for the next oldest tasks. + let token = claim.claim_token.clone(); + let ids: Vec = claim.tasks.iter().map(|t| t.id.clone()).collect(); + queue.complete_bulk(&ids, &token).await.unwrap(); + + let claim = queue + .claim_bulk( + &pool, + "test", + "w", + &ClaimOptions { + limit: 5, + max_concurrency: Some(2), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 2); + let mut claimed_ids: Vec = claim.tasks.iter().map(|t| t.id.clone()).collect(); + claimed_ids.sort(); + let mut expected_ids = vec![tasks[2].id.clone(), tasks[3].id.clone()]; + expected_ids.sort(); + assert_eq!(claimed_ids, expected_ids); +} + +#[tokio::test] +async fn test_concurrent_claims_respect_concurrency_cap() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 5); + queue.enqueue_bulk(&tasks).await.unwrap(); + + let workers: Vec<_> = (0..10) + .map(|i| { + let queue = queue.clone(); + let pool = pool.clone(); + async move { + queue + .claim_bulk( + &pool, + "test", + &format!("worker-{}", i), + &ClaimOptions { + limit: 5, + max_concurrency: Some(1), + ..Default::default() + }, + ) + .await + .unwrap() + } + }) + .collect(); + + let claims: Vec<_> = join_all(workers).await; + let claimed: Vec = claims + .iter() + .flat_map(|c| c.tasks.iter().map(|t| t.id.clone())) + .collect(); + assert_eq!( + claimed.len(), + 1, + "cap of 1 admits exactly one task across concurrent claims" + ); + + // The slot stays occupied until the running task completes. + let holder = claims.iter().find(|c| !c.tasks.is_empty()).unwrap(); + let blocked = queue + .claim_bulk( + &pool, + "test", + "w", + &ClaimOptions { + limit: 5, + max_concurrency: Some(1), + ..Default::default() + }, + ) + .await + .unwrap(); + assert!(blocked.tasks.is_empty()); + + queue + .complete_bulk(&[holder.tasks[0].id.clone()], &holder.claim_token) + .await + .unwrap(); + + let claim = queue + .claim_bulk( + &pool, + "test", + "w", + &ClaimOptions { + limit: 5, + max_concurrency: Some(1), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1); +} + +#[tokio::test] +async fn test_claim_by_candidate_ids() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 3); + queue.enqueue_bulk(&tasks).await.unwrap(); + + // Claim only the third task by id, ignoring the earlier ones. + let claim = queue + .claim_bulk( + &pool, + "test", + "w", + &ClaimOptions { + candidate_ids: Some(vec![tasks[2].id.clone()]), + limit: 1, + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1); + assert_eq!(claim.tasks[0].id, tasks[2].id); +} + +#[tokio::test] +async fn test_heartbeat_renews_and_fences_stale_token() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 1); + queue.enqueue_bulk(&tasks).await.unwrap(); + + let claim = queue + .claim_bulk( + &pool, + "test", + "w", + &ClaimOptions { + limit: 1, + lease_seconds: 60, + ..Default::default() + }, + ) + .await + .unwrap(); + + assert!( + queue + .heartbeat(&tasks[0].id, &claim.claim_token, 120) + .await + .unwrap() + ); + + // Wrong token never renews. + assert!( + !queue + .heartbeat(&tasks[0].id, "01J00000000000000000000002", 120) + .await + .unwrap() + ); + + // An expired lease is not renewable even with the right token. + sqlx::query("UPDATE tasks SET lease_expires_at = NOW() - INTERVAL '1 second'") + .execute(&pool) + .await + .unwrap(); + assert!( + !queue + .heartbeat(&tasks[0].id, &claim.claim_token, 120) + .await + .unwrap() + ); +} + +#[tokio::test] +async fn test_terminal_writes_are_fenced_by_claim_token() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 2); + queue.enqueue_bulk(&tasks).await.unwrap(); + + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(2)) + .await + .unwrap(); + let ids: Vec = claim.tasks.iter().map(|t| t.id.clone()).collect(); + + // Wrong token: no rows affected. + assert_eq!( + queue + .complete_bulk(&ids, "01J00000000000000000000002") + .await + .unwrap(), + 0 + ); + let failed = queue + .fail_bulk(&ids, "01J00000000000000000000002", "stale worker", true, 0) + .await + .unwrap(); + assert!(failed.is_empty()); + + // Right token: both complete. + assert_eq!( + queue.complete_bulk(&ids, &claim.claim_token).await.unwrap(), + 2 + ); +} + +#[tokio::test] +async fn test_lease_expiry_recovery() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 1); + queue.enqueue_bulk(&tasks).await.unwrap(); + + let claim = queue + .claim_bulk( + &pool, + "test", + "w", + &ClaimOptions { + limit: 1, + lease_seconds: 60, + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1); + + sqlx::query("UPDATE tasks SET lease_expires_at = NOW() - INTERVAL '1 second'") + .execute(&pool) + .await + .unwrap(); + + let recovered = queue.recover_expired().await.unwrap(); + assert_eq!(recovered.len(), 1); + let task = &recovered[0]; + assert_eq!(task.status, TaskStatus::Pending); + assert!(task.claim_token.is_none()); + assert!(task.claimed_by.is_none()); + assert!(task.lease_expires_at.is_none()); + assert_eq!( + task.last_error.as_deref(), + Some("task lease expired; requeued") + ); + + // The requeued task is claimable again on attempt 2. + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(1)) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1); + assert_eq!(claim.tasks[0].attempt_count, 2); +} + +#[tokio::test] +async fn test_fail_retries_then_dead_letters() { + let (_env, pool, queue) = new_queue().await; + + let mut task = EnqueueTaskRequest::new("test", serde_json::json!({})); + task.max_attempts = 3; + queue.enqueue_bulk(&[task.clone()]).await.unwrap(); + + for attempt in 1..=3 { + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(1)) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1, "claim on attempt {}", attempt); + assert_eq!(claim.tasks[0].attempt_count, attempt); + + let failed = queue + .fail_bulk( + &[task.id.clone()], + &claim.claim_token, + &format!("error {}", attempt), + true, + 0, + ) + .await + .unwrap(); + let (id, status) = &failed[0]; + assert_eq!(id, &task.id); + if attempt < 3 { + assert_eq!( + *status, + TaskStatus::Pending, + "retry after attempt {}", + attempt + ); + } else { + assert_eq!( + *status, + TaskStatus::DeadLetter, + "dead-letter on final attempt" + ); + } + } + + // No claims remain. + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(1)) + .await + .unwrap(); + assert!(claim.tasks.is_empty()); + + let stats = queue.stats(Some("test")).await.unwrap(); + assert_eq!( + stats + .iter() + .find(|s| s.status == TaskStatus::DeadLetter) + .map(|s| s.count), + Some(1) + ); +} + +#[tokio::test] +async fn test_retry_delay_schedules_available_at() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 1); + queue.enqueue_bulk(&tasks).await.unwrap(); + + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(1)) + .await + .unwrap(); + queue + .fail_bulk( + &[tasks[0].id.clone()], + &claim.claim_token, + "slow down", + true, + 3600, + ) + .await + .unwrap(); + + // Not claimable while the retry delay is pending. + let blocked = queue + .claim_bulk(&pool, "test", "w", &claim_opts(1)) + .await + .unwrap(); + assert!(blocked.tasks.is_empty()); + + sqlx::query("UPDATE tasks SET available_at = NOW() - INTERVAL '1 second'") + .execute(&pool) + .await + .unwrap(); + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(1)) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1); +} + +#[tokio::test] +async fn test_non_retryable_failure_dead_letters_immediately() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 1); + queue.enqueue_bulk(&tasks).await.unwrap(); + + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(1)) + .await + .unwrap(); + let failed = queue + .fail_bulk( + &[tasks[0].id.clone()], + &claim.claim_token, + "permanent error", + false, + 0, + ) + .await + .unwrap(); + assert_eq!(failed[0].1, TaskStatus::DeadLetter); + + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(1)) + .await + .unwrap(); + assert!(claim.tasks.is_empty()); +} + +#[tokio::test] +async fn test_stats_counts_by_type_and_status() { + let (_env, pool, queue) = new_queue().await; + + let mut alpha = enqueue_tasks("alpha", 2); + let mut beta = enqueue_tasks("beta", 1); + let mut all = Vec::new(); + all.append(&mut alpha); + all.append(&mut beta); + queue.enqueue_bulk(&all).await.unwrap(); + + let claim = queue + .claim_bulk(&pool, "alpha", "w", &claim_opts(1)) + .await + .unwrap(); + queue + .complete_bulk(&[claim.tasks[0].id.clone()], &claim.claim_token) + .await + .unwrap(); + + let stats = queue.stats(None).await.unwrap(); + let count = |task_type: &str, status: TaskStatus| { + stats + .iter() + .find(|s| s.task_type == task_type && s.status == status) + .map(|s| s.count) + .unwrap_or(0) + }; + assert_eq!(count("alpha", TaskStatus::Pending), 1); + assert_eq!(count("alpha", TaskStatus::Completed), 1); + assert_eq!(count("beta", TaskStatus::Pending), 1); + + let beta_stats = queue.stats(Some("beta")).await.unwrap(); + assert_eq!(beta_stats.len(), 1); + assert_eq!(beta_stats[0].status, TaskStatus::Pending); +} + +#[tokio::test] +async fn test_cleanup_deletes_old_terminal_tasks() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 2); + queue.enqueue_bulk(&tasks).await.unwrap(); + + let claim = queue + .claim_bulk(&pool, "test", "w", &claim_opts(2)) + .await + .unwrap(); + let ids: Vec = claim.tasks.iter().map(|t| t.id.clone()).collect(); + queue.complete_bulk(&ids, &claim.claim_token).await.unwrap(); + + // Nothing old yet (completed_at was just set to now). + assert_eq!( + queue + .cleanup(OffsetDateTime::now_utc() - Duration::hours(1)) + .await + .unwrap(), + 0 + ); + + sqlx::query("UPDATE tasks SET completed_at = NOW() - INTERVAL '7 days'") + .execute(&pool) + .await + .unwrap(); + assert_eq!(queue.cleanup(OffsetDateTime::now_utc()).await.unwrap(), 2); + + // Pending tasks are never cleaned up. + let tasks = enqueue_tasks("test", 1); + queue.enqueue_bulk(&tasks).await.unwrap(); + assert_eq!(queue.cleanup(OffsetDateTime::now_utc()).await.unwrap(), 0); +} + +#[tokio::test] +async fn test_constraints_reject_invalid_lifecycle_states() { + let (_env, pool, _queue) = new_queue().await; + + // running requires claim token, claimed_by, and lease. + let err = sqlx::query( + r#" + INSERT INTO tasks (id, task_type, payload, status) + VALUES ('01J00000000000000000000001', 'test', '{}', 'running') + "#, + ) + .execute(&pool) + .await; + assert!(err.is_err()); + + // Unknown status. + let err = sqlx::query( + r#" + INSERT INTO tasks (id, task_type, payload, status) + VALUES ('01J00000000000000000000002', 'test', '{}', 'bogus') + "#, + ) + .execute(&pool) + .await; + assert!(err.is_err()); + + // attempt_count cannot exceed max_attempts. + let err = sqlx::query( + r#" + INSERT INTO tasks (id, task_type, payload, attempt_count, max_attempts) + VALUES ('01J00000000000000000000003', 'test', '{}', 5, 3) + "#, + ) + .execute(&pool) + .await; + assert!(err.is_err()); + + // payload must be a JSON object. + let err = sqlx::query( + r#" + INSERT INTO tasks (id, task_type, payload) + VALUES ('01J00000000000000000000004', 'test', '"not an object"'::jsonb) + "#, + ) + .execute(&pool) + .await; + assert!(err.is_err()); + + // id must be a 26-char ULID. + let err = sqlx::query( + r#" + INSERT INTO tasks (id, task_type, payload) + VALUES ('short', 'test', '{}') + "#, + ) + .execute(&pool) + .await; + assert!(err.is_err()); +} + +#[tokio::test] +async fn test_recover_expired_exhausted_attempts_dead_letters() { + let (_env, pool, queue) = new_queue().await; + + let mut task = EnqueueTaskRequest::new("test", serde_json::json!({})); + task.max_attempts = 1; + queue.enqueue_bulk(&[task.clone()]).await.unwrap(); + + let claim = queue + .claim_bulk( + &pool, + "test", + "w", + &ClaimOptions { + limit: 1, + lease_seconds: 60, + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1); + + // Model a task with a prior failure, then an expired lease: exhausted + // recovery must record the lease expiry as the latest event. + sqlx::query( + "UPDATE tasks SET lease_expires_at = NOW() - INTERVAL '1 second', last_error = 'previous failure'", + ) + .execute(&pool) + .await + .unwrap(); + + let recovered = queue.recover_expired().await.unwrap(); + assert_eq!(recovered.len(), 1); + assert_eq!(recovered[0].status, TaskStatus::DeadLetter); + assert!(recovered[0].completed_at.is_some()); + assert_eq!( + recovered[0].last_error.as_deref(), + Some("task lease expired; retries exhausted") + ); +} + +#[tokio::test] +async fn test_claim_validates_identity() { + let (_env, pool, queue) = new_queue().await; + + assert!( + queue + .claim_bulk(&pool, " ", "w", &claim_opts(1)) + .await + .is_err() + ); + assert!( + queue + .claim_bulk(&pool, "test", " ", &claim_opts(1)) + .await + .is_err() + ); +} + +#[tokio::test] +async fn test_sql_rejects_null_required_arguments() { + let (_env, pool, _queue) = new_queue().await; + + // NULL limit must not be interpreted as "no limit". + let err = sqlx::query( + "SELECT * FROM task_claim_bulk('test', NULL, NULL, NULL, NULL, 60, '01J00000000000000000000001', 'w')", + ) + .fetch_all(&pool) + .await; + assert!(err.is_err()); + + // NULL retryable must not silently dead-letter. + let err = sqlx::query( + "SELECT * FROM task_fail_bulk(ARRAY['01J00000000000000000000000'], '01J00000000000000000000001', 'boom', NULL, 0)", + ) + .fetch_all(&pool) + .await; + assert!(err.is_err()); + + // Blank claimed_by is rejected. + let err = sqlx::query( + "SELECT * FROM task_claim_bulk('test', NULL, 1, NULL, NULL, 60, '01J00000000000000000000001', '')", + ) + .fetch_all(&pool) + .await; + assert!(err.is_err()); + + // Blank task_type is rejected. + let err = sqlx::query( + "SELECT * FROM task_claim_bulk('', NULL, 1, NULL, NULL, 60, '01J00000000000000000000001', 'w')", + ) + .fetch_all(&pool) + .await; + assert!(err.is_err()); + + // NULL lease_seconds is rejected. + let err = sqlx::query( + "SELECT * FROM task_claim_bulk('test', NULL, 1, NULL, NULL, NULL, '01J00000000000000000000001', 'w')", + ) + .fetch_all(&pool) + .await; + assert!(err.is_err()); + + // Malformed claim token in complete is rejected. + let err = + sqlx::query("SELECT task_complete_bulk(ARRAY['01J00000000000000000000000'], 'short')") + .fetch_all(&pool) + .await; + assert!(err.is_err()); +} + +#[tokio::test] +async fn test_terminal_write_respects_wall_clock_lease_expiry() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 1); + queue.enqueue_bulk(&tasks).await.unwrap(); + + let claim = queue + .claim_bulk( + &pool, + "test", + "w", + &ClaimOptions { + limit: 1, + lease_seconds: 1, + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1); + + // Open a transaction, wait past the wall-clock lease expiry, then try to + // complete from inside that transaction. NOW() is frozen at transaction + // start, so fencing must use the statement timestamp of the write. + let mut tx = pool.begin().await.unwrap(); + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + let completed: i64 = sqlx::query_scalar("SELECT task_complete_bulk($1, $2)") + .bind(&[claim.tasks[0].id.clone()]) + .bind(&claim.claim_token) + .fetch_one(&mut *tx) + .await + .unwrap(); + assert_eq!( + completed, 0, + "lease expired in wall time; the terminal write must be fenced" + ); + tx.rollback().await.unwrap(); +} + +#[tokio::test] +async fn test_claim_lease_starts_after_advisory_lock_wait() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 1); + queue.enqueue_bulk(&tasks).await.unwrap(); + + // Hold the per-task-type advisory lock so the capped claim blocks on it + // for longer than the requested lease. + let mut lock_tx = pool.begin().await.unwrap(); + sqlx::query("SELECT pg_advisory_xact_lock(hashtext('test')::bigint)") + .execute(&mut *lock_tx) + .await + .unwrap(); + + let queue = queue.clone(); + let pool = pool.clone(); + let claim_handle = tokio::spawn(async move { + queue + .claim_bulk( + &pool, + "test", + "w", + &ClaimOptions { + limit: 1, + max_concurrency: Some(1), + lease_seconds: 1, + ..Default::default() + }, + ) + .await + .unwrap() + }); + + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + lock_tx.commit().await.unwrap(); + + let claim = claim_handle.await.unwrap(); + assert_eq!(claim.tasks.len(), 1); + let lease = claim.tasks[0].lease_expires_at.unwrap(); + assert!( + lease > OffsetDateTime::now_utc(), + "lease must start after the advisory lock wait, not be already expired" + ); +} + +#[tokio::test] +async fn test_terminal_write_waits_for_row_lock_then_fences() { + let (_env, pool, queue) = new_queue().await; + + let tasks = enqueue_tasks("test", 1); + queue.enqueue_bulk(&tasks).await.unwrap(); + + let claim = queue + .claim_bulk( + &pool, + "test", + "w", + &ClaimOptions { + limit: 1, + lease_seconds: 1, + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(claim.tasks.len(), 1); + let task_id = claim.tasks[0].id.clone(); + let claim_token = claim.claim_token.clone(); + + // Hold a row lock on the claimed task so the terminal write blocks on it + // until after the lease expires. + let mut lock_tx = pool.begin().await.unwrap(); + sqlx::query("UPDATE tasks SET last_error = 'held' WHERE id = $1") + .bind(&task_id) + .execute(&mut *lock_tx) + .await + .unwrap(); + + let queue = queue.clone(); + let complete_handle = + tokio::spawn(async move { queue.complete_bulk(&[task_id], &claim_token).await.unwrap() }); + + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + lock_tx.commit().await.unwrap(); + + let completed = complete_handle.await.unwrap(); + assert_eq!( + completed, 0, + "lease expired while waiting on the row lock; complete must be fenced" + ); +}