[iris] Add event-replay testing system for transitions equivalence (#5165)

## Summary

Adds a pytest-driven replay framework under
`lib/iris/tests/cluster/controller/replay/` that drives a fresh
`ControllerDB` through curated mutation sequences and asserts the
final DB state against committed goldens. **Pure test infrastructure
— no production-code changes.**

### Workflow

Generate goldens once on `main` (this PR's commit):

```
uv run pytest lib/iris/tests/cluster/controller/replay/ --update-goldens
```

The committed golden tree then acts as a behavioral fingerprint of
`main`. Any branch that preserves observable DB state for all 13
scenarios keeps the tests green without regeneration.

Author: rjpower

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

@@ -0,0 +1,2 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0

lib/iris/tests/cluster/controller/replay/conftest.py

@@ -0,0 +1,14 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Pytest hooks for the replay golden tests."""
+
+
+def pytest_addoption(parser) -> None:
+    """Register ``--update-goldens`` to regenerate the committed golden files."""
+    parser.addoption(
+        "--update-goldens",
+        action="store_true",
+        default=False,
+        help="Regenerate replay golden files instead of asserting against them.",
+    )

lib/iris/tests/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/tests/cluster/controller/replay/events.py

@@ -0,0 +1,209 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""``IrisEvent`` dataclass union + ``apply_event`` dispatcher.
+
+Each variant captures the arguments of one public mutation method on
+``ControllerTransitions``. ``apply_event`` opens a write transaction
+and invokes the matching method.
+
+Multi-transaction orchestrators (``fail_workers``, ``prune_old_data``)
+and ``*_for_test`` helpers are intentionally excluded — scenarios call
+those methods directly when needed.
+"""
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any
+
+from iris.cluster.controller.schema import EndpointRow
+from iris.cluster.controller.transitions import (
+    Assignment,
+    ControllerTransitions,
+    HeartbeatApplyRequest,
+    ReservationClaim,
+    TaskUpdate,
+)
+from iris.cluster.types import JobName, WorkerId
+from iris.rpc import controller_pb2, job_pb2
+from rigging.timing import Timestamp
+
+
+@dataclass(frozen=True, slots=True)
+class SubmitJob:
+    job_id: JobName
+    request: controller_pb2.Controller.LaunchJobRequest
+    ts: Timestamp
+
+
+@dataclass(frozen=True, slots=True)
+class CancelJob:
+    job_id: JobName
+    reason: str
+
+
+@dataclass(frozen=True, slots=True)
+class RegisterOrRefreshWorker:
+    worker_id: WorkerId
+    address: str
+    metadata: job_pb2.WorkerMetadata
+    ts: Timestamp
+    slice_id: str = ""
+    scale_group: str = ""
+
+
+@dataclass(frozen=True, slots=True)
+class QueueAssignments:
+    assignments: list[Assignment]
+    direct_dispatch: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class ApplyTaskUpdates:
+    request: HeartbeatApplyRequest
+
+
+@dataclass(frozen=True, slots=True)
+class ApplyHeartbeatsBatch:
+    requests: list[HeartbeatApplyRequest]
+
+
+@dataclass(frozen=True, slots=True)
+class PreemptTask:
+    task_id: JobName
+    reason: str
+
+
+@dataclass(frozen=True, slots=True)
+class CancelTasksForTimeout:
+    task_ids: frozenset[JobName]
+    reason: str
+
+
+@dataclass(frozen=True, slots=True)
+class MarkTaskUnschedulable:
+    task_id: JobName
+    reason: str
+
+
+@dataclass(frozen=True, slots=True)
+class RemoveFinishedJob:
+    job_id: JobName
+
+
+@dataclass(frozen=True, slots=True)
+class RemoveWorker:
+    worker_id: WorkerId
+
+
+@dataclass(frozen=True, slots=True)
+class UpdateWorkerPings:
+    snapshots: Mapping[WorkerId, job_pb2.WorkerResourceSnapshot | None]
+
+
+@dataclass(frozen=True, slots=True)
+class DrainForDirectProvider:
+    max_promotions: int = 16
+
+
+@dataclass(frozen=True, slots=True)
+class ApplyDirectProviderUpdates:
+    updates: list[TaskUpdate]
+
+
+@dataclass(frozen=True, slots=True)
+class BufferDirectKill:
+    task_id: str
+
+
+@dataclass(frozen=True, slots=True)
+class AddEndpoint:
+    endpoint: EndpointRow
+
+
+@dataclass(frozen=True, slots=True)
+class RemoveEndpoint:
+    endpoint_id: str
+
+
+@dataclass(frozen=True, slots=True)
+class ReplaceReservationClaims:
+    claims: dict[WorkerId, ReservationClaim]
+
+
+IrisEvent = (
+    SubmitJob
+    | CancelJob
+    | RegisterOrRefreshWorker
+    | QueueAssignments
+    | ApplyTaskUpdates
+    | ApplyHeartbeatsBatch
+    | PreemptTask
+    | CancelTasksForTimeout
+    | MarkTaskUnschedulable
+    | RemoveFinishedJob
+    | RemoveWorker
+    | UpdateWorkerPings
+    | DrainForDirectProvider
+    | ApplyDirectProviderUpdates
+    | BufferDirectKill
+    | AddEndpoint
+    | RemoveEndpoint
+    | ReplaceReservationClaims
+)
+
+
+def apply_event(transitions: ControllerTransitions, event: IrisEvent) -> Any:
+    """Dispatch ``event`` to the matching method on ``ControllerTransitions``.
+
+    Each transition method opens its own transaction — the dispatcher is
+    a thin match on the event variant. Returns whatever the underlying
+    method returns (``TxResult``, ``SubmitJobResult``,
+    ``DirectProviderBatch``, etc.).
+    """
+    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__}")