mz_cluster_controller/

lib.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! The cluster controller: the single decision-maker for the replica set of
//! every managed cluster.
//!
//! It is a **reconciler**. Each tick it reads desired cluster state and live
//! signals through the [`ClusterControllerCtx`] boundary, runs a set of pure
//! [`Strategy`]s, unions their desired contributions, diffs that against the
//! actual replica set, and emits the create/drop and durable-state-write
//! [`Decision`]s that close the gap. It holds no in-memory state: the source of
//! truth is always the catalog plus live signals, pulled fresh each tick.
//!
//! The crate is **pure**. It depends only on primitive id/shape types and the
//! [`ClusterControllerCtx`] trait, never on the adapter or catalog. That
//! boundary is what makes the controller testable against a fake
//! implementation and extractable later without touching controller code.
//!
//! A tick runs two phases per cluster, `update_state` then `desired_replicas`
//! (see [`ClusterController::reconcile`]). Every [`Decision`] carries the
//! durable state it was derived from, and the apply path transacts it only if
//! that state still holds (compare-and-append). So a create or drop derived
//! from a pre-`ALTER` snapshot can never reshape the replica set against the
//! config the `ALTER` has since established. Applies are per cluster, so one
//! cluster's rejection does not block the others, and commands name explicit
//! replicas, so re-emitting one across a lagging view or a restart is a no-op.
//!
//! [`ClusterControllerCtx`]: crate::ctx::ClusterControllerCtx

pub mod ctx;
pub mod strategy;

use std::collections::{BTreeMap, BTreeSet};

use mz_adapter_types::dyncfgs::{DEFAULT_HYDRATION_BURST_LINGER, ENABLE_HYDRATION_BURST};
use mz_controller_types::ClusterId;
use mz_dyncfg::ConfigSet;
use mz_ore::soft_panic_or_log;

use crate::ctx::{
    ApplyOutcome, BurstAudit, BurstFinishCause, BurstWrite, ClusterControllerCtx, ClusterState,
    CreateReason, Decision, ObservedReplica, ReconfigurationAudit, ReconfigurationRecord,
    ReconfigurationStatus, ReconfigurationWrite, RefreshWindowInputs, ReplicaShape, StateWrite,
};
use crate::strategy::{
    BaselineStrategy, ConfigSignals, DesiredReplica, GracefulReconfigurationStrategy,
    HydrationBurstStrategy, LiveSignals, OnRefreshStrategy, SignalRequest, Strategy,
};

/// The cluster controller. Holds the (stateless) set of strategies and drives a
/// reconcile tick against a [`ClusterControllerCtx`].
pub struct ClusterController {
    strategies: Vec<Box<dyn Strategy>>,
    /// The dyncfgs the config signals are latched from each tick. A shared
    /// handle, so a flipped flag takes effect on the next tick.
    dyncfgs: ConfigSet,
}

impl ClusterController {
    /// A controller with the full set of strategies. Each strategy's rustdoc
    /// describes when it engages.
    pub fn new(dyncfgs: ConfigSet) -> Self {
        Self {
            strategies: vec![
                Box::new(BaselineStrategy),
                Box::new(GracefulReconfigurationStrategy),
                Box::new(OnRefreshStrategy),
                Box::new(HydrationBurstStrategy),
            ],
            dyncfgs,
        }
    }

    /// The tick's config signals, latched from the dyncfgs so every strategy
    /// decides against one consistent config per tick.
    fn config_signals(&self) -> ConfigSignals {
        ConfigSignals {
            burst_enabled: ENABLE_HYDRATION_BURST.get(&self.dyncfgs),
            default_burst_linger: DEFAULT_HYDRATION_BURST_LINGER.get(&self.dyncfgs),
        }
    }

