[iris] Add event-replay testing system for transitions equivalence

Introduces a curated replay harness for ControllerTransitions:

- IrisEvent: a frozen dataclass union, one variant per public mutation
  method on ControllerTransitions, with a thin dispatcher that calls
  the matching method.
- SQL trace hook: ControllerDB now exposes a class-level
  _trace_callback slot that _configure() registers on every
  connection (production never sets it).
- Deterministic DB dump: per-table JSON sorted by primary key for
  state-equivalence assertions.
- 13 curated scenarios exercising submit/assign/run/succeed,
  retries, preemption, worker-failure cascade, coscheduled timeout,
  direct-provider lifecycle, prune sweep, endpoints, and
  reservation claims.
- pytest goldens (DB state strict, SQL trace informational) under
  lib/iris/tests/cluster/controller/replay/golden/, with a
  --update-goldens regeneration flag.
- CLI: ``uv run python -m iris.cluster.controller.replay.run
  [--seed=DIR] [--scenario=NAMES] --out=DIR`` writes db.json and
  sql.txt per scenario.

Companion to the rjpower/iris-sql-store branch: this PR delivers the
testing infrastructure that proves the in-flight refactor preserves
DB-state equivalence across all curated scenarios.

Author: rjpower

lib/iris/src/iris/cluster/controller/db.py

@@ -291,6 +291,14 @@ class ControllerDB:
     AUTH_DB_FILENAME = "auth.sqlite3"
     PROFILES_DB_FILENAME = "profiles.sqlite3"
 
+    # Class-level SQL trace hook: when set, every connection (writer + read pool)
+    # opened via _configure() registers this callback via
+    # ``conn.set_trace_callback``. Exposed for the replay-system tests
+    # (``iris.cluster.controller.replay``); production never sets this. Must be
+    # set BEFORE ``ControllerDB(...)`` is constructed so the writer connection
+    # picks it up — connections opened earlier will not retroactively trace.
+    _trace_callback: Callable[[str], None] | None = None
+
     def __init__(self, db_dir: Path):
         import time
 
@@ -412,6 +420,8 @@ def _configure(conn: sqlite3.Connection) -> None:
         conn.execute("PRAGMA synchronous = NORMAL")
         conn.execute("PRAGMA busy_timeout = 5000")
         conn.execute("PRAGMA foreign_keys = ON")
+        if (cb := ControllerDB._trace_callback) is not None:
+            conn.set_trace_callback(cb)
 
     def optimize(self) -> None:
         """Run PRAGMA optimize to refresh statistics for tables with stale data.

lib/iris/src/iris/cluster/controller/replay/__init__.py

@@ -0,0 +1,69 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Event-replay testing system for ControllerTransitions.
+
+This package defines a frozen ``IrisEvent`` union — one variant per public
+mutation method on ``ControllerTransitions`` — together with a dispatcher,
+a SQLite trace hook, a deterministic DB dump, and a curated set of
+scenarios. Used both as a pytest golden suite and as a CLI for diffing
+DB state across branches/checkpoints.
+"""
+
+from iris.cluster.controller.replay.db_dump import deterministic_dump
+from iris.cluster.controller.replay.dispatcher import apply_event
+from iris.cluster.controller.replay.events import (
+    AddEndpoint,
+    ApplyDirectProviderUpdates,
+    ApplyHeartbeatsBatch,
+    ApplyTaskUpdates,
+    BufferDirectKill,
+    CancelJob,
+    CancelTasksForTimeout,
+    DrainForDirectProvider,
+    IrisEvent,
+    MarkTaskUnschedulable,
+    PreemptTask,
+    QueueAssignments,
+    RegisterOrRefreshWorker,
+    RemoveEndpoint,
+    RemoveFinishedJob,
+    RemoveWorker,
+    ReplaceReservationClaims,
+    SubmitJob,
+    UpdateWorkerPings,
+)
+from iris.cluster.controller.replay.scenarios import (
+    SCENARIO_NAMES,
+    SCENARIOS,
+    run_scenario,
+)
+from iris.cluster.controller.replay.sql_trace import sql_tracing
+
+__all__ = [
+    "SCENARIOS",
+    "SCENARIO_NAMES",
+    "AddEndpoint",
+    "ApplyDirectProviderUpdates",
+    "ApplyHeartbeatsBatch",
+    "ApplyTaskUpdates",
+    "BufferDirectKill",
+    "CancelJob",
+    "CancelTasksForTimeout",
+    "DrainForDirectProvider",
+    "IrisEvent",
+    "MarkTaskUnschedulable",
+    "PreemptTask",
+    "QueueAssignments",
+    "RegisterOrRefreshWorker",
+    "RemoveEndpoint",
+    "RemoveFinishedJob",
+    "RemoveWorker",
+    "ReplaceReservationClaims",
+    "SubmitJob",
+    "UpdateWorkerPings",
+    "apply_event",
+    "deterministic_dump",
+    "run_scenario",
+    "sql_tracing",
+]

