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_controller_types::ClusterId;
use mz_ore::soft_panic_or_log;

use crate::ctx::{
    ApplyOutcome, ClusterControllerCtx, ClusterState, Decision, ObservedReplica,
    ReconfigurationAudit, ReconfigurationRecord, ReconfigurationStatus, ReconfigurationWrite,
    ReplicaShape, StateWrite,
};
use crate::strategy::{
    BaselineStrategy, DesiredReplica, GracefulReconfigurationStrategy, LiveSignals, 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>>,
}

impl Default for ClusterController {
    fn default() -> Self {
        Self::new()
    }
}

impl ClusterController {
    /// A controller with the implicit baseline and the graceful-reconfiguration
    /// strategy. The baseline holds the steady set. Graceful engages only while a
    /// `reconfiguration` record is in flight.
    pub fn new() -> Self {
        Self {
            strategies: vec![
                Box::new(BaselineStrategy),
                Box::new(GracefulReconfigurationStrategy),
            ],
        }
    }

    /// 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 signals = self.fetch_signals(ctx, &states).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 write = self.merge_state_writes(state, &signals[&state.cluster_id], 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).await;
            (states, signals)
        } else {
            (states, signals)
        };
        let now = ctx.now();
        for state in &states {
            if rejected.contains(&state.cluster_id) {
                continue;
            }
            let decisions = self.collect_replica_decisions(state, &signals[&state.cluster_id], 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. The graceful reconfiguration is the most expendable:
    /// a discretionary user change that fails cleanly (audited, and the wait-shim
    /// reports a timeout) and can be retried, while aborting it leaves the
    /// cluster running at its realized shape. The baseline is never shed, it is
    /// the committed floor.
    ///
    /// We shed one strategy per exhausted apply. If that was not enough, the
    /// next tick recomputes and sheds the next one.
    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: 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.
    /// Every field is owned by exactly one strategy, so by design it 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,
        now: mz_repr::Timestamp,
    ) -> StateWrite {
        let writes: Vec<StateWrite> = self
            .strategies
            .iter()
            .map(|strategy| strategy.update_state(state, signals, 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 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,
            ),
            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(", "),
            );
        }

        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
    /// ([`Strategy::signal_request`]), so the kernel stays ignorant of when a
    /// strategy engages. Signals are fetched per cluster and only where
    /// requested: a steady cluster is never probed, keeping the ctx seam
    /// pay-for-what-you-use. The returned map has an entry for every state.
    async fn fetch_signals(
        &self,
        ctx: &mut dyn ClusterControllerCtx,
        states: &[ClusterState],
    ) -> BTreeMap<ClusterId, LiveSignals> {
        let mut signals = BTreeMap::new();
        for state in states {
            let request = self
                .strategies
                .iter()
                .fold(SignalRequest::default(), |acc, strategy| {
                    acc.union(strategy.signal_request(state))
                });
            let mut live = LiveSignals::default();
            if request.hydration {
                let replica_ids: Vec<_> = state.replicas.iter().map(|r| r.replica_id).collect();
                if !replica_ids.is_empty() {
                    live.hydrated_replicas =
                        ctx.hydrated_replicas(state.cluster_id, &replica_ids).await;
                }
            }
            signals.insert(state.cluster_id, live);
        }
        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,
        now: mz_repr::Timestamp,
    ) -> Vec<Decision> {
        // Each strategy's contribution, tagged with the strategy name for
        // attribution.
        let contributions: Vec<(&'static str, Vec<DesiredReplica>)> = self
            .strategies
            .iter()
            .map(|strategy| {
                (
                    strategy.name(),
                    strategy.desired_replicas(state, signals, 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 names of the strategies that desired the shape. Drops
///   carry no attribution, because a drop happens exactly when no strategy
///   desires the replica.
fn reconcile_replicas(
    state: &ClusterState,
    contributions: &[(&'static str, Vec<DesiredReplica>)],
) -> Vec<Decision> {
    // Desired count per shape = max over strategies of how many that strategy
    // wants of the shape, and the union of which strategies want it.
    let mut desired: Vec<DesiredShape> = Vec::new();
    for (name, slots) in contributions {
        // How many of each shape this strategy wants.
        let mut per_shape: Vec<(ReplicaShape, usize)> = Vec::new();
        for slot in slots {
            match per_shape.iter_mut().find(|(s, _)| s.matches(&slot.shape)) {
                Some((_, count)) => *count += 1,
                None => per_shape.push((slot.shape.clone(), 1)),
            }
        }
        for (shape, count) in per_shape {
            match desired.iter_mut().find(|d| d.shape.matches(&shape)) {
                Some(existing) => {
                    existing.count = existing.count.max(count);
                    if !existing.reasons.contains(name) {
                        existing.reasons.push(*name);
                    }
                }
                None => desired.push(DesiredShape {
                    shape,
                    count,
                    reasons: vec![*name],
                }),
            }
        }
    }

    // Bucket the actual replicas by shape.
    let mut actual_by_shape: Vec<(ReplicaShape, Vec<&ObservedReplica>)> = Vec::new();
    for replica in &state.replicas {
        match actual_by_shape
            .iter_mut()
            .find(|(s, _)| s.matches(&replica.shape))
        {
            Some((_, replicas)) => replicas.push(replica),
            None => actual_by_shape.push((replica.shape.clone(), vec![replica])),
        }
    }

    let mut decisions = Vec::new();

    // Track existing names so freshly-created replicas avoid collisions.
    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(),
                reasons: d.reasons.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 which strategies wanted it.
struct DesiredShape {
    shape: ReplicaShape,
    count: usize,
    reasons: Vec<&'static str>,
}

/// 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;