    /// Run one reconcile tick over every managed cluster the ctx reports.
    ///
    /// See the module docs for the two-phase structure. Both phases apply per
    /// cluster, so a compare-and-append rejection on one cluster never blocks
    /// progress on the others.
    pub async fn reconcile(&self, ctx: &mut dyn ClusterControllerCtx) {
        let cluster_ids = ctx.managed_cluster_ids().await;
        if cluster_ids.is_empty() {
            return;
        }

        // Phase 1: update_state. We merge every strategy's write for a cluster
        // into one compare-and-append, applied per cluster and independently of
        // other clusters. Two separate decisions live here.
        //
        // Per cluster, not one batch per tick: a write rejected because a
        // concurrent `ALTER` moved the cluster off its `expected` rejects only
        // that cluster and leaves the rest free to progress. One batched apply
        // would let a single mid-`ALTER` cluster sink the whole tick, the failure
        // mode at large cluster counts where some cluster is almost always
        // mid-`ALTER`.
        //
        // Merged across strategies, not one apply per strategy: every strategy
        // for a cluster shares the same start-of-tick `expected`, so applying
        // them one at a time would let the first write move the cluster off that
        // `expected` and reject all the rest, serializing a cluster's disjoint
        // writes one-per-tick. Merging lands them together under one guard. We
        // still rely on the compare-and-append, not the merge, for `ALTER`
        // safety, which is why the merged write carries the cluster's `expected`.
        // See `merge_state_writes` for the join and its conflict handling.
        let states = ctx.cluster_states(&cluster_ids).await;
        let config = self.config_signals();
        let signals = self.fetch_signals(ctx, &states, &config).await;
        let now = ctx.now();
        // Set when we issue any phase-1 apply, applied or rejected. Either way
        // the durable state may have moved (our write, or the concurrent `ALTER`
        // that rejected it), so phase 2 re-reads.
        let mut phase_1_wrote = false;
        // Clusters whose phase-1 write was rejected. We skip their phase 2 this
        // tick. Proceeding would be safe (we re-read below and every create/drop
        // is guard-checked), but a cluster that just lost a race is likely still
        // settling, so we let it recompute next tick instead of emitting work
        // that is probably about to go stale.
        let mut rejected = BTreeSet::new();
        for state in &states {
            let Some(signals) = signals.get(&state.cluster_id) else {
                continue;
            };
            let write = self.merge_state_writes(state, signals, &config, now);
            if write.is_empty() {
                continue;
            }
            phase_1_wrote = true;
            let decision = Decision::UpdateClusterState {
                cluster_id: state.cluster_id,
                expected: state.expected(),
                write,
            };
            // A phase-1 batch carries no creates, so it cannot exhaust the
            // resource budget. Treat any non-applied outcome as a rejection.
            if ctx.apply(vec![decision]).await != ApplyOutcome::Applied {
                rejected.insert(state.cluster_id);
            }
        }

        // Phase 2: desired_replicas. The barrier exists so that a cut-over a
        // phase-1 write performed is visible before we diff the replica set
        // against the realized config. We re-read (and re-enrich) only if phase 1
        // wrote. The first read is otherwise still current. A stale diff is
        // harmless: every create/drop carries its `expected` and is guard-rejected
        // if the durable state has since diverged.
        let (states, signals) = if phase_1_wrote {
            let states = ctx.cluster_states(&cluster_ids).await;
            let signals = self.fetch_signals(ctx, &states, &config).await;
            (states, signals)
        } else {
            (states, signals)
        };
        let now = ctx.now();
        for state in &states {
            if rejected.contains(&state.cluster_id) {
                continue;
            }
            let Some(signals) = signals.get(&state.cluster_id) else {
                continue;
            };
            let decisions = self.collect_replica_decisions(state, signals, &config, now);
            if decisions.is_empty() {
                continue;
            }
            // Per-cluster apply: a guard failure here is isolated to this cluster,
            // and benign anyway since every command names an explicit replica and
            // is reconciled away next tick. We do not retry within the tick.
            match ctx.apply(decisions).await {
                ApplyOutcome::Applied | ApplyOutcome::Rejected => {}
                ApplyOutcome::ResourceExhausted => {
                    // The batch exceeded the resource budget. Retrying cannot make
                    // the transient peak smaller, so shed the cluster's most
                    // expendable transient strategy and recompute next tick.
                    //
                    // The failed apply rolled back without changing durable state,
                    // so this tick's `expected` witness is still current, unless a
                    // concurrent user `ALTER` re-targeted the record, in which case
                    // the guard rejects the shed and that new reconfiguration is
                    // left to converge instead of being clobbered.
                    if let Some(shed) = Self::shed_decision(state) {
                        let _ = ctx.apply(vec![shed]).await;
                    }
                }
            }
        }
    }

