[iris] Lift transactions to entrypoints; ControllerTransitions methods take cur

Transactions used to be opened ~22 times inside ControllerTransitions
methods. That had three problems: (1) RPC handlers couldn't compose
multiple transitions atomically — the read-validate-then-write pattern
left a TOCTOU window; (2) transitions.py had a self._db backdoor next
to its self._store, blurring the layering rule; (3) the per-method
``with self._db.transaction()`` boilerplate said the same thing 22
times.

Move tx scope to the caller. Every "one-event" transition now takes a
``cur: TransactionCursor`` (or ``snap: Tx`` for read paths); the
RPC handlers in service.py and the loop iterations in controller.py
each open exactly one ``with self._store.transaction() as cur:``
around their work. Multi-tx orchestrators (``fail_workers`` chunked,
``prune_old_data`` per-row) keep their own internal tx loops because
that's their reason for existing. Test helpers prefixed
``_for_test`` keep their auto-tx wrappers since tests call them ad-hoc.

Post-commit work that used to live at the tail of methods (attribute
cache invalidation in register/remove worker, the ``worker_registered``
audit log) moves into ``cur.on_commit(...)`` hooks so it only fires
once the data is durable — and a rolled-back transaction can't leave
the in-memory scheduling cache ahead of the DB row.

The ``self._db`` field on ``ControllerTransitions`` becomes a thin
``@property`` that delegates to ``self._store._db``. Production code
in transitions.py never touches it; the property only exists so
existing test helpers that read directly via ``state._db.snapshot()``
keep working. ``ControllerStore.optimize()`` is added so transitions'
post-prune ``self._store.optimize()`` call has a place to land.

Tests: ~190 call sites in tests/ updated to the cur-passing form (most
mechanically via a regex-based transformer). conftest helpers
``submit_job`` / ``submit_direct_job`` / ``register_worker`` /
``dispatch_task`` / ``transition_task`` / ``sync_k8s`` open a tx and
pass cur. ``state._db.snapshot()`` reads in conftest switch to
``state._db.read_snapshot()`` so they don't try to BEGIN on the writer
connection from inside an open writer tx.

Verification: 850 controller-suite tests pass (one pre-existing
failure unchanged); pyrefly clean for the controller package; full
iris test suite green.

Author: rjpower

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

@@ -1456,13 +1456,16 @@ def _sync_direct_provider(self) -> None:
         assert isinstance(self._provider, K8sTaskProvider)
         provider = self._provider
         max_promotions = self._promotion_bucket.available
-        batch = self._transitions.drain_for_direct_provider(
-            max_promotions=max_promotions,
-        )
+        with self._store.transaction() as cur:
+            batch = self._transitions.drain_for_direct_provider(
+                cur,
+                max_promotions=max_promotions,
+            )
         if batch.tasks_to_run:
             self._promotion_bucket.try_acquire(len(batch.tasks_to_run))
         result = provider.sync(batch)
-        tx_result = self._transitions.apply_direct_provider_updates(result.updates)
+        with self._store.transaction() as cur:
+            tx_result = self._transitions.apply_direct_provider_updates(cur, result.updates)
         self._provider_scheduling_events = list(result.scheduling_events) if result.scheduling_events else []
         self._provider_capacity = result.capacity
         if tx_result.tasks_to_kill:
@@ -1625,7 +1628,8 @@ def _cleanup_stale_claims(self, claims: dict[WorkerId, ReservationClaim] | None
         for wid in stale:
             del claims[wid]
         if stale and persisted:
-            self._transitions.replace_reservation_claims(claims)
+            with self._store.transaction() as cur:
+                self._transitions.replace_reservation_claims(cur, claims)
             log_event("reservation_claims_cleaned", "controller", count=len(stale))
         return bool(stale)
 
@@ -1673,7 +1677,8 @@ def _claim_workers_for_reservations(self, claims: dict[WorkerId, ReservationClai
                     changed = True
                     break
         if changed and persisted:
-            self._transitions.replace_reservation_claims(claims)
+            with self._store.transaction() as cur:
+                self._transitions.replace_reservation_claims(cur, claims)
             log_event("reservation_claims_updated", "controller", total_claims=len(claims))
         return changed
 
@@ -1769,7 +1774,8 @@ def _refresh_reservation_claims(self) -> dict[WorkerId, ReservationClaim]:
             if self._config.dry_run:
                 logger.info("[DRY-RUN] Would update %d reservation claims", len(claims))
             else:
-                self._transitions.replace_reservation_claims(claims)
+                with self._store.transaction() as cur:
+                    self._transitions.replace_reservation_claims(cur, claims)
         return claims
 
     def _read_scheduling_state(self) -> _SchedulingStateRead:
@@ -1982,7 +1988,10 @@ def _apply_preemptions(
             )
             preemptions = _run_preemption_pass(unscheduled, running_info, context)
             for preemptor_name, victim_id in preemptions:
-                preempt_result = self._transitions.preempt_task(victim_id, reason=f"Preempted by {preemptor_name}")
+                with self._store.transaction() as cur:
+                    preempt_result = self._transitions.preempt_task(
+                        cur, victim_id, reason=f"Preempted by {preemptor_name}"
+                    )
                 self.kill_tasks_on_workers(preempt_result.tasks_to_kill)
             if preemptions:
                 logger.info("Preemption pass: %d tasks preempted", len(preemptions))
@@ -2052,7 +2061,8 @@ def _enforce_execution_timeouts(self) -> None:
         for task in timed_out:
             logger.warning("Task %s exceeded execution timeout, killing", task.task_id)
         task_ids = {t.task_id for t in timed_out}
-        result = self._transitions.cancel_tasks_for_timeout(task_ids, reason="Execution timeout exceeded")
+        with self._store.transaction() as cur:
+            result = self._transitions.cancel_tasks_for_timeout(cur, task_ids, reason="Execution timeout exceeded")
         if result.tasks_to_kill:
             self.kill_tasks_on_workers(result.tasks_to_kill, result.task_kill_workers)
 
@@ -2067,10 +2077,12 @@ def _mark_task_unschedulable(self, task: TaskRow) -> None:
         else:
             timeout = None
         logger.warning(f"Task {task.task_id} exceeded scheduling timeout ({timeout}), marking as UNSCHEDULABLE")
-        result = self._transitions.mark_task_unschedulable(
-            task.task_id,
-            reason=f"Scheduling timeout exceeded ({timeout})",
-        )
+        with self._store.transaction() as cur:
+            result = self._transitions.mark_task_unschedulable(
+                cur,
+                task.task_id,
+                reason=f"Scheduling timeout exceeded ({timeout})",
+            )
         if result.tasks_to_kill:
             self.kill_tasks_on_workers(result.tasks_to_kill, result.task_kill_workers)
 
@@ -2099,8 +2111,9 @@ def kill_tasks_on_workers(
             self._stop_tasks_direct(task_ids, task_kill_workers)
             return
         # K8s: buffer direct kills for the provider sync loop.
-        for task_id in task_ids:
-            self._transitions.buffer_direct_kill(task_id.to_wire())
+        with self._store.transaction() as cur:
+            for task_id in task_ids:
+                self._transitions.buffer_direct_kill(cur, task_id.to_wire())
 
     # =========================================================================
     # Worker lifecycle RPC dispatch (StartTasks / StopTasks / Ping / PollTasks)
@@ -2116,7 +2129,8 @@ def _dispatch_assignments_direct(
                 logger.info("[DRY-RUN] Would assign task %s to worker %s", task_id, worker_id)
             return
         command = [Assignment(task_id=task_id, worker_id=worker_id) for task_id, worker_id in assignments]
-        result = self._transitions.queue_assignments(command, direct_dispatch=True)
+        with self._store.transaction() as cur:
+            result = self._transitions.queue_assignments(cur, command, direct_dispatch=True)
 
         # Group StartTasks payloads by (worker_id, address)
         by_worker: dict[tuple[WorkerId, str], list[job_pb2.RunTaskRequest]] = {}
@@ -2245,7 +2259,8 @@ def _run_ping_loop(self, stop_event: threading.Event) -> None:
                         self._health.ping(result.worker_id, healthy=True)
                         ping_snapshots[result.worker_id] = result.resource_snapshot if update_resources else None
 
-                self._transitions.update_worker_pings(ping_snapshots)
+                with self._store.transaction() as cur:
+                    self._transitions.update_worker_pings(cur, ping_snapshots)
 
                 unhealthy = self._health.workers_over_threshold()
                 if unhealthy:
@@ -2268,7 +2283,8 @@ def _poll_all_workers(self) -> None:
         """Poll all workers for task state and feed results into the updater queue."""
         if self._config.dry_run:
             return
-        running, addresses = self._transitions.get_running_tasks_for_poll()
+        with self._store.read_snapshot() as snap:
+            running, addresses = self._transitions.get_running_tasks_for_poll(snap)
         if not running:
             return
         poll_results = self._provider.poll_workers(running, addresses)
@@ -2296,7 +2312,8 @@ def _run_task_updater_loop(self, stop_event: threading.Event) -> None:
             if not requests or stop_event.is_set():
                 continue
             try:
-                results = self._transitions.apply_heartbeats_batch(requests)
+                with self._store.transaction() as cur:
+                    results = self._transitions.apply_heartbeats_batch(cur, requests)
                 all_tasks_to_kill: set[JobName] = set()
                 all_task_kill_workers: dict[JobName, WorkerId] = {}
                 for result in results:

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

@@ -1161,19 +1161,22 @@ def launch_job(
                 if not is_job_finished(existing_job.state):
                     return controller_pb2.Controller.LaunchJobResponse(job_id=job_id.to_wire())
                 # Job finished, replace it (KEEP only preserves running jobs)
-                self._transitions.remove_finished_job(job_id)
+                with self._store.transaction() as cur:
+                    self._transitions.remove_finished_job(cur, job_id)
             elif policy == job_pb2.EXISTING_JOB_POLICY_RECREATE:
-                if not is_job_finished(existing_job.state):
-                    self._transitions.cancel_job(job_id, "Replaced by new submission")
-                self._transitions.remove_finished_job(job_id)
+                with self._store.transaction() as cur:
+                    if not is_job_finished(existing_job.state):
+                        self._transitions.cancel_job(cur, job_id, "Replaced by new submission")
+                    self._transitions.remove_finished_job(cur, job_id)
             elif is_job_finished(existing_job.state):
                 # Default/UNSPECIFIED: replace finished jobs
                 logger.info(
                     "Replacing finished job %s (state=%s) with new submission",
                     job_id,
                     job_pb2.JobState.Name(existing_job.state),
                 )
-                self._transitions.remove_finished_job(job_id)
+                with self._store.transaction() as cur:
+                    self._transitions.remove_finished_job(cur, job_id)
             else:
                 raise ConnectError(Code.ALREADY_EXISTS, f"Job {job_id} already exists and is still running")
 
@@ -1228,7 +1231,8 @@ def launch_job(
                     f"Job {job_id} is unschedulable: {error} (constraints: {constraints})",
                 )
 
-        self._transitions.submit_job(job_id, request, Timestamp.now())
+        with self._store.transaction() as cur:
+            self._transitions.submit_job(cur, job_id, request, Timestamp.now())
         self._controller.wake()
 
         with self._db.read_snapshot() as q:
@@ -1380,7 +1384,8 @@ def terminate_job(
         self._authorize_job_owner(job_id)
         # cancel_job uses a recursive CTE to walk the full subtree in a single
         # transaction, so there is no need to recurse manually.
-        result = self._transitions.cancel_job(job_id, reason="Terminated by user")
+        with self._store.transaction() as cur:
+            result = self._transitions.cancel_job(cur, job_id, reason="Terminated by user")
         if result.tasks_to_kill:
             self._controller.kill_tasks_on_workers(result.tasks_to_kill, result.task_kill_workers)
         return job_pb2.Empty()
@@ -1625,14 +1630,16 @@ def register(
             )
         worker_id = WorkerId(request.worker_id)
 
-        self._transitions.register_or_refresh_worker(
-            worker_id=worker_id,
-            address=request.address,
-            metadata=request.metadata,
-            ts=Timestamp.now(),
-            slice_id=request.slice_id,
-            scale_group=request.scale_group,
-        )
+        with self._store.transaction() as cur:
+            self._transitions.register_or_refresh_worker(
+                cur,
+                worker_id=worker_id,
+                address=request.address,
+                metadata=request.metadata,
+                ts=Timestamp.now(),
+                slice_id=request.slice_id,
+                scale_group=request.scale_group,
+            )
 
         logger.info("Worker registered: %s at %s", worker_id, request.address)
         return controller_pb2.Controller.RegisterResponse(
@@ -1711,7 +1718,9 @@ def register_endpoint(
             registered_at=Timestamp.now(),
         )
 
-        if not self._transitions.add_endpoint(endpoint):
+        with self._store.transaction() as cur:
+            added = self._transitions.add_endpoint(cur, endpoint)
+        if not added:
             raise ConnectError(
                 Code.FAILED_PRECONDITION,
                 f"Task {request.task_id} is already terminal; endpoint not registered",
@@ -1725,7 +1734,8 @@ def unregister_endpoint(
         ctx: Any,
     ) -> job_pb2.Empty:
         """Unregister a service endpoint. Idempotent."""
-        self._transitions.remove_endpoint(request.endpoint_id)
+        with self._store.transaction() as cur:
+            self._transitions.remove_endpoint(cur, request.endpoint_id)
         return job_pb2.Empty()
 
     def list_endpoints(
@@ -2683,12 +2693,14 @@ def update_task_status(
         """
         updates = task_updates_from_proto(request.updates)
         if updates:
-            self._transitions.apply_task_updates(
-                HeartbeatApplyRequest(
-                    worker_id=WorkerId(request.worker_id),
-                    worker_resource_snapshot=None,
-                    updates=updates,
+            with self._store.transaction() as cur:
+                self._transitions.apply_task_updates(
+                    cur,
+                    HeartbeatApplyRequest(
+                        worker_id=WorkerId(request.worker_id),
+                        worker_resource_snapshot=None,
+                        updates=updates,
+                    ),
                 )
-            )
             self._controller.wake()
         return controller_pb2.Controller.UpdateTaskStatusResponse()

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

@@ -2070,3 +2070,6 @@ def transaction(self):
 
     def read_snapshot(self):
         return self._db.read_snapshot()
+
+    def optimize(self) -> None:
+        self._db.optimize()