feat(dispatch): job lifecycle enqueue/claim/renew/reclaim with Lua (slice 2/6)

dispatch/job.py

@@ -0,0 +1,182 @@
+"""Job lifecycle: enqueue / claim / renew / reclaim (spec §7.2, §7.3, §8, §9).
+
+This is the Python skin over the Lua state transitions in ``scripts.py``. Every
+transition is atomic in Redis; this layer only picks the ``now`` clock (Python
+``time.time()``, passed into every script via ARGV — never a server-side TIME),
+walks ``jobs:leased`` for the time-gated orphan scan, and shapes payloads.
+
+Slice 2 of the delivery plan (§15.2) — dark: nothing calls these yet. The
+``complete``/``abort``/JE-landing transitions are slice 3; presigned ``code_url``
+signing on the payload is slice 4. Both are called out at their seams below.
+"""
+
+import time
+from typing import Dict, Optional
+
+from ulid import ULID
+
+from mongo.utils import RedisCache
+from . import params
+from . import redis_keys
+from . import scripts
+
+JOB_ID_PREFIX = 'jb_'
+
+# One shared RedisCache instance, mirroring runner.py: under fakeredis each
+# RedisCache owns an isolated dataset, so enqueue/claim/renew must share one to
+# stay coherent; in production they all share the pooled real Redis anyway.
+_cache: Optional[RedisCache] = None
+
+
+def _redis():
+    global _cache
+    if _cache is None:
+        _cache = RedisCache()
+    return _cache.client
+
+
+def _now() -> float:
+    return time.time()
+
+
+def _decode(value) -> Optional[str]:
+    if value is None:
+        return None
+    return value.decode() if isinstance(value, bytes) else value
+
+
+def enqueue_job(
+    submission_id: str,
+    problem_id: str,
+    language: int,
+    code_minio_path: str,
+    checker: str,
+    tasks_meta_json: str,
+) -> str:
+    """Enqueue a fresh judging job for ``submission_id`` (spec §8, §9).
+
+    Writes the job hash, points ``submission:<sid>:current_job`` at it (last
+    enqueue wins — that is exactly how a rejudge supersedes an in-flight job; the
+    stale one is destroyed lazily on the dispatch side by ``claim_pending``, INV4,
+    never proactively here), then LPUSHes onto the FIFO pending queue.
+    """
+    now = _now()
+    job_id = JOB_ID_PREFIX + str(ULID())
+
+    client = _redis()
+    job_key = redis_keys.job(job_id)
+
+    pipe = client.pipeline()
+    pipe.hset(
+        job_key,
+        mapping={
+            'submission_id': submission_id,
+            'problem_id': problem_id,
+            'language': language,
+            'code_minio_path': code_minio_path,
+            'checker': checker,
+            'tasks_meta_json': tasks_meta_json,
+            'leased_by': '',
+            'lease_deadline': '0',
+            'state': 'pending',
+            'attempts': 0,
+            'created_at': repr(now),
+            'last_error': '',
+        },
+    )
+    pipe.set(redis_keys.submission_current_job(submission_id), job_id)
+    pipe.lpush(redis_keys.JOBS_PENDING, job_id)
+    pipe.execute()
+
+    return job_id
+
+
+def claim_next_job(runner_id: str) -> Optional[Dict]:
+    """Hand ``runner_id`` its next job, or None when there is nothing to do.
+
+    Order matches spec §7.3: first a time-gated orphan scan (at most one runner
+    per ``ORPHAN_SCAN_INTERVAL_SEC`` window), whose first successful reclaim goes
+    straight to this caller; otherwise fall through to the pending queue.
+    """
+    now = _now()
+    client = _redis()
+
+    reclaimed = _orphan_scan(client, runner_id, now)
+    if reclaimed is not None:
+        return _payload(client, reclaimed)
+
+    job_id = scripts.load(client).claim_pending(
+        keys=[redis_keys.JOBS_PENDING, redis_keys.JOBS_LEASED],
+        args=[now, runner_id, params.LEASE_TTL_SEC],
+    )
+    if job_id is None:
+        return None
+    return _payload(client, _decode(job_id))
+
+
+def renew_lease(runner_id: str, job_id: str) -> bool:
+    """Extend the lease on ``job_id`` iff ``runner_id`` still owns it (spec §7.2).
+
+    Thin wrapper over ``renew_lease`` Lua; heartbeat wiring is slice 4.
+    """
+    now = _now()
+    result = scripts.load(_redis()).renew_lease(
+        keys=[redis_keys.JOBS_LEASED,
+              redis_keys.job(job_id)],
+        args=[job_id, runner_id, now, params.LEASE_TTL_SEC],
+    )
+    return result == 1
+
+
+def _orphan_scan(client, runner_id: str, now: float) -> Optional[str]:
+    """Time-gated sweep of ``jobs:leased``; return an id reclaimed for the caller.
+
+    The ``SET NX EX`` gate on ``dispatch:last_recovery`` lets at most one scan run
+    per window (spec §7.3 step 1). Within a winning scan, each expired candidate
+    is offered to ``reclaim_expired``: the first ``1`` (reclaimed to this caller)
+    is returned immediately; ``-1`` (attempts exhausted) candidates are skipped
+    and left in place — the JE landing sweep that removes them is slice 3.
+    """
+    acquired = client.set(
+        redis_keys.DISPATCH_LAST_RECOVERY,
+        '1',
+        nx=True,
+        ex=params.ORPHAN_SCAN_INTERVAL_SEC,
+    )
+    if not acquired:
+        return None
+
+    reclaim = scripts.load(client).reclaim_expired
+    for member in client.smembers(redis_keys.JOBS_LEASED):
+        job_id = _decode(member)
+        result = reclaim(
+            keys=[redis_keys.JOBS_LEASED,
+                  redis_keys.job(job_id)],
+            args=[
+                job_id,
+                runner_id,
+                now,
+                params.LEASE_TTL_SEC,
+                params.MAX_ATTEMPTS,
+            ],
+        )
+        if result == 1:
+            return job_id
+    return None
+
+
+def _payload(client, job_id: str) -> Optional[Dict]:
+    """Shape a job hash into a claim payload (spec §7.3), or None if it vanished.
+
+    Returns ``job_id`` plus the raw hash fields the HTTP layer will need. Typed
+    coercion and the presigned ``code_url`` are slice 4's concern, not this one's.
+    """
+    raw = client.hgetall(redis_keys.job(job_id))
+    if not raw:
+        return None
+    fields: Dict[str, Optional[str]] = {
+        _decode(k): _decode(v)
+        for k, v in raw.items()
+    }
+    fields['job_id'] = job_id
+    return fields

