Skip to main content

mz_adapter/coord/
cluster_controller.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Driver and glue for the [`mz_cluster_controller`] reconciler.
11//!
12//! The controller crate is pure: it knows nothing about the Coordinator. This
13//! module is the half of the [`ClusterControllerCtx`] boundary that does: it runs
14//! the controller as a **separate task** and implements the ctx by marshaling
15//! each pull/apply to the Coordinator over the internal command channel, because
16//! the catalog and the live compute/storage signals are reachable only from the
17//! coordinator loop. The two whole-tick reads are batched; the per-cluster live
18//! signals are pulled on demand, so a tick's round-trips scale with the number of
19//! managed clusters that need a live signal, not with a constant.
20//!
21//! Everything here is gated by [`ENABLE_CLUSTER_CONTROLLER`] (default on). With
22//! the gate off the task does not tick, so the legacy scheduling and graceful
23//! paths remain the sole writers of the replica set. With the gate on the
24//! controller owns the *user* managed-cluster replica set; the legacy entry
25//! points no-op. (System/builtin clusters are excluded here. Their config-implied
26//! replicas are materialized by `reconcile_builtin_cluster_replicas` at catalog
27//! open, which derives the same target from the same config.)
28
29use std::collections::BTreeSet;
30use std::sync::Arc;
31use std::time::Duration;
32
33use mz_adapter_types::dyncfgs::{CLUSTER_CONTROLLER_TICK_INTERVAL, ENABLE_CLUSTER_CONTROLLER};
34use mz_catalog::memory::objects::{ClusterConfig, ClusterVariant};
35use mz_cluster_controller::ClusterController;
36use mz_cluster_controller::ctx::{
37    ApplyOutcome, AvailabilityZones, ClusterControllerCtx, ClusterState, CreateReason, Decision,
38    ExpectedClusterState, ObservedReplica, OnTimeout, ReconfigurationRecord, ReconfigurationStatus,
39    ReconfigurationTarget, RefreshMvInfo, RefreshWindowInputs, ReplicaShape, StateWrite,
40};
41use mz_compute_types::config::ComputeReplicaConfig;
42use mz_controller::clusters::ClusterStatus;
43use mz_controller_types::{ClusterId, ReplicaId};
44use mz_ore::task::spawn;
45use mz_repr::Timestamp;
46use tokio::sync::{mpsc, oneshot};
47use tracing::{debug, warn};
48
49use crate::catalog::{DropObjectInfo, Op, ReplicaCreateDropReason};
50use crate::coord::{ClusterReplicaStatuses, Coordinator, Message};
51use crate::error::AdapterError;
52
53/// A request the controller task marshals to the Coordinator to satisfy one
54/// [`ClusterControllerCtx`] call. Each variant carries a oneshot for the reply.
55///
56/// `ManagedClusterIds` and `ClusterStates` are the per-tick batched reads. The
57/// `ClusterStates` reply also carries `now`. `HydratedReplicas` is a
58/// per-cluster live signal a strategy pulls on demand.
59#[derive(Debug)]
60pub enum ClusterControllerRequest {
61    /// The ids of all *user* managed clusters the controller owns this tick.
62    /// System/builtin clusters are excluded. `reconcile_builtin_cluster_replicas`
63    /// has already materialized their config-implied replicas at catalog open.
64    ManagedClusterIds { tx: oneshot::Sender<Vec<ClusterId>> },
65    /// A consistent durable view of the given clusters and their replicas, plus
66    /// the current time.
67    ClusterStates {
68        clusters: Vec<ClusterId>,
69        tx: oneshot::Sender<(Vec<ClusterState>, Timestamp)>,
70    },
71    /// Of `replicas` on `cluster`, which are online and have all current
72    /// collections hydrated.
73    HydratedReplicas {
74        cluster_id: ClusterId,
75        replicas: Vec<ReplicaId>,
76        tx: oneshot::Sender<BTreeSet<ReplicaId>>,
77    },
78    /// Whether the cluster has any hydratable (dataflow-backed) objects bound to
79    /// it.
80    HasHydratableObjects {
81        cluster_id: ClusterId,
82        tx: oneshot::Sender<bool>,
83    },
84    /// The refresh-window live signals for one scheduled cluster (read ts,
85    /// compaction estimate, bound REFRESH MVs). `None` for a cluster that is not
86    /// scheduled `ON REFRESH`.
87    RefreshWindowInputs {
88        cluster_id: ClusterId,
89        tx: oneshot::Sender<Option<RefreshWindowInputs>>,
90    },
91    /// Apply a tick's batch of decisions under their compare-and-append guards.
92    Apply {
93        decisions: Vec<Decision>,
94        tx: oneshot::Sender<ApplyOutcome>,
95    },
96    /// The current configured reconcile cadence. Read once per tick so a runtime
97    /// change to `cluster_controller_tick_interval` takes effect without a
98    /// restart.
99    TickInterval { tx: oneshot::Sender<Duration> },
100}
101
102struct ReplicaHydrationCheck {
103    replica_id: ReplicaId,
104    compute_hydrated: oneshot::Receiver<bool>,
105}
106
107/// The controller-task side of the boundary: a [`ClusterControllerCtx`] that
108/// marshals every call to the Coordinator over `internal_cmd_tx`.
109struct CoordCtx {
110    internal_cmd_tx: mpsc::UnboundedSender<Message>,
111    /// Latched `now` from the most recent batched read, returned by
112    /// [`ClusterControllerCtx::now`] so a strategy and the kernel see a single
113    /// consistent time per phase.
114    now: Timestamp,
115}
116
117impl CoordCtx {
118    /// Send a request and await its reply. Returns `None` if the Coordinator has
119    /// gone away (shutdown), which the caller treats as "nothing to do".
120    async fn request<T>(
121        &self,
122        make: impl FnOnce(oneshot::Sender<T>) -> ClusterControllerRequest,
123    ) -> Option<T> {
124        let (tx, rx) = oneshot::channel();
125        if self
126            .internal_cmd_tx
127            .send(Message::ClusterControllerRequest(make(tx)))
128            .is_err()
129        {
130            return None;
131        }
132        rx.await.ok()
133    }
134}
135
136#[async_trait::async_trait]
137impl ClusterControllerCtx for CoordCtx {
138    fn now(&self) -> Timestamp {
139        self.now
140    }
141
142    async fn managed_cluster_ids(&mut self) -> Vec<ClusterId> {
143        self.request(|tx| ClusterControllerRequest::ManagedClusterIds { tx })
144            .await
145            .unwrap_or_default()
146    }
147
148    async fn cluster_states(&mut self, clusters: &[ClusterId]) -> Vec<ClusterState> {
149        let clusters = clusters.to_vec();
150        match self
151            .request(|tx| ClusterControllerRequest::ClusterStates { clusters, tx })
152            .await
153        {
154            Some((states, now)) => {
155                self.now = now;
156                states
157            }
158            None => Vec::new(),
159        }
160    }
161
162    async fn hydrated_replicas(
163        &mut self,
164        cluster_id: ClusterId,
165        replicas: &[ReplicaId],
166    ) -> BTreeSet<ReplicaId> {
167        let replicas = replicas.to_vec();
168        self.request(|tx| ClusterControllerRequest::HydratedReplicas {
169            cluster_id,
170            replicas,
171            tx,
172        })
173        .await
174        .unwrap_or_default()
175    }
176
177    async fn has_hydratable_objects(&mut self, cluster_id: ClusterId) -> bool {
178        self.request(|tx| ClusterControllerRequest::HasHydratableObjects { cluster_id, tx })
179            .await
180            // A lost reply means shutdown; "no objects" arms nothing, which is
181            // the safe answer.
182            .unwrap_or(false)
183    }
184
185    async fn refresh_window_inputs(
186        &mut self,
187        cluster_id: ClusterId,
188    ) -> Option<RefreshWindowInputs> {
189        self.request(|tx| ClusterControllerRequest::RefreshWindowInputs { cluster_id, tx })
190            .await
191            .flatten()
192    }
193
194    async fn apply(&mut self, decisions: Vec<Decision>) -> ApplyOutcome {
195        self.request(|tx| ClusterControllerRequest::Apply { decisions, tx })
196            .await
197            // A lost reply means shutdown; treat as rejected so we make no
198            // further claims about the catalog state.
199            .unwrap_or(ApplyOutcome::Rejected)
200    }
201}
202
203impl Coordinator {
204    /// Spawn the cluster controller task.
205    ///
206    /// The task ticks at [`CLUSTER_CONTROLLER_TICK_INTERVAL`] and reconciles when
207    /// [`ENABLE_CLUSTER_CONTROLLER`] is on; while the gate is off it ticks but
208    /// each tick is an early no-op. Both the gate and the interval are re-read
209    /// each tick (the interval via a [`ClusterControllerRequest::TickInterval`]
210    /// round-trip), so a runtime change to either takes effect without a restart.
211    /// It owns the controller and a [`CoordCtx`] that marshals back to this
212    /// Coordinator.
213    ///
214    /// The interval is the fallback cadence: `reconcile_now` cuts the
215    /// sleep short after a catalog transaction changes durable cluster state.
216    /// The notification only wakes the task. The tick still pulls fresh state
217    /// through the coordinator loop, and the controller's own applies wake it
218    /// again at the cost of one no-op tick.
219    pub(crate) fn spawn_cluster_controller_task(&self) {
220        let internal_cmd_tx = self.internal_cmd_tx.clone();
221        let reconcile_now = Arc::clone(&self.reconcile_now);
222        // A shared handle: dyncfg updates land in the same underlying values, so
223        // the controller task sees flag flips without any push.
224        let dyncfgs = self.catalog().system_config().dyncfgs().clone();
225
226        spawn(|| "cluster_controller", async move {
227            let controller = ClusterController::new(dyncfgs);
228            let mut ctx = CoordCtx {
229                internal_cmd_tx,
230                now: Timestamp::MIN,
231            };
232
233            loop {
234                // Re-read the cadence each tick so a runtime change takes effect.
235                // A lost reply means the Coordinator is gone; stop ticking.
236                let Some(interval) = ctx
237                    .request(|tx| ClusterControllerRequest::TickInterval { tx })
238                    .await
239                else {
240                    break;
241                };
242                tokio::select! {
243                    _ = tokio::time::sleep(interval.max(Duration::from_millis(1))) => {}
244                    _ = reconcile_now.notified() => {}
245                }
246
247                if ctx.internal_cmd_tx.is_closed() {
248                    // Coordinator gone; stop ticking.
249                    break;
250                }
251                controller.reconcile(&mut ctx).await;
252            }
253        });
254    }
255
256    /// Handle one [`ClusterControllerRequest`] on the coordinator loop.
257    ///
258    /// The controller is inactive when the gate is off, or while the deployment
259    /// is in read-only mode (a 0dt upgrade, where it must not write the catalog).
260    /// When inactive, reads report no managed clusters (so the controller finds
261    /// nothing to reconcile) and applies are rejected: the task still wakes each
262    /// tick and sends one `ManagedClusterIds` request, but that request
263    /// early-returns here and no catalog state is read or written, so the legacy
264    /// paths remain the sole writers of the replica set. The task keeps ticking,
265    /// so the controller reactivates on its own once the deployment promotes out
266    /// of read-only mode.
267    #[mz_ore::instrument(level = "debug")]
268    pub(crate) async fn handle_cluster_controller_request(
269        &mut self,
270        request: ClusterControllerRequest,
271    ) {
272        let active = ENABLE_CLUSTER_CONTROLLER.get(self.catalog().system_config().dyncfgs())
273            && !self.controller.read_only();
274
275        match request {
276            ClusterControllerRequest::ManagedClusterIds { tx } => {
277                let ids = if active {
278                    self.catalog()
279                        .clusters()
280                        // Only *user* managed clusters. System/builtin clusters
281                        // (mz_system, mz_catalog_server, …) are also managed, and
282                        // `reconcile_builtin_cluster_replicas` materializes their
283                        // config-implied replica set at catalog open, so those
284                        // replicas exist before the controller could ever tick.
285                        // Both derive the target from the cluster's
286                        // `replication_factor`, so extending ownership here would
287                        // converge rather than conflict.
288                        .filter(|c| c.is_managed() && c.id.is_user())
289                        .map(|c| c.id)
290                        .collect()
291                } else {
292                    Vec::new()
293                };
294                let _ = tx.send(ids);
295            }
296            ClusterControllerRequest::ClusterStates { clusters, tx } => {
297                let now = Timestamp::from(self.now());
298                // Only ever asked about clusters the controller is reconciling
299                // this tick, which the inactive `ManagedClusterIds` gate above
300                // makes empty, so no guard is needed here.
301                let states: Vec<_> = clusters
302                    .into_iter()
303                    .filter_map(|id| self.observe_cluster_state(id))
304                    .collect();
305                let _ = tx.send((states, now));
306            }
307            ClusterControllerRequest::HydratedReplicas {
308                cluster_id,
309                replicas,
310                tx,
311            } => {
312                let checks = self.start_hydration_checks(cluster_id, replicas);
313                // Start the controller calls on the coordinator loop, then wait
314                // for compute's replies off-loop. The compute check can wait on
315                // the compute instance task.
316                spawn(|| "cluster_controller_hydration_probe", async move {
317                    let mut hydrated = BTreeSet::new();
318                    for check in checks {
319                        if check.compute_hydrated.await.unwrap_or(false) {
320                            hydrated.insert(check.replica_id);
321                        }
322                    }
323                    let _ = tx.send(hydrated);
324                });
325            }
326            ClusterControllerRequest::HasHydratableObjects { cluster_id, tx } => {
327                let _ = tx.send(self.cluster_has_hydratable_objects(cluster_id));
328            }
329            ClusterControllerRequest::RefreshWindowInputs { cluster_id, tx } => {
330                // Gather the catalog- and storage-derived inputs on the loop,
331                // then complete the reply from a spawned task: the oracle
332                // read is a network round-trip (to the Postgres/CRDB-backed
333                // timestamp oracle) and must never run on the serial
334                // coordinator loop. The legacy `check_refresh_policy` makes
335                // the same split.
336                match self.refresh_window_catalog_inputs(cluster_id) {
337                    None => {
338                        let _ = tx.send(None);
339                    }
340                    Some((compaction_estimate, refresh_mvs)) => {
341                        let oracle = self.get_local_timestamp_oracle();
342                        // NOTE: this is one oracle read per scheduled cluster
343                        // per tick, and the controller awaits each pull before
344                        // the next, so the reads are sequential and the
345                        // batching oracle cannot coalesce them. Fine at the
346                        // tick cadence for realistic scheduled-cluster counts.
347                        // TODO: hoist to one read per tick if that stops
348                        // holding.
349                        spawn(|| "cluster_controller_refresh_window_read_ts", async move {
350                            let read_ts = oracle.read_ts().await;
351                            let _ = tx.send(Some(RefreshWindowInputs {
352                                read_ts,
353                                compaction_estimate,
354                                refresh_mvs,
355                            }));
356                        });
357                    }
358                }
359            }
360            ClusterControllerRequest::Apply { decisions, tx } => {
361                let outcome = if active {
362                    self.apply_cluster_decisions(decisions).await
363                } else {
364                    ApplyOutcome::Rejected
365                };
366                let _ = tx.send(outcome);
367            }
368            ClusterControllerRequest::TickInterval { tx } => {
369                let interval =
370                    CLUSTER_CONTROLLER_TICK_INTERVAL.get(self.catalog().system_config().dyncfgs());
371                let _ = tx.send(interval);
372            }
373        }
374    }
375
376    /// Build the controller's view of one managed cluster from the catalog.
377    /// Returns `None` for a missing or unmanaged cluster.
378    fn observe_cluster_state(&self, cluster_id: ClusterId) -> Option<ClusterState> {
379        let cluster = self.catalog().try_get_cluster(cluster_id)?;
380        let ClusterVariant::Managed(managed) = &cluster.config.variant else {
381            return None;
382        };
383        // The witness fields come from the same projection the compare-and-append
384        // check uses, so the state a decision is derived from and the state the
385        // apply path checks against cannot drift.
386        let expected = crate::catalog::cluster_state::project_expected(managed);
387
388        // All replicas, with the raw traits the controller's ownership test
389        // (`ObservedReplica::owned_shape`) classifies on.
390        let replicas = cluster
391            .replicas()
392            .map(|replica| ObservedReplica {
393                replica_id: replica.replica_id,
394                name: replica.name.clone(),
395                shape: replica_shape(&replica.config),
396                internal: replica.config.location.internal(),
397                billed_as: replica.config.location.billed_as().is_some(),
398                pending: replica.config.location.pending(),
399            })
400            .collect();
401
402        Some(ClusterState {
403            cluster_id,
404            size: expected.size,
405            replication_factor: expected.replication_factor,
406            availability_zones: expected.availability_zones.0,
407            logging: expected.logging,
408            arrangement_compression: expected.arrangement_compression,
409            schedule: expected.schedule,
410            auto_scaling_policy: expected.auto_scaling_policy,
411            reconfiguration: expected.reconfiguration,
412            burst: expected.burst,
413            replicas,
414        })
415    }
416
417    /// Whether the cluster has any hydratable objects bound to it, backing the
418    /// controller's [`ClusterControllerCtx::has_hydratable_objects`] pull (see
419    /// the trait method for the approximation contract and why mismatches with
420    /// the hydration check are self-healing).
421    fn cluster_has_hydratable_objects(&self, cluster_id: ClusterId) -> bool {
422        let Some(cluster) = self.catalog().try_get_cluster(cluster_id) else {
423            return false;
424        };
425        cluster
426            .bound_objects
427            .iter()
428            .any(|id| self.catalog().get_entry(id).item().is_hydratable())
429    }
430
431    /// Starts per-replica hydration checks for `cluster_id`.
432    ///
433    /// Returns only checks for replicas whose processes are all online, that
434    /// are already storage-hydrated, and that are known to the compute
435    /// controller. The compute receiver completes off the coordinator loop.
436    fn start_hydration_checks(
437        &self,
438        cluster_id: ClusterId,
439        replicas: Vec<ReplicaId>,
440    ) -> Vec<ReplicaHydrationCheck> {
441        use mz_catalog::memory::objects::CatalogItem;
442
443        // Materialized views pinned to a replica (via `IN CLUSTER ... REPLICA`)
444        // are only ever installed on that replica, so any other replica can
445        // never report them hydrated. Collect the cluster's pinned MVs once,
446        // then exclude the ones pinned elsewhere from each replica's hydration
447        // check. Otherwise a graceful reconfiguration's cut-over to a fresh
448        // replica set would wait forever for the new replicas to hydrate an MV
449        // bound to a replica being replaced, then roll back at the deadline
450        // (leaving the old replica, and the targeted MV, in place). Indexes
451        // cannot be replica-pinned, so MVs are the only case.
452        let pinned_mvs: Vec<(ReplicaId, mz_repr::GlobalId)> = self
453            .catalog()
454            .try_get_cluster(cluster_id)
455            .into_iter()
456            .flat_map(|cluster| cluster.bound_objects.iter())
457            .filter_map(|id| match self.catalog().get_entry(id).item() {
458                CatalogItem::MaterializedView(mv) => mv
459                    .target_replica
460                    .map(|target| (target, mv.global_id_writes())),
461                _ => None,
462            })
463            .collect();
464
465        let mut checks = Vec::new();
466        for replica_id in replicas {
467            // Skip replicas that are not online. We wait for a replica to be
468            // online even when it has no objects that need hydration on it
469            // (e.g. a single-replica source).
470            let status = self
471                .cluster_replica_statuses
472                .try_get_cluster_replica_statuses(cluster_id, replica_id)
473                .map(ClusterReplicaStatuses::cluster_replica_status);
474            if !matches!(status, Some(ClusterStatus::Online)) {
475                continue;
476            }
477            let exclude: BTreeSet<mz_repr::GlobalId> = pinned_mvs
478                .iter()
479                .filter(|(target, _)| *target != replica_id)
480                .map(|(_, id)| *id)
481                .collect();
482            let compute_fut = match self.controller.compute.collections_hydrated_for_replicas(
483                cluster_id,
484                vec![replica_id],
485                exclude.clone(),
486            ) {
487                Ok(fut) => fut,
488                // The replica is not known to the compute controller. Treat it
489                // as not hydrated.
490                Err(_) => continue,
491            };
492            let storage_hydrated = match self.controller.storage.collections_hydrated_on_replicas(
493                Some(vec![replica_id]),
494                &cluster_id,
495                &exclude,
496            ) {
497                Ok(hydrated) => hydrated,
498                Err(_) => continue,
499            };
500            if storage_hydrated {
501                checks.push(ReplicaHydrationCheck {
502                    replica_id,
503                    compute_hydrated: compute_fut,
504                });
505            }
506        }
507        checks
508    }
509
510    /// The catalog- and storage-derived refresh-window signals for one scheduled
511    /// cluster (the system compaction estimate and each bound REFRESH
512    /// materialized view's storage write frontier and refresh schedule), or
513    /// `None` if the cluster is missing, unmanaged, or not scheduled `ON
514    /// REFRESH`. These are the same signals the legacy `check_refresh_policy`
515    /// reads.
516    ///
517    /// The oracle read timestamp completing [`RefreshWindowInputs`] is
518    /// deliberately not fetched here: this runs on the coordinator loop, and
519    /// the oracle read is a network round-trip the request handler performs on
520    /// a spawned task instead.
521    ///
522    /// The MV write frontier is carried through with full fidelity as the
523    /// `Antichain` the storage controller reports, matching the legacy refresh
524    /// policy; the on-refresh strategy compares against it directly.
525    fn refresh_window_catalog_inputs(
526        &self,
527        cluster_id: ClusterId,
528    ) -> Option<(Duration, Vec<RefreshMvInfo>)> {
529        use mz_catalog::memory::objects::CatalogItem;
530
531        let cluster = self.catalog().try_get_cluster(cluster_id)?;
532        let ClusterVariant::Managed(managed) = &cluster.config.variant else {
533            return None;
534        };
535        if !matches!(
536            managed.schedule,
537            mz_sql::plan::ClusterSchedule::Refresh { .. }
538        ) {
539            return None;
540        }
541
542        let refresh_mvs = cluster
543            .bound_objects
544            .iter()
545            .filter_map(|id| {
546                let CatalogItem::MaterializedView(mv) = self.catalog().get_entry(id).item() else {
547                    return None;
548                };
549                let refresh_schedule = mv.refresh_schedule.clone()?;
550                // The storage controller knows about every MV in the catalog. The
551                // write frontier is passed through with full fidelity as the
552                // `Antichain` reported here.
553                let (_since, write_frontier) = self
554                    .controller
555                    .storage
556                    .collection_frontiers(mv.global_id_writes())
557                    .expect("storage controller knows about catalog MVs");
558                Some(RefreshMvInfo {
559                    id: mv.global_id_writes(),
560                    write_frontier,
561                    refresh_schedule,
562                })
563            })
564            .collect();
565
566        let compaction_estimate = self
567            .catalog()
568            .system_config()
569            .cluster_refresh_mv_compaction_estimate();
570
571        Some((compaction_estimate, refresh_mvs))
572    }
573
574    /// Apply one batch of decisions under their compare-and-append guards.
575    ///
576    /// The kernel calls this once per tick phase: a phase-1 batch is all
577    /// `UpdateClusterState`, a phase-2 batch is all create/drop. Either batch may
578    /// in principle be mixed; this handles both. The work is staged across four
579    /// steps: collect the per-cluster guards, pre-allocate the ids the creates
580    /// need, build the mutation ops, then commit ops and guards in one
581    /// transaction (see [`Self::commit_with_checks`] for why the guard holds).
582    /// Any step that finds the batch incoherent rejects it, and the controller
583    /// recomputes next tick.
584    async fn apply_cluster_decisions(&mut self, decisions: Vec<Decision>) -> ApplyOutcome {
585        let checks = Self::partition_checks(&decisions);
586
587        // Pre-allocate replica ids before the apply transaction (each allocation
588        // is its own durable commit, so it cannot happen inside the transaction).
589        let Some(replica_ids) = self.allocate_replica_ids_for_creates(&decisions).await else {
590            return ApplyOutcome::Rejected;
591        };
592
593        let Some(mutations) = self.build_mutation_ops(decisions, replica_ids) else {
594            return ApplyOutcome::Rejected;
595        };
596        if mutations.is_empty() {
597            // Nothing to apply, so the checks guard nothing. Skip the transaction
598            // rather than commit a check-only batch, which would still cost a
599            // durable round-trip.
600            return ApplyOutcome::Applied;
601        }
602
603        self.commit_with_checks(checks, mutations).await
604    }
605
606    /// The compare-and-append guards for a decision batch: one
607    /// `(cluster_id, expected)` per distinct cluster, in first-seen order. All
608    /// of a cluster's decisions in a tick come from one snapshot, so they share
609    /// one `expected` witness and one guard covers them all.
610    fn partition_checks(decisions: &[Decision]) -> Vec<(ClusterId, ExpectedClusterState)> {
611        let mut checks: Vec<(ClusterId, ExpectedClusterState)> = Vec::new();
612        let mut seen_clusters = BTreeSet::new();
613        for decision in decisions {
614            let (cluster_id, expected) = match decision {
615                Decision::CreateReplica {
616                    cluster_id,
617                    expected,
618                    ..
619                }
620                | Decision::DropReplica {
621                    cluster_id,
622                    expected,
623                    ..
624                }
625                | Decision::UpdateClusterState {
626                    cluster_id,
627                    expected,
628                    ..
629                } => (*cluster_id, expected),
630            };
631            if seen_clusters.insert(cluster_id) {
632                checks.push((cluster_id, expected.clone()));
633            } else {
634                debug_assert!(
635                    checks
636                        .iter()
637                        .any(|(c, e)| *c == cluster_id && e == expected),
638                    "decisions for a cluster in one tick must share one expected witness",
639                );
640            }
641        }
642        checks
643    }
644
645    /// Pre-allocate one replica id per `CreateReplica` decision, in the order the
646    /// creates appear (which is the order [`Self::build_mutation_ops`] consumes
647    /// them). Returns `None` if any allocation fails, which rejects the batch.
648    ///
649    /// `Op::CreateClusterReplica` carries a pre-allocated id, so we allocate
650    /// out-of-band here, before the apply transaction. Each allocation commits
651    /// durably, so we take a fresh write ts per allocation: two commits must not
652    /// share a timestamp.
653    async fn allocate_replica_ids_for_creates(
654        &mut self,
655        decisions: &[Decision],
656    ) -> Option<Vec<ReplicaId>> {
657        let mut replica_ids = Vec::new();
658        for decision in decisions {
659            let Decision::CreateReplica { cluster_id, .. } = decision else {
660                continue;
661            };
662            let id_ts = self.get_catalog_write_ts().await;
663            let result = self
664                .catalog()
665                .allocate_replica_ids(*cluster_id, 1, id_ts)
666                .await;
667            match result {
668                Ok(ids) => {
669                    replica_ids.push(ids.into_iter().next().expect("allocated one replica id"))
670                }
671                Err(err) => {
672                    warn!(%cluster_id, "cluster controller could not allocate replica id: {err}");
673                    return None;
674                }
675            }
676        }
677        Some(replica_ids)
678    }
679
680    /// Turn a decision batch into the catalog mutation ops to transact, consuming
681    /// the `replica_ids` pre-allocated for the creates (one per `CreateReplica`,
682    /// in order). Returns `None` if a target cluster has vanished or gone
683    /// unmanaged, which makes the batch incoherent and rejects it.
684    fn build_mutation_ops(
685        &self,
686        decisions: Vec<Decision>,
687        replica_ids: Vec<ReplicaId>,
688    ) -> Option<Vec<Op>> {
689        let mut replica_ids = replica_ids.into_iter();
690        let mut mutations = Vec::new();
691        let mut drops = Vec::new();
692        for decision in decisions {
693            match decision {
694                Decision::UpdateClusterState {
695                    cluster_id, write, ..
696                } => match self.build_update_cluster_config_op(cluster_id, &write) {
697                    Some(op) => mutations.push(op),
698                    // The cluster vanished. The batch is no longer coherent.
699                    None => return None,
700                },
701                Decision::CreateReplica {
702                    cluster_id,
703                    name,
704                    shape,
705                    reason,
706                    ..
707                } => {
708                    let replica_id = replica_ids.next().expect("one pre-allocated id per create");
709                    let reason = audit_reason_for_create(reason);
710                    match self.build_create_replica_op(cluster_id, replica_id, name, &shape, reason)
711                    {
712                        Ok(Some(op)) => mutations.push(op),
713                        Ok(None) => return None,
714                        Err(err) => {
715                            warn!(%cluster_id, "cluster controller could not build replica create: {err}");
716                            return None;
717                        }
718                    }
719                }
720                Decision::DropReplica {
721                    cluster_id,
722                    replica_id,
723                    ..
724                } => {
725                    // The replica may have vanished since the decisions were
726                    // derived (a user DDL landed between the tick's read and
727                    // this apply). The in-transaction witness check would
728                    // reject such a stale batch, but resource-limit validation
729                    // runs before the transaction and panics on a missing
730                    // replica, so reject the batch here instead.
731                    if self
732                        .catalog()
733                        .try_get_cluster_replica(cluster_id, replica_id)
734                        .is_none()
735                    {
736                        return None;
737                    }
738                    drops.push(DropObjectInfo::ClusterReplica((
739                        cluster_id,
740                        replica_id,
741                        ReplicaCreateDropReason::Retired,
742                    )));
743                }
744            }
745        }
746        if !drops.is_empty() {
747            mutations.push(Op::DropObjects(drops));
748        }
749        Some(mutations)
750    }
751
752    /// Prepend the per-cluster compare-and-append `checks` to `mutations` and
753    /// transact them together.
754    ///
755    /// The checks run inside the transaction, before any mutation, so they cannot
756    /// be separated from the commit they guard. A cluster whose durable state has
757    /// diverged from what the decisions were derived from (e.g. a user `ALTER`
758    /// landed mid-tick) aborts the whole batch, so a stale create or drop can
759    /// never reshape the replica set against the config the `ALTER` has since
760    /// established (in particular, a stale drop cannot retire a replica the
761    /// `ALTER` has just made desired). On rejection nothing is applied.
762    async fn commit_with_checks(
763        &mut self,
764        checks: Vec<(ClusterId, ExpectedClusterState)>,
765        mutations: Vec<Op>,
766    ) -> ApplyOutcome {
767        let mut ops: Vec<Op> = checks
768            .into_iter()
769            .map(|(cluster_id, expected)| Op::CheckClusterState {
770                cluster_id,
771                expected,
772            })
773            .collect();
774        ops.extend(mutations);
775
776        match self.catalog_transact(None, ops).await {
777            Ok(()) => ApplyOutcome::Applied,
778            Err(AdapterError::ClusterStateChanged { .. }) => {
779                // A concurrent `ALTER` moved a cluster's durable state out from
780                // under the decisions. Expected, so the controller recomputes
781                // next tick.
782                ApplyOutcome::Rejected
783            }
784            Err(AdapterError::ReadOnly) => {
785                // The controller is quiesced while read-only (see
786                // `handle_cluster_controller_request`), so this is normally
787                // unreachable; if reached it's expected and not actionable, not
788                // a failure to surface.
789                debug!("cluster controller apply skipped in read-only mode");
790                ApplyOutcome::Rejected
791            }
792            Err(AdapterError::ResourceExhaustion { .. }) => {
793                // The batch cannot fit the resource budget. Report the fact and
794                // leave the reaction (what, if anything, to shed) to the kernel.
795                debug!("cluster controller apply exceeded the resource budget");
796                ApplyOutcome::ResourceExhausted
797            }
798            Err(err) => {
799                warn!("cluster controller apply failed: {err}");
800                ApplyOutcome::Rejected
801            }
802        }
803    }
804
805    /// Build an [`Op::UpdateClusterConfig`] that applies `write`'s deltas to the
806    /// cluster's current in-memory config, or `None` if the cluster is gone or
807    /// unmanaged. The write was guard-checked against the same state, so this is
808    /// the realized cut-over / record write.
809    fn build_update_cluster_config_op(
810        &self,
811        cluster_id: ClusterId,
812        write: &StateWrite,
813    ) -> Option<Op> {
814        let cluster = self.catalog().try_get_cluster(cluster_id)?;
815        let mut config = cluster.config.clone();
816        let ClusterConfig {
817            variant: ClusterVariant::Managed(managed),
818            ..
819        } = &mut config
820        else {
821            return None;
822        };
823        // Exhaustive destructure of the source (no `..`): a field added to
824        // `StateWrite` is a compile error here until it's overlaid onto the
825        // managed config. We cannot destructure `managed` itself. It carries
826        // fields the controller does not model (`workload_class`,
827        // `optimizer_feature_overrides`) that this overlay must leave untouched.
828        let StateWrite {
829            new_size,
830            new_replication_factor,
831            new_availability_zones,
832            new_logging,
833            new_arrangement_compression,
834            reconfiguration,
835            burst,
836        } = write;
837        if let Some(size) = new_size {
838            managed.size = size.clone();
839        }
840        if let Some(rf) = new_replication_factor {
841            managed.replication_factor = *rf;
842        }
843        if let Some(azs) = new_availability_zones {
844            managed.availability_zones = azs.clone();
845        }
846        if let Some(logging) = new_logging {
847            managed.logging = logging.clone();
848        }
849        if let Some(arrangement_compression) = new_arrangement_compression {
850            managed.arrangement_compression = *arrangement_compression;
851        }
852        if let Some(reconfiguration) = reconfiguration {
853            managed.reconfiguration = reconfiguration.record.as_ref().map(memory_reconfiguration);
854        }
855        if let Some(burst) = burst {
856            managed.burst = burst.record.as_ref().map(memory_burst);
857        }
858        // The audit intents travel with the write, declared by the strategy at
859        // the decision point. We pass them through untouched so the events are
860        // emitted in the same catalog transaction as the state they describe.
861        let reconfiguration_audit = write.reconfiguration.as_ref().and_then(|w| w.audit);
862        let burst_audit = write.burst.as_ref().and_then(|w| w.audit);
863        Some(Op::UpdateClusterConfig {
864            id: cluster_id,
865            name: cluster.name.clone(),
866            config,
867            reconfiguration_audit,
868            burst_audit,
869        })
870    }
871
872    /// Build an [`Op::CreateClusterReplica`] for a desired replica `shape` on
873    /// `cluster_id` with the pre-allocated `replica_id`, attributed to `reason`.
874    /// Returns `Ok(None)` if the cluster is gone or unmanaged.
875    fn build_create_replica_op(
876        &self,
877        cluster_id: ClusterId,
878        replica_id: ReplicaId,
879        name: String,
880        shape: &ReplicaShape,
881        reason: ReplicaCreateDropReason,
882    ) -> Result<Option<Op>, mz_catalog::memory::error::Error> {
883        let Some(cluster) = self.catalog().try_get_cluster(cluster_id) else {
884            return Ok(None);
885        };
886        if !cluster.is_managed() {
887            return Ok(None);
888        }
889        let owner_id = cluster.owner_id;
890
891        let location = mz_catalog::durable::ReplicaLocation::Managed {
892            // Concretized from the cluster config below; left empty here.
893            availability_zones: Vec::new(),
894            billed_as: None,
895            internal: false,
896            size: shape.size.clone(),
897            pending: false,
898        };
899        let azs: Option<&[String]> = if shape.availability_zones.0.is_empty() {
900            None
901        } else {
902            Some(&shape.availability_zones.0)
903        };
904        let location = self.catalog().concretize_replica_location(
905            location,
906            &self
907                .catalog()
908                .get_role_allowed_cluster_sizes(&Some(owner_id)),
909            azs,
910            false,
911        )?;
912
913        let config = mz_controller::clusters::ReplicaConfig {
914            location,
915            compute: ComputeReplicaConfig {
916                logging: shape.logging.clone(),
917                arrangement_compression: shape.arrangement_compression,
918            },
919        };
920
921        Ok(Some(Op::CreateClusterReplica {
922            cluster_id,
923            replica_id,
924            name,
925            config,
926            owner_id,
927            reason,
928        }))
929    }
930}
931
932/// Map a create decision's [`CreateReason`] to the audit reason carried on the
933/// create event. `Baseline` audits [`ReplicaCreateDropReason::Manual`], the tag
934/// for replicas the user's own cluster config calls for. The match is
935/// exhaustive, so a new `CreateReason` variant is a compile error here instead
936/// of a silent `Manual`.
937///
938/// Drops never come through here: a drop happens exactly when no strategy
939/// desires the replica, so it carries no attribution and is uniformly audited
940/// [`ReplicaCreateDropReason::Retired`].
941fn audit_reason_for_create(reason: CreateReason) -> ReplicaCreateDropReason {
942    match reason {
943        CreateReason::Baseline => ReplicaCreateDropReason::Manual,
944        CreateReason::GracefulReconfiguration => ReplicaCreateDropReason::GracefulReconfiguration,
945        CreateReason::HydrationBurst => ReplicaCreateDropReason::HydrationBurst,
946        CreateReason::OnRefresh(decision) => ReplicaCreateDropReason::OnRefresh(decision),
947    }
948}
949
950/// Map an in-memory replica config to a [`ReplicaShape`], or `None` for an
951/// unmanaged replica (which the controller does not own).
952fn replica_shape(config: &mz_controller::clusters::ReplicaConfig) -> Option<ReplicaShape> {
953    use mz_controller::clusters::ReplicaLocation;
954    let ReplicaLocation::Managed(managed) = &config.location else {
955        return None;
956    };
957    Some(ReplicaShape {
958        size: managed.size.clone(),
959        availability_zones: AvailabilityZones(managed.availability_zones.clone()),
960        logging: config.compute.logging.clone(),
961        arrangement_compression: config.compute.arrangement_compression,
962    })
963}
964
965fn on_timeout_from_controller(action: OnTimeout) -> mz_sql::plan::OnTimeoutAction {
966    match action {
967        OnTimeout::Commit => mz_sql::plan::OnTimeoutAction::Commit,
968        OnTimeout::Rollback => mz_sql::plan::OnTimeoutAction::Rollback,
969    }
970}
971
972fn memory_reconfiguration(
973    record: &ReconfigurationRecord,
974) -> mz_catalog::memory::objects::ReconfigurationState {
975    // Destructure the source (no `..`): a field added to the controller type is a
976    // compile error here until it's carried across. The target is the same.
977    let ReconfigurationRecord {
978        target,
979        deadline,
980        on_timeout,
981        status,
982    } = record;
983    let ReconfigurationTarget {
984        size,
985        replication_factor,
986        availability_zones,
987        logging,
988        arrangement_compression,
989    } = target;
990    mz_catalog::memory::objects::ReconfigurationState {
991        target: mz_catalog::memory::objects::ReconfigurationTarget {
992            size: size.clone(),
993            replication_factor: *replication_factor,
994            availability_zones: availability_zones.0.clone(),
995            logging: logging.clone(),
996            arrangement_compression: *arrangement_compression,
997        },
998        deadline: *deadline,
999        on_timeout: on_timeout_from_controller(*on_timeout),
1000        status: status_from_controller(*status),
1001    }
1002}
1003
1004fn status_from_controller(
1005    status: ReconfigurationStatus,
1006) -> mz_catalog::memory::objects::ReconfigurationStatus {
1007    match status {
1008        ReconfigurationStatus::InProgress => {
1009            mz_catalog::memory::objects::ReconfigurationStatus::InProgress
1010        }
1011        ReconfigurationStatus::Finalized => {
1012            mz_catalog::memory::objects::ReconfigurationStatus::Finalized
1013        }
1014        ReconfigurationStatus::TimedOut => {
1015            mz_catalog::memory::objects::ReconfigurationStatus::TimedOut
1016        }
1017        ReconfigurationStatus::Cancelled => {
1018            mz_catalog::memory::objects::ReconfigurationStatus::Cancelled
1019        }
1020        ReconfigurationStatus::ResourceExhausted => {
1021            mz_catalog::memory::objects::ReconfigurationStatus::ResourceExhausted
1022        }
1023    }
1024}
1025
1026fn memory_burst(
1027    record: &mz_cluster_controller::ctx::BurstRecord,
1028) -> mz_catalog::memory::objects::BurstState {
1029    // Destructure the source (no `..`): a field added to the controller type is a
1030    // compile error here until it's carried across.
1031    let mz_cluster_controller::ctx::BurstRecord {
1032        burst_size,
1033        linger_duration,
1034        steady_hydrated_at,
1035    } = record;
1036    mz_catalog::memory::objects::BurstState {
1037        burst_size: burst_size.clone(),
1038        linger_duration: *linger_duration,
1039        steady_hydrated_at: *steady_hydrated_at,
1040    }
1041}
1042
1043#[cfg(test)]
1044mod tests {
1045    use super::*;
1046
1047    #[mz_ore::test]
1048    fn test_audit_reason_for_create() {
1049        use ReplicaCreateDropReason as Reason;
1050        use mz_cluster_controller::ctx::RefreshWindowDecision;
1051        use mz_repr::GlobalId;
1052
1053        // Each variant maps to its own audit reason, with the baseline
1054        // auditing `Manual`.
1055        assert!(matches!(
1056            audit_reason_for_create(CreateReason::Baseline),
1057            Reason::Manual
1058        ));
1059        assert!(matches!(
1060            audit_reason_for_create(CreateReason::GracefulReconfiguration),
1061            Reason::GracefulReconfiguration
1062        ));
1063        assert!(matches!(
1064            audit_reason_for_create(CreateReason::HydrationBurst),
1065            Reason::HydrationBurst
1066        ));
1067
1068        // The on-refresh reason carries the create's window decision through
1069        // to the audit detail intact.
1070        let decision = RefreshWindowDecision {
1071            objects_needing_refresh: vec![GlobalId::User(1)],
1072            objects_needing_compaction: vec![GlobalId::User(2)],
1073            hydration_time_estimate: Duration::from_secs(7),
1074        };
1075        match audit_reason_for_create(CreateReason::OnRefresh(decision.clone())) {
1076            Reason::OnRefresh(carried) => assert_eq!(carried, decision),
1077            other => panic!("expected an on-refresh reason, got {other:?}"),
1078        }
1079    }
1080}