    /// The decision that sheds this cluster's most expendable transient strategy
    /// after a resource-exhausted apply, or `None` if nothing sheddable is
    /// active.
    ///
    /// The strategy to shed is chosen by presence, ranked by expendability, not
    /// by which create failed: validation is aggregate, and the strategy worth
    /// giving up may be one whose replicas already materialized rather than one
    /// in the failed batch. A graceful reconfiguration is a discretionary user
    /// change that fails cleanly. The failure is audited, the wait-shim reports
    /// insufficient resources, and the cluster keeps running at its realized
    /// shape. The baseline is never shed because it is the committed floor. A
    /// hydration burst remains armed because no durable state records that the
    /// unchanged policy should suppress it.
    ///
    /// Without an active graceful reconfiguration there is nothing to shed. The
    /// next tick retries the desired replica set.
    fn shed_decision(state: &ClusterState) -> Option<Decision> {
        let record = state.reconfiguration.as_ref()?;
        if !record.is_in_progress() {
            return None;
        }
        Some(Decision::UpdateClusterState {
            cluster_id: state.cluster_id,
            expected: state.expected(),
            write: StateWrite {
                reconfiguration: Some(ReconfigurationWrite {
                    record: Some(ReconfigurationRecord {
                        status: ReconfigurationStatus::ResourceExhausted,
                        ..record.clone()
                    }),
                    audit: Some(ReconfigurationAudit::ResourceExhausted),
                }),
                ..Default::default()
            },
        })
    }

    /// Merge every strategy's [`Strategy::update_state`] for one cluster into the
    /// single [`StateWrite`] the tick applies under one compare-and-append.
    ///
    /// The merge is a per-field join, independent of the order strategies run
    /// in, before the cutover adjustment below: a field set by exactly one
    /// strategy is taken as-is, a field no strategy sets is left unchanged, and
    /// a field set to the same value by several is that value.
    ///
    /// Two strategies setting one field to *different* values is a conflict.
    /// The strategies keep every field single-writer at any given moment:
    /// most fields are owned by exactly one strategy outright, and
    /// `new_replication_factor`, which both the graceful cut-over and the
    /// on-refresh normalization write, is time-shared (on-refresh skips its
    /// normalization while a reconfiguration record is in progress). So by
    /// design a conflict cannot happen and the merge is really a disjoint
    /// union. We treat a conflict as an invariant violation rather than a
    /// condition to resolve: there is no safety-meaningful winner to pick for
    /// a contended `size` or record, so we trip [`soft_panic_or_log!`] (a
    /// panic under test/CI soft assertions, a logged error in production) and
    /// leave the field unchanged, the only outcome that cannot make things
    /// worse. A persistent conflict then freezes that field and keeps tripping
    /// the alarm, which is the point: surface the design bug loudly instead of
    /// silently picking an arbitrary value.
    fn merge_state_writes(
        &self,
        state: &ClusterState,
        signals: &LiveSignals,
        config: &ConfigSignals,
        now: mz_repr::Timestamp,
    ) -> StateWrite {
        let writes: Vec<StateWrite> = self
            .strategies
            .iter()
            .map(|strategy| strategy.update_state(state, signals, config, now))
            .filter(|write| !write.is_empty())
            .collect();

        let mut conflicts: Vec<&'static str> = Vec::new();
        // Exhaustive construction (every field named, no `..`): a field added to
        // `StateWrite` is a compile error here until its join is spelled out.
        let mut merged = StateWrite {
            new_size: join(
                "size",
                writes.iter().map(|w| w.new_size.clone()),
                &mut conflicts,
            ),
            new_replication_factor: join(
                "replication_factor",
                writes.iter().map(|w| w.new_replication_factor),
                &mut conflicts,
            ),
            new_availability_zones: join(
                "availability_zones",
                writes.iter().map(|w| w.new_availability_zones.clone()),
                &mut conflicts,
            ),
            new_logging: join(
                "logging",
                writes.iter().map(|w| w.new_logging.clone()),
                &mut conflicts,
            ),
            new_arrangement_compression: join(
                "arrangement_compression",
                writes.iter().map(|w| w.new_arrangement_compression),
                &mut conflicts,
            ),
            reconfiguration: join(
                "reconfiguration",
                writes.iter().map(|w| w.reconfiguration.clone()),
                &mut conflicts,
            ),
            burst: join(
                "burst",
                writes.iter().map(|w| w.burst.clone()),
                &mut conflicts,
            ),
        };

        if !conflicts.is_empty() {
            soft_panic_or_log!(
                "cluster {:?}: strategies produced conflicting state writes for \
                 field(s) {}; leaving those fields unchanged. Strategies must own \
                 disjoint `StateWrite` fields.",
                state.cluster_id,
                conflicts.join(", "),
            );
        }

        // A graceful cutover replaces the steady set used to establish burst
        // readiness. Lag allowances are not transitive, so the new steady set
        // must establish its own readiness directly against the burst helpers.
        if merged.new_size.is_some() {
            if let Some(record) = &state.burst {
                let explicit_cleanup = merged.burst.as_ref().is_some_and(|write| {
                    write.audit
                        == Some(BurstAudit::Finished {
                            cause: BurstFinishCause::NoLongerWarranted,
                        })
                });
                if !explicit_cleanup {
                    merged.burst = Some(BurstWrite {
                        record: Some(crate::ctx::BurstRecord {
                            steady_ready_at: None,
                            ..record.clone()
                        }),
                        audit: None,
                    });
                }
            }
        }

        merged
    }