dispatch/scripts.py

@@ -0,0 +1,170 @@
+"""Lua scripts — one per job state transition (spec §8, §9).
+
+Each script is atomic on the Redis side, so a state transition can never be
+observed half-applied even under concurrent runners (INV2/INV3/INV4/INV5).
+Every script takes ``now`` via ARGV — never ``redis.call('TIME')`` — so tests
+are deterministic and fakeredis (which has no server clock for scripts) works.
+
+The scripts here cover slice 2 of the delivery plan (§15.2): ``claim_pending``,
+``renew_lease``, ``reclaim_expired``. ``abort_requeue`` belongs to slice 3 and
+is deliberately absent.
+
+Key handling: static keys (``jobs:pending``, ``jobs:leased``, a specific
+``job:<jb>`` hash) are passed in as KEYS by the Python layer via redis_keys.
+``claim_pending`` is the sole exception — it discovers a job id by RPOP and so
+must build that job's ``job:<jb>`` hash key and its
+``submission:<sid>:current_job`` pointer key from inside Lua; those two literal
+templates mirror ``redis_keys.job`` / ``redis_keys.submission_current_job`` and
+are the only place key names are spelled outside redis_keys.
+
+Corrupted state (partial Redis data loss — never produced by normal operation)
+is reaped lazily, in passing: the job hash is the single source of truth and
+the pending list / leased set are only indexes, so a script that stumbles on a
+dangling entry destroys or unindexes it and moves on instead of erroring. The
+affected submission stays Pending and is rescued by rejudge (spec §12).
+"""
+
+from dataclasses import dataclass
+from typing import Any
+
+# claim_pending — pending → leased, with dispatch-side currency destroy (INV4).
+#
+# KEYS[1] = jobs:pending (LIST)   KEYS[2] = jobs:leased (SET)
+# ARGV[1] = now   ARGV[2] = runner_id   ARGV[3] = lease_ttl_sec
+#
+# Loop popping the FIFO tail: a job whose hash has evaporated or lost its
+# submission_id (currency would be unverifiable — corrupted beyond judging) is
+# destroyed and skipped; a job that is no longer its submission's current_job is
+# destroyed on the spot (the dispatch-side half of INV4 — a superseded rejudge
+# job never runs); the first live current job is leased (attempts+1,
+# state=leased) and its id returned. A missing attempts field falls back to 0
+# instead of erroring — the id is already popped at that point, so a script
+# error would silently lose the job. Empty queue → nil.
+CLAIM_PENDING = '''
+local now = tonumber(ARGV[1])
+local runner = ARGV[2]
+local ttl = tonumber(ARGV[3])
+while true do
+  local jb = redis.call('RPOP', KEYS[1])
+  if not jb then
+    return nil
+  end
+  local job_key = 'job:' .. jb
+  local sid = redis.call('HGET', job_key, 'submission_id')
+  if not sid then
+    -- Hash evaporated (DEL is a no-op then) or corrupted beyond judging (no
+    -- submission_id means currency cannot be verified): destroy, move on.
+    redis.call('DEL', job_key)
+  else
+    local current = redis.call('GET', 'submission:' .. sid .. ':current_job')
+    if current == jb then
+      local attempts = (tonumber(redis.call('HGET', job_key, 'attempts')) or 0) + 1
+      redis.call('HSET', job_key,
+        'attempts', attempts,
+        'leased_by', runner,
+        'lease_deadline', now + ttl,
+        'state', 'leased')
+      redis.call('SADD', KEYS[2], jb)
+      return jb
+    else
+      redis.call('DEL', job_key)
+    end
+  end
+end
+'''
+
+# renew_lease — heartbeat extends an owned lease (spec §7.2).
+#
+# KEYS[1] = jobs:leased (SET)   KEYS[2] = job:<jb> (HASH)
+# ARGV[1] = job_id   ARGV[2] = runner_id   ARGV[3] = now   ARGV[4] = lease_ttl_sec
+#
+# Only when the job is still leased AND owned by the caller: push the deadline to
+# now+ttl, return 1. Anything else (not leased, wrong owner, evaporated hash) →
+# 0. An evaporated hash additionally reaps the ghost member from jobs:leased —
+# the hash is the source of truth, the set only an index. Never creates or
+# resurrects a key (SREM only removes) — HSET runs only on the confirmed-owner
+# path, where the hash provably exists.
+RENEW_LEASE = '''
+if redis.call('SISMEMBER', KEYS[1], ARGV[1]) == 0 then
+  return 0
+end
+if redis.call('EXISTS', KEYS[2]) == 0 then
+  -- Hash is the source of truth; a member whose hash evaporated is a ghost --
+  -- reap it so the orphan scan stops rescanning it forever.
+  redis.call('SREM', KEYS[1], ARGV[1])
+  return 0
+end
+local leased_by = redis.call('HGET', KEYS[2], 'leased_by')
+if leased_by == ARGV[2] then
+  redis.call('HSET', KEYS[2], 'lease_deadline', tonumber(ARGV[3]) + tonumber(ARGV[4]))
+  return 1
+end
+return 0
+'''
+
+# reclaim_expired — expired lease → new leaseholder, or converge (spec §7.3, §9).
+#
+# KEYS[1] = jobs:leased (SET)   KEYS[2] = job:<jb> (HASH)
+# ARGV[1] = job_id   ARGV[2] = new_runner   ARGV[3] = now
+# ARGV[4] = lease_ttl_sec   ARGV[5] = max_attempts
+#
+# Eligibility is decided ONLY by lease_deadline vs now (INV2) — runner liveness
+# never enters. If expired and attempts < max: attempts+1, hand to new_runner,
+# fresh deadline, return 1. If expired and attempts already >= max: return -1 and
+# leave the job UNTOUCHED (INV5 boundary; the JE landing sweep is slice 3, not
+# this script's job). Not leased / not expired → 0. An evaporated hash → 0 and
+# the ghost member is reaped from jobs:leased (hash is the source of truth); a
+# hash that still exists but lost its lease_deadline field stays tracked — a
+# runner may well be judging it, and tearing down live state risks losing the
+# job entirely.
+RECLAIM_EXPIRED = '''
+if redis.call('SISMEMBER', KEYS[1], ARGV[1]) == 0 then
+  return 0
+end
+if redis.call('EXISTS', KEYS[2]) == 0 then
+  -- Hash is the source of truth; a member whose hash evaporated is a ghost --
+  -- reap it so the orphan scan stops rescanning it forever.
+  redis.call('SREM', KEYS[1], ARGV[1])
+  return 0
+end
+local deadline = redis.call('HGET', KEYS[2], 'lease_deadline')
+if not deadline then
+  -- Hash alive but the field is gone: leave it tracked (see header note).
+  return 0
+end
+local now = tonumber(ARGV[3])
+if tonumber(deadline) >= now then
+  return 0
+end
+local attempts = tonumber(redis.call('HGET', KEYS[2], 'attempts'))
+if attempts >= tonumber(ARGV[5]) then
+  return -1
+end
+redis.call('HSET', KEYS[2],
+  'attempts', attempts + 1,
+  'leased_by', ARGV[2],
+  'lease_deadline', now + tonumber(ARGV[4]),
+  'state', 'leased')
+return 1
+'''
+
+
+@dataclass(frozen=True)
+class Scripts:
+    claim_pending: Any
+    renew_lease: Any
+    reclaim_expired: Any
+
+
+def load(client) -> Scripts:
+    """Register the per-transition scripts on ``client`` and return them.
+
+    ``register_script`` is cheap (it only stores the body and computes the SHA
+    lazily on first EVALSHA), so binding to the currently active shared client on
+    each use keeps scripts coherent with the ``_cache`` reset done per test.
+    """
+    return Scripts(
+        claim_pending=client.register_script(CLAIM_PENDING),
+        renew_lease=client.register_script(RENEW_LEASE),
+        reclaim_expired=client.register_script(RECLAIM_EXPIRED),
+    )