trading-system-docs/tolaria/c5-sync-service/job_claim.py

"""
AUTH.3E — job_claim.py
======================
Atomarer Claim / Lease / Concurrency für den Executor Command Channel.

Garantien:
  * Claim ist ATOMAR (SQLite-Transaktion mit Status-Bedingung) — kein TOCTOU.
  * Zwei SAVE-Worker: nur einer kann denselben Job claimen.
  * Zwei DELETE-Worker: nur einer kann denselben DELETE-Job claimen.
  * Ein bereits gestarteter DELETE wird bei Lease-Expiry NICHT blind wiederholt
    (-> OUTCOME_UNKNOWN, nicht zurück zu READY).
  * Stale Claims werden via Lease-Expiry erkannt.

Isoliert implementiert (KEIN produktiver Container).
"""

from __future__ import annotations

import sqlite3
import time
import uuid
from typing import Any, Dict, Optional

from job_state_machine import ST_CLAIMED, ST_READY

# ---------------------------------------------------------------------------
# Fehler
# ---------------------------------------------------------------------------
class ClaimError(Exception):
    """Basis-Fehler für Claim/Lease."""


class JobAlreadyClaimedError(ClaimError):
    """Job ist nicht claimbar (bereits geclaimt oder Lease aktiv)."""


class JobNotFoundError(ClaimError):
    pass


# ---------------------------------------------------------------------------
# Atomarer Claim
# ---------------------------------------------------------------------------
def atomic_claim(
    conn: sqlite3.Connection,
    job_id: str,
    worker_id: str,
    lease_seconds: int = 60,
) -> Dict[str, Any]:
    """
    Führt einen atomaren Claim aus: READY -> CLAIMED, nur wenn Lease abgelaufen
    oder nie gesetzt. Kein TOCTOU (Status-Bedingung in der UPDATE-WHERE-Klausel).

    Rückgabe: der geclaimte Job (als dict).
    Wirft JobAlreadyClaimedError, wenn der Job nicht claimbar ist.
    """
    now = int(time.time() * 1000)
    lease_until = now + lease_seconds * 1000
    claim_id = str(uuid.uuid4())
    try:
        cur = conn.execute(
            """
            UPDATE jobs
            SET state = ?, worker_id = ?, claim_id = ?, claimed_at = ?,
                lease_until = ?, attempt_count = attempt_count + 1, updated_at = ?
            WHERE job_id = ?
              AND state = ?
              AND (lease_until IS NULL OR lease_until < ?)
            """,
            (ST_CLAIMED, worker_id, claim_id, now, lease_until, now,
             job_id, ST_READY, now),
        )
        conn.commit()
    except sqlite3.IntegrityError as e:
        raise ClaimError(f"atomic_claim failed: {e}") from e
    if cur.rowcount == 0:
        raise JobAlreadyClaimedError(f"job not claimable: {job_id}")
    row = conn.execute(
        "SELECT * FROM jobs WHERE job_id = ?", (job_id,)
    ).fetchone()
    if row is None:
        raise JobNotFoundError(f"job not found after claim: {job_id}")
    return dict(row)


# ---------------------------------------------------------------------------
# Lease
# ---------------------------------------------------------------------------
def renew_lease(
    conn: sqlite3.Connection,
    job_id: str,
    lease_seconds: int = 60,
) -> None:
    """Verlängert die Lease eines CLAIMED/EXECUTING-Jobs."""
    now = int(time.time() * 1000)
    lease_until = now + lease_seconds * 1000
    conn.execute(
        "UPDATE jobs SET lease_until = ?, updated_at = ? WHERE job_id = ?",
        (lease_until, now, job_id),
    )
    conn.commit()


def is_lease_expired(lease_until: Optional[int], now: Optional[int] = None) -> bool:
    """True, wenn die Lease abgelaufen ist (oder nie gesetzt)."""
    if lease_until is None:
        return True
    now = now if now is not None else int(time.time() * 1000)
    return lease_until < now


# ---------------------------------------------------------------------------
# Crash-Recovery
# ---------------------------------------------------------------------------
def recover_stale_claims(
    conn: sqlite3.Connection,
    job_type: str,
    now: Optional[int] = None,
) -> int:
    """
    Findet stale CLAIMED-Jobs (Lease abgelaufen) und überführt sie:
      * SAVE: zurück zu READY (wieder claimbar, RETRYABLE)
      * DELETE: zu OUTCOME_UNKNOWN (NICHT blind wiederholen)

    Rückgabe: Anzahl der recovered Jobs.
    """
    from job_state_machine import ST_OUTCOME_UNKNOWN

    now = now if now is not None else int(time.time() * 1000)
    stale = conn.execute(
        "SELECT job_id FROM jobs WHERE state = ? AND lease_until IS NOT NULL AND lease_until < ?",
        (ST_CLAIMED, now),
    ).fetchall()
    recovered = 0
    for row in stale:
        job_id = row["job_id"]
        if job_type == "C5_SAVE_OBJECT":
            conn.execute(
                "UPDATE jobs SET state = ?, updated_at = ? WHERE job_id = ?",
                (ST_READY, now, job_id),
            )
        else:
            conn.execute(
                "UPDATE jobs SET state = ?, updated_at = ? WHERE job_id = ?",
                (ST_OUTCOME_UNKNOWN, now, job_id),
            )
        recovered += 1
    conn.commit()
    return recovered