    /// Fetch the live signals the strategies declared they need for `states`.
    ///
    /// Each strategy names its needs as a pure function of the durable state
    /// and the tick's config signals ([`Strategy::signal_request`]), so the
    /// kernel stays ignorant of when a strategy engages. Signals are fetched
    /// only where requested: a steady cluster is never probed, keeping the ctx
    /// seam pay-for-what-you-use. Refresh-window inputs are fetched as one batch
    /// so every scheduled cluster shares one oracle read per phase. The returned
    /// map omits a state when one of its required inputs was unavailable, which
    /// causes the reconciliation phase to skip that cluster.
    async fn fetch_signals(
        &self,
        ctx: &mut dyn ClusterControllerCtx,
        states: &[ClusterState],
        config: &ConfigSignals,
    ) -> BTreeMap<ClusterId, LiveSignals> {
        let mut signals = BTreeMap::new();
        let mut refresh_window_clusters = Vec::new();
        for state in states {
            let request = self
                .strategies
                .iter()
                .fold(SignalRequest::default(), |acc, strategy| {
                    acc.union(strategy.signal_request(state, config))
                });
            let mut live = LiveSignals::default();
            if request.hydratable_objects {
                live.has_hydratable_objects = ctx.has_hydratable_objects(state.cluster_id).await;
            }
            if request.hydration || !request.readiness.is_empty() {
                let replica_ids: Vec<_> = state
                    .replicas
                    .iter()
                    .filter(|r| r.owned_shape().is_some())
                    .map(|r| r.replica_id)
                    .collect();
                if !replica_ids.is_empty() {
                    if request.hydration {
                        live.hydrated_replicas =
                            ctx.hydrated_replicas(state.cluster_id, &replica_ids).await;
                    }
                    for reference in request.readiness {
                        let ready = ctx
                            .ready_replicas(state.cluster_id, &replica_ids, &reference)
                            .await;
                        live.ready_replicas.insert(reference, ready);
                    }
                }
            }
            if request.refresh_window {
                refresh_window_clusters.push(state.cluster_id);
            }
            signals.insert(state.cluster_id, live);
        }
        if !refresh_window_clusters.is_empty() {
            match ctx.refresh_window_inputs(&refresh_window_clusters).await {
                Some(batch) => {
                    let read_ts = batch.read_ts;
                    let mut cluster_inputs = batch.cluster_inputs;
                    for cluster_id in refresh_window_clusters {
                        let Some(inputs) = cluster_inputs.remove(&cluster_id) else {
                            signals.remove(&cluster_id);
                            continue;
                        };
                        let live = signals
                            .get_mut(&cluster_id)
                            .expect("signal entry inserted for requested cluster");
                        live.refresh_window = Some(RefreshWindowInputs {
                            read_ts,
                            compaction_estimate: inputs.compaction_estimate,
                            refresh_mvs: inputs.refresh_mvs,
                        });
                    }
                }
                None => {
                    for cluster_id in refresh_window_clusters {
                        signals.remove(&cluster_id);
                    }
                }
            }
        }
        signals
    }

    /// Diff the unioned desired set against the actual replicas of one cluster
    /// and emit the create/drop decisions that close the gap.
    fn collect_replica_decisions(
        &self,
        state: &ClusterState,
        signals: &LiveSignals,
        config: &ConfigSignals,
        now: mz_repr::Timestamp,
    ) -> Vec<Decision> {
        let contributions: Vec<Vec<DesiredReplica>> = self
            .strategies
            .iter()
            .map(|strategy| strategy.desired_replicas(state, signals, config, now))
            .collect();

        reconcile_replicas(state, &contributions)
    }
}

/// Join one `StateWrite` field across the strategies that set it: `None` if
/// none did, the common value if one or more set it to the same value, and
/// `None` with `field` pushed onto `conflicts` if two set it to different
/// values. The result and the conflict signal depend only on the set of values,
/// not the order they arrive in.
fn join<T: PartialEq>(
    field: &'static str,
    values: impl IntoIterator<Item = Option<T>>,
    conflicts: &mut Vec<&'static str>,
) -> Option<T> {
    let mut merged: Option<T> = None;
    for value in values.into_iter().flatten() {
        match &merged {
            None => merged = Some(value),
            Some(existing) if *existing == value => {}
            // Two strategies disagree on this field. Record it and leave the
            // field unchanged; merge_state_writes raises the alarm.
            Some(_) => {
                conflicts.push(field);
                return None;
            }
        }
    }
    merged
}

/// The pure multiset union/diff kernel for one cluster: given each strategy's
/// desired replica slots and the actual replicas, match slots to replicas by
/// shape and emit the creates and drops that close the gap.
///
/// Semantics:
/// - The desired set is the multiset **union** of every strategy's slots: a
///   given shape is desired `max` over strategies (not the sum), since a replica
///   of that shape satisfies every strategy that wants one. This is what makes a
///   replica survive iff *some* strategy desires its shape.
/// - For each shape, if actual count < desired count we create the difference;
///   if actual count > desired count we drop the difference, picking specific
///   excess replicas. A replica of a shape no strategy desires is dropped.
/// - Creates carry the winning [`CreateReason`] among the slots that
///   desired the shape (see [`CreateReason::outranks`]). Drops carry no
///   attribution. A drop happens exactly when no strategy desires the replica.
fn reconcile_replicas(
    state: &ClusterState,
    contributions: &[Vec<DesiredReplica>],
) -> Vec<Decision> {
    // Desired count per shape = max over strategies of how many that strategy
    // wants of the shape, carrying the highest-ranking reason among the
    // slots.
    let mut desired: Vec<DesiredShape> = Vec::new();
    for slots in contributions {
        // How many of each shape this strategy wants, and the winning reason
        // among the shape's slots.
        let mut per_shape: Vec<(ReplicaShape, usize, CreateReason)> = Vec::new();
        for slot in slots {
            match per_shape
                .iter_mut()
                .find(|(s, _, _)| s.matches(&slot.shape))
            {
                Some((_, count, reason)) => {
                    *count += 1;
                    if slot.reason.outranks(reason) {
                        *reason = slot.reason.clone();
                    }
                }
                None => per_shape.push((slot.shape.clone(), 1, slot.reason.clone())),
            }
        }
        for (shape, count, reason) in per_shape {
            match desired.iter_mut().find(|d| d.shape.matches(&shape)) {
                Some(existing) => {
                    existing.count = existing.count.max(count);
                    if reason.outranks(&existing.reason) {
                        existing.reason = reason;
                    }
                }
                None => desired.push(DesiredShape {
                    shape,
                    count,
                    reason,
                }),
            }
        }
    }

    // Bucket the controller-owned replicas by shape. Replicas the controller
    // does not own (see `ObservedReplica::owned_shape`) are invisible to the
    // desired/actual diff: neither counted toward a shape nor dropped.
    let mut actual_by_shape: Vec<(ReplicaShape, Vec<&ObservedReplica>)> = Vec::new();
    for replica in &state.replicas {
        let Some(shape) = replica.owned_shape() else {
            continue;
        };
        match actual_by_shape.iter_mut().find(|(s, _)| s.matches(shape)) {
            Some((_, replicas)) => replicas.push(replica),
            None => actual_by_shape.push((shape.clone(), vec![replica])),
        }
    }

    let mut decisions = Vec::new();

    // Every observed replica occupies a name, owned or not, so a generated
    // name never collides with a replica already on the cluster.
    let used_names: Vec<&str> = state.replicas.iter().map(|r| r.name.as_str()).collect();
    let mut name_gen = ReplicaNameGen::new(&used_names);

    // The compare-and-append witness for every create/drop this tick emits for
    // the cluster: the apply path rejects the batch if the cluster's durable
    // state has diverged from what we diffed against (e.g. a concurrent `ALTER`),
    // so a stale create/drop can never reshape the replica set against the new
    // config.
    let expected = state.expected();

    // Creates: for each desired shape, fill the gap below its desired count.
    for d in &desired {
        let actual_count = actual_by_shape
            .iter()
            .find(|(s, _)| s.matches(&d.shape))
            .map(|(_, replicas)| replicas.len())
            .unwrap_or(0);
        for _ in actual_count..d.count {
            decisions.push(Decision::CreateReplica {
                cluster_id: state.cluster_id,
                name: name_gen.next_name(),
                shape: d.shape.clone(),
                // Multiple creates of one shape in a tick share the merged
                // reason.
                reason: d.reason.clone(),
                expected: expected.clone(),
            });
        }
    }

    // Drops: any actual replica beyond the desired count for its shape, plus
    // every replica of a shape no strategy desires.
    for (shape, replicas) in &actual_by_shape {
        let desired_count = desired
            .iter()
            .find(|d| d.shape.matches(shape))
            .map(|d| d.count)
            .unwrap_or(0);
        for replica in replicas.iter().skip(desired_count) {
            decisions.push(Decision::DropReplica {
                cluster_id: state.cluster_id,
                replica_id: replica.replica_id,
                expected: expected.clone(),
            });
        }
    }

    decisions
}

/// A shape the union desires, how many, and the highest-ranking reason of
/// the strategies that wanted it.
struct DesiredShape {
    shape: ReplicaShape,
    count: usize,
    reason: CreateReason,
}

/// Generates deterministic fresh replica names that avoid a set of in-use names.
///
/// The controller derives names from the observed actual set rather than
/// renaming existing replicas, which keeps re-emission harmless. The concrete
/// naming convention (the `rNN` managed-replica scheme) is the environment's; the
/// kernel only needs distinct, stable-per-tick names, so it uses a simple
/// monotonic scheme starting past the highest observed `rNN` index, and never
/// below `r1` since managed-replica names are 1-based.
struct ReplicaNameGen {
    next: u32,
    used: BTreeSet<String>,
}

impl ReplicaNameGen {
    fn new(used: &[&str]) -> Self {
        let mut highest = 1;
        for name in used {
            if let Some(idx) = name.strip_prefix('r').and_then(|n| n.parse::<u32>().ok()) {
                highest = highest.max(idx + 1);
            }
        }
        Self {
            next: highest,
            used: used.iter().map(|n| n.to_string()).collect(),
        }
    }

    fn next_name(&mut self) -> String {
        loop {
            let name = format!("r{}", self.next);
            self.next += 1;
            if !self.used.contains(&name) {
                self.used.insert(name.clone());
                return name;
            }
        }
    }
}

#[cfg(test)]
mod tests;