"""Worker dequeue helper. The worker picks one pending ``workflow_attempts`` row, marks it ``in_progress`` with lock metadata, and returns it. The atomic "pick + mark" is the core multi-machine coordination primitive. Per plan v5/v6: - Postgres: ``SELECT … FOR UPDATE SKIP LOCKED LIMIT 1`` is the canonical primitive. Multiple workers race; only one acquires each row. - SQLite: no SKIP LOCKED. Workers serialise via ``BEGIN IMMEDIATE`` + ``UPDATE … WHERE status='pending'`` (RETURNING clause; SQLite 3.35+). For single-machine SQLite tests this is fine; multi-machine SQLite over NFS is explicitly NOT supported (per plan v9: Postgres is required for multi-machine). Pickup-count guard (v6 blocker): the dequeue UPDATE bumps ``pickup_count``. A separate query (master-side) finds attempts where ``pickup_count >= MAX_PICKUPS`` and transitions the workflow to STUCK. """ from __future__ import annotations import os from dataclasses import dataclass from datetime import datetime, timezone from sqlalchemy import bindparam, text from sqlalchemy.engine import Engine from sqlalchemy.orm import Session from .models import WorkflowAttempt # Per-attempt pickup limit. If a worker dequeues an attempt and # fails (status='failed' outcome='worker-internal-error'), it can be # re-pended by the reaper; the reaper bumps pickup_count. Beyond # MAX_PICKUPS, the master transitions # the workflow to STUCK with reason='attempt-pickup-exhausted'. DEFAULT_MAX_PICKUPS = int(os.environ.get("CONTROLLER_MAX_ATTEMPT_PICKUPS", "3")) @dataclass class DequeueResult: """What the worker gets back from dequeue. ``acquired`` is True iff a row was claimed. On False, no rows available (worker sleeps + retries) — ``reason`` distinguishes 'empty queue' from 'all candidate rows exceeded pickup limit'. """ acquired: bool attempt_id: int | None = None workflow_id: int | None = None role: str | None = None tier: int | None = None pickup_count: int | None = None reason: str | None = None def _now() -> datetime: return datetime.now(timezone.utc) def dequeue_one( session: Session, *, instance_id: str, roles: list[str], max_pickups: int = DEFAULT_MAX_PICKUPS, ) -> DequeueResult: """Atomically pick the oldest pending attempt this worker can handle and stamp it as in_progress with the worker's identity. ``instance_id`` format: ``{hostname}/{pid}/{worker_uuid}`` (v9: slash-delimited so IPv6 hostnames don't trip parsing). ``roles`` is the list of role values the worker accepts. Returns ``DequeueResult(acquired=False)`` if no eligible row. Returns ``DequeueResult(acquired=True, attempt_id=..., ...)`` on success — the row is now locked under ``instance_id`` and the worker holds the heartbeat responsibility. Dialect routing: - Postgres: ``SELECT … FOR UPDATE SKIP LOCKED`` (atomic dequeue across racing workers — multi-machine safe). - SQLite: ``UPDATE … WHERE attempt_id = (SELECT … LIMIT 1) RETURNING``. Under SQLite's default BEGIN DEFERRED isolation, concurrent dequeues from multiple processes CAN race — the loser hits SQLITE_BUSY (and retries via ``busy_timeout`` = 5s). This is acceptable for single-host dev. MULTI-MACHINE DEPLOYMENTS MUST USE POSTGRES; see ``db/session.py`` + the deploy RUNBOOK. The session is committed by the caller (``session_scope``). """ if not roles: return DequeueResult(acquired=False, reason="no_roles_configured") dialect = session.bind.dialect.name if session.bind else "sqlite" now = _now() if dialect == "postgresql": return _dequeue_postgres(session, instance_id, roles, max_pickups, now) # SQLite (or any other; fall through with a generic implementation) return _dequeue_sqlite(session, instance_id, roles, max_pickups, now) def _dequeue_postgres( session: Session, instance_id: str, roles: list[str], max_pickups: int, now: datetime, ) -> DequeueResult: """Postgres path: SELECT FOR UPDATE SKIP LOCKED + UPDATE.""" # Two-statement transactional dequeue. We use raw SQL because # SQLAlchemy ORM's auto-flush behavior would re-query. select_sql = text(""" SELECT attempt_id, workflow_id, role, tier, pickup_count FROM workflow_attempts WHERE status = 'pending' AND role = ANY(:roles) AND pickup_count < :max_pickups ORDER BY created_at ASC LIMIT 1 FOR UPDATE SKIP LOCKED """).bindparams(bindparam("roles"), bindparam("max_pickups")) row = session.execute( select_sql, {"roles": roles, "max_pickups": max_pickups} ).first() if row is None: return DequeueResult(acquired=False, reason="no_pending_eligible") update_sql = text(""" UPDATE workflow_attempts SET status = 'in_progress', locked_by_instance = :instance, locked_at = :now, lock_heartbeat_at = :now, started_at = :now WHERE attempt_id = :attempt_id """) # Note: pickup_count is NOT bumped here. It records the number of # times the attempt has been REAPED (reset to pending after a stale # heartbeat). A worker that crashes mid-attempt bumps pickup_count # via the reaper; a worker that simply picks up an attempt and runs # it cleanly does NOT bump. Otherwise N normal pickups would STUCK # the workflow at MAX_PICKUPS even with no failures. session.execute( update_sql, {"instance": instance_id, "now": now, "attempt_id": row.attempt_id}, ) return DequeueResult( acquired=True, attempt_id=row.attempt_id, workflow_id=row.workflow_id, role=row.role, tier=row.tier, pickup_count=row.pickup_count or 0, ) def _dequeue_sqlite( session: Session, instance_id: str, roles: list[str], max_pickups: int, now: datetime, ) -> DequeueResult: """SQLite path: serialise on BEGIN IMMEDIATE + UPDATE. SQLite 3.35+ supports RETURNING, which makes this atomic in one statement. Older SQLite would need a SELECT-then-UPDATE, accepting a tight TOCTOU window — fine for tests but the test DB pins sqlite ≥3.35 in practice (Python 3.13 ships 3.43+). """ # SQLite doesn't have ANY(:roles); emit an IN-clause via parameter # expansion. role_clauses = ",".join(f":role_{i}" for i in range(len(roles))) update_sql = text(f""" UPDATE workflow_attempts SET status = 'in_progress', locked_by_instance = :instance, locked_at = :now, lock_heartbeat_at = :now, started_at = :now WHERE attempt_id = ( SELECT attempt_id FROM workflow_attempts WHERE status = 'pending' AND role IN ({role_clauses}) AND pickup_count < :max_pickups ORDER BY created_at ASC LIMIT 1 ) RETURNING attempt_id, workflow_id, role, tier, pickup_count """) # pickup_count is NOT bumped on dequeue — see the postgres path # for the rationale. params = { "instance": instance_id, "now": now, "max_pickups": max_pickups, **{f"role_{i}": r for i, r in enumerate(roles)}, } row = session.execute(update_sql, params).first() if row is None: return DequeueResult(acquired=False, reason="no_pending_eligible") return DequeueResult( acquired=True, attempt_id=row.attempt_id, workflow_id=row.workflow_id, role=row.role, tier=row.tier, pickup_count=row.pickup_count, )