lib/iris/src/iris/cluster/controller/replay/__main__.py

@@ -0,0 +1,9 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Entry point so ``python -m iris.cluster.controller.replay`` works."""
+
+from iris.cluster.controller.replay.run import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

lib/iris/src/iris/cluster/controller/replay/db_dump.py

@@ -0,0 +1,77 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Deterministic table-by-table JSON dump of the controller DB.
+
+Produces a stable mapping ``{table_name: [row_dicts]}`` with rows sorted
+by primary key. This is the canonical "did the DB end up in the same
+state?" check used by the replay golden tests; SQL traces are reviewed
+informationally.
+
+Excludes SQLite internal tables and ``schema_migrations`` (which is
+populated unconditionally on ``apply_migrations`` and adds noise without
+testing any controller behavior).
+"""
+
+import base64
+from typing import Any
+
+from iris.cluster.controller.db import ControllerDB
+
+EXCLUDED_TABLES: frozenset[str] = frozenset({"schema_migrations"})
+"""Tables ignored by ``deterministic_dump`` — schema bookkeeping only."""
+
+
+def _encode(value: Any) -> Any:
+    """Render a SQLite cell value as a JSON-serializable scalar.
+
+    ``bytes`` columns (notably ``job_workdir_files.data``) become
+    base64 ASCII so the dump round-trips through ``json.dumps``.
+    """
+    if isinstance(value, (bytes, bytearray, memoryview)):
+        return base64.b64encode(bytes(value)).decode("ascii")
+    return value
+
+
+def _list_user_tables(db: ControllerDB) -> list[str]:
+    with db.read_snapshot() as snap:
+        rows = snap.fetchall(
+            "SELECT name FROM sqlite_master WHERE type = 'table' AND name NOT LIKE 'sqlite_%' ORDER BY name"
+        )
+    return [str(row["name"]) for row in rows if str(row["name"]) not in EXCLUDED_TABLES]
+
+
+def _primary_key_columns(db: ControllerDB, table: str) -> list[str]:
+    """Return the table's primary key columns in declared order.
+
+    Falls back to *all* columns (also in declared order) when the table
+    has no PK so dumps remain stable regardless of insert order.
+    """
+    with db.read_snapshot() as snap:
+        rows = snap.fetchall(f"PRAGMA table_info({table})")
+    pk_cols = sorted(
+        ((int(row["pk"]), str(row["name"]), int(row["cid"])) for row in rows if int(row["pk"]) > 0),
+        key=lambda triple: triple[0],
+    )
+    if pk_cols:
+        return [name for _, name, _ in pk_cols]
+    # No PK declared — sort by every column for determinism.
+    return [str(row["name"]) for row in sorted(rows, key=lambda r: int(r["cid"]))]
+
+
+def deterministic_dump(db: ControllerDB) -> dict[str, list[dict[str, Any]]]:
+    """Dump every user table as ``{table: [row_dicts]}``.
+
+    Rows are returned as ordinary dicts in column-declaration order and
+    sorted by primary key (or every column when no PK exists). Bytes
+    columns are base64-encoded. Used as the canonical state-equivalence
+    check for replay scenarios.
+    """
+    out: dict[str, list[dict[str, Any]]] = {}
+    for table in _list_user_tables(db):
+        pk = _primary_key_columns(db, table)
+        order = ", ".join(pk)
+        with db.read_snapshot() as snap:
+            rows = snap.fetchall(f"SELECT * FROM {table} ORDER BY {order}")
+        out[table] = [{key: _encode(row[key]) for key in row.keys()} for row in rows]
+    return out

lib/iris/src/iris/cluster/controller/replay/dispatcher.py

@@ -0,0 +1,91 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Dispatch an :class:`IrisEvent` to the matching ``ControllerTransitions`` method.
+
+This module targets ``main`` semantics: each transition method opens its
+own ``with self._db.transaction()`` block and does not take a ``cur``
+argument. The branch ``rjpower/iris-sql-store`` rewrites this single
+file into the cur-passing form.
+"""
+
+from typing import Any
+
+from iris.cluster.controller.replay.events import (
+    AddEndpoint,
+    ApplyDirectProviderUpdates,
+    ApplyHeartbeatsBatch,
+    ApplyTaskUpdates,
+    BufferDirectKill,
+    CancelJob,
+    CancelTasksForTimeout,
+    DrainForDirectProvider,
+    IrisEvent,
+    MarkTaskUnschedulable,
+    PreemptTask,
+    QueueAssignments,
+    RegisterOrRefreshWorker,
+    RemoveEndpoint,
+    RemoveFinishedJob,
+    RemoveWorker,
+    ReplaceReservationClaims,
+    SubmitJob,
+    UpdateWorkerPings,
+)
+from iris.cluster.controller.transitions import ControllerTransitions
+
+
+def apply_event(transitions: ControllerTransitions, event: IrisEvent) -> Any:
+    """Dispatch ``event`` to the matching method on ``transitions``.
+
+    Returns whatever the underlying method returns (``TxResult``,
+    ``SubmitJobResult``, ``DirectProviderBatch``, etc.). The dispatcher
+    is intentionally thin — its purpose is to keep scenarios free of
+    branch-specific calling-convention details.
+    """
+    match event:
+        case SubmitJob(job_id, request, ts):
+            return transitions.submit_job(job_id, request, ts)
+        case CancelJob(job_id, reason):
+            return transitions.cancel_job(job_id, reason)
+        case RegisterOrRefreshWorker(worker_id, address, metadata, ts, slice_id, scale_group):
+            return transitions.register_or_refresh_worker(
+                worker_id=worker_id,
+                address=address,
+                metadata=metadata,
+                ts=ts,
+                slice_id=slice_id,
+                scale_group=scale_group,
+            )
+        case QueueAssignments(assignments, direct_dispatch):
+            return transitions.queue_assignments(assignments, direct_dispatch=direct_dispatch)
+        case ApplyTaskUpdates(request):
+            return transitions.apply_task_updates(request)
+        case ApplyHeartbeatsBatch(requests):
+            return transitions.apply_heartbeats_batch(requests)
+        case PreemptTask(task_id, reason):
+            return transitions.preempt_task(task_id, reason)
+        case CancelTasksForTimeout(task_ids, reason):
+            return transitions.cancel_tasks_for_timeout(set(task_ids), reason)
+        case MarkTaskUnschedulable(task_id, reason):
+            return transitions.mark_task_unschedulable(task_id, reason)
+        case RemoveFinishedJob(job_id):
+            return transitions.remove_finished_job(job_id)
+        case RemoveWorker(worker_id):
+            return transitions.remove_worker(worker_id)
+        case UpdateWorkerPings(snapshots):
+            return transitions.update_worker_pings(snapshots)
+        case DrainForDirectProvider(max_promotions):
+            return transitions.drain_for_direct_provider(max_promotions)
+        case ApplyDirectProviderUpdates(updates):
+            return transitions.apply_direct_provider_updates(updates)
+        case BufferDirectKill(task_id):
+            return transitions.buffer_direct_kill(task_id)
+        case AddEndpoint(endpoint):
+            return transitions.add_endpoint(endpoint)
+        case RemoveEndpoint(endpoint_id):
+            return transitions.remove_endpoint(endpoint_id)
+        case ReplaceReservationClaims(claims):
+            return transitions.replace_reservation_claims(claims)
+        case _:
+            raise TypeError(f"unhandled IrisEvent variant: {type(event).__name__}")