mz_cluster_controller/

ctx.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 boundary between the controller and its environment.
//!
//! [`ClusterControllerCtx`] is the single, strategy-agnostic boundary through which
//! the controller pulls the signals a tick examines and applies the catalog
//! mutations it derives. The signals in are primitive and carry no per-strategy
//! state; the decisions out are primitive catalog mutations plus per-tick audit
//! attribution. A create carries the [`CreateReason`] of the winning strategy
//! behind it, which the environment turns into the audit event. The controller
//! crate knows nothing about the Coordinator. The Coordinator implements this
//! trait, which is what makes the controller testable against a fake
//! implementation and extractable later without touching controller code.
//!
//! The interface is **pull-based**: a tick fetches only the signals it actually
//! examines (no eager all-clusters-all-replicas snapshot is pushed in), and the
//! controller drives what is fetched. Read methods are batched so a separate-task
//! deployment can bound its round-trips to the Coordinator.

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

use async_trait::async_trait;
use mz_compute_types::config::ComputeReplicaLogging;
use mz_controller_types::{ClusterId, ReplicaId};
use mz_repr::refresh_schedule::RefreshSchedule;
use mz_repr::{GlobalId, Timestamp};
use timely::progress::Antichain;

// The compare-and-append witness types, and the replica shape they pair with,
// live in `mz-adapter-types` so the catalog transaction that applies a
// decision can share them without depending on this crate. They are part of the
// ctx vocabulary, so re-export them here.
pub use mz_adapter_types::cluster_state::{
    AutoScalingPolicy, AvailabilityZones, BurstAudit, BurstFinishCause, BurstRecord,
    ClusterSchedule, ExpectedClusterState, OnHydrationPolicy, OnTimeout, ReconfigurationAudit,
    ReconfigurationRecord, ReconfigurationStatus, ReconfigurationTarget, ReplicaShape,
};

/// A replica that actually exists on a cluster, as observed through the ctx.
/// Every replica physically on the cluster appears here, whether or not the
/// controller owns it; [`Self::owned_shape`] is the ownership test.
#[derive(Clone, Debug)]
pub struct ObservedReplica {
    pub replica_id: ReplicaId,
    pub name: String,
    /// `None` for a replica with an unmanaged location, which has no managed
    /// shape to reconcile against.
    pub shape: Option<ReplicaShape>,
    /// Created with `INTERNAL`.
    pub internal: bool,
    /// Carries a `BILLED AS` override.
    pub billed_as: bool,
    /// The `-pending` target of an in-flight graceful reconfiguration.
    pub pending: bool,
}

impl ObservedReplica {
    /// The replica's shape if the controller owns it, `None` otherwise.
    ///
    /// INTERNAL / BILLED AS replicas are manually managed: a user can attach
    /// one to any managed cluster, outside the replication-factor domain. A
    /// pending replica is owned by the reconfiguration sequencer path until
    /// finalize (retiring it would defeat the zero-downtime resize creating
    /// it). The controller must neither count such a replica toward a desired
    /// shape nor drop it as excess, but their names still block the name
    /// generator, since every replica observed here occupies a name.
    pub fn owned_shape(&self) -> Option<&ReplicaShape> {
        if self.internal || self.billed_as || self.pending {
            return None;
        }
        self.shape.as_ref()
    }
}

/// One REFRESH materialized view bound to a scheduled cluster, as the on-refresh
/// strategy needs to see it.
///
/// `write_frontier` is the MV's storage write frontier, carried with full
/// fidelity as the `Antichain` the storage controller reports. The strategy
/// compares it against the read timestamp (`less_than`) to decide whether the MV
/// still needs a refresh. For the compaction window it reads the frontier's lone
/// element via `as_option` to find the previous refresh time, falling back to the
/// schedule's last refresh on the empty/sealed frontier `[]`. The frontier of a
/// single-input total-order MV holds at most one element.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RefreshMvInfo {
    /// The MV's writes-`GlobalId`: the identity the window decision records in
    /// [`RefreshWindowDecision`] so the audit log can say which MVs kept the
    /// cluster on.
    pub id: GlobalId,
    pub write_frontier: Antichain<Timestamp>,
    pub refresh_schedule: RefreshSchedule,
}

/// Why a strategy desires a replica slot: the audit attribution a create
/// decision carries. When several strategies desire the same shape the
/// winning reason is decided by [`CreateReason::outranks`], since the audit
/// event carries exactly one reason.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CreateReason {
    /// The implicit baseline: the user's own cluster config calls for the
    /// replica. The environment audits it as a manual create.
    Baseline,
    /// The graceful-reconfiguration strategy converging an in-flight
    /// background `ALTER CLUSTER`.
    GracefulReconfiguration,
    /// The hydration-burst strategy accelerating hydration.
    HydrationBurst,
    /// The on-refresh strategy holding a scheduled cluster on inside a
    /// refresh window. Embeds the window decision behind the create, which
    /// the environment renders into the audit event's detail. Embedding it
    /// makes "the decision detail appears iff the create audits the schedule
    /// reason" structural: when another reason wins the precedence, the
    /// decision is discarded with it.
    OnRefresh(RefreshWindowDecision),
}

impl CreateReason {
    /// Whether this reason wins over `other` when several strategies desire
    /// the same shape, so the create audits this reason.
    ///
    /// A pairwise check rather than `Ord`: an `Ord` ranking by variant would
    /// have to call two `OnRefresh` reasons with different window decisions
    /// equal while `Eq` says they differ, violating the `Ord` contract.
    pub fn outranks(&self, other: &CreateReason) -> bool {
        self.rank() > other.rank()
    }

    fn rank(&self) -> u8 {
        match self {
            // Graceful wins over burst when both desire a shape (their shapes
            // differ in practice, so this is a stable tie-break), both win
            // over on-refresh, and the baseline loses to everything: any
            // strategy's reason beats the implicit "the config calls for it".
            CreateReason::Baseline => 0,
            CreateReason::OnRefresh(_) => 1,
            CreateReason::HydrationBurst => 2,
            CreateReason::GracefulReconfiguration => 3,
        }
    }
}

/// The on-refresh strategy's per-tick window decision: which bound REFRESH MVs
/// keep a scheduled cluster on, and why. The window is open iff either list is
/// non-empty, so an open window always has an explanation.
///
/// Carried inside [`CreateReason::OnRefresh`] on the create decisions the open
/// window produces. The environment converts it to the audit log's
/// `scheduling_policies` detail. Plain ids and durations so the controller crate
/// stays free of audit-log vocabulary.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RefreshWindowDecision {
    /// MVs whose write frontier has not yet passed the (hydration-adjusted) read
    /// timestamp: a refresh is due or imminent.
    pub objects_needing_refresh: Vec<GlobalId>,
    /// MVs estimated to still need Persist compaction after their last refresh.
    pub objects_needing_compaction: Vec<GlobalId>,
    /// The cluster's `HYDRATION TIME ESTIMATE` the refresh window was widened by.
    pub hydration_time_estimate: Duration,
}

impl RefreshWindowDecision {
    /// Whether the refresh window is open: some MV still needs a refresh or
    /// compaction time.
    pub fn window_open(&self) -> bool {
        !self.objects_needing_refresh.is_empty() || !self.objects_needing_compaction.is_empty()
    }
}

/// The catalog and storage inputs for one scheduled cluster's refresh window.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RefreshWindowClusterInputs {
    /// How long after a refresh an MV is estimated to still need Persist
    /// compaction, which also keeps the cluster on.
    pub compaction_estimate: Duration,
    /// The REFRESH MVs bound to the cluster.
    pub refresh_mvs: Vec<RefreshMvInfo>,
}

/// Refresh-window inputs gathered for one reconciliation phase.
///
/// The top-level timestamp makes sharing one oracle read across every included
/// cluster structural. A cluster absent from `cluster_inputs` has unavailable
/// inputs and must not be reconciled during the phase.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RefreshWindowInputsBatch {
    /// The local oracle read timestamp for every cluster in the batch.
    pub read_ts: Timestamp,
    /// The available catalog and storage inputs, keyed by cluster.
    pub cluster_inputs: BTreeMap<ClusterId, RefreshWindowClusterInputs>,
}

/// The fulfilled live signal the on-refresh strategy uses for one cluster.
///
/// Pulled on demand only for scheduled clusters. A MANUAL cluster carries
/// `None` and is never probed.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RefreshWindowInputs {
    /// The shared local oracle read timestamp the window decision uses.
    pub read_ts: Timestamp,
    /// How long after a refresh an MV is estimated to still need Persist
    /// compaction, which also keeps the cluster on.
    pub compaction_estimate: Duration,
    /// The REFRESH MVs bound to the cluster.
    pub refresh_mvs: Vec<RefreshMvInfo>,
}

/// The durable state of a single managed cluster plus its observed replicas, as
/// pulled through the ctx for one reconcile tick.
///
/// This is the input every strategy reads. Unmanaged clusters are not
/// controller-owned and are not represented here.
///
/// The `size`, `replication_factor`, `availability_zones`, `logging`, and
/// `arrangement_compression` fields together are the realized config the
/// cluster is currently serving. The implicit baseline desires
/// `replication_factor` replicas at that shape.
#[derive(Clone, Debug)]
pub struct ClusterState {
    pub cluster_id: ClusterId,
    pub size: String,
    pub replication_factor: u32,
    pub availability_zones: Vec<String>,
    pub logging: ComputeReplicaLogging,
    pub arrangement_compression: bool,
    /// The cluster's scheduling policy. Drives whether the implicit baseline owns
    /// the replica set (MANUAL) or the on-refresh strategy does (REFRESH).
    pub schedule: ClusterSchedule,
    pub auto_scaling_policy: Option<AutoScalingPolicy>,
    /// Latest graceful reconfiguration record, if one has been written.
    pub reconfiguration: Option<ReconfigurationRecord>,
    /// In-flight hydration burst, if any.
    pub burst: Option<BurstRecord>,
    /// The replicas that actually exist on the cluster, owned or not.
    pub replicas: Vec<ObservedReplica>,
}

impl ClusterState {
    /// The shape the implicit baseline desires: the realized config.
    pub fn realized_shape(&self) -> ReplicaShape {
        ReplicaShape {
            size: self.size.clone(),
            availability_zones: AvailabilityZones(self.availability_zones.clone()),
            logging: self.logging.clone(),
            arrangement_compression: self.arrangement_compression,
        }
    }

    /// The compare-and-append witness for decisions derived from this state: the
    /// durable fields a concurrent `ALTER` could change out from under a tick.
    pub fn expected(&self) -> ExpectedClusterState {
        ExpectedClusterState {
            size: self.size.clone(),
            replication_factor: self.replication_factor,
            availability_zones: AvailabilityZones(self.availability_zones.clone()),
            logging: self.logging.clone(),
            arrangement_compression: self.arrangement_compression,
            schedule: self.schedule,
            auto_scaling_policy: self.auto_scaling_policy.clone(),
            reconfiguration: self.reconfiguration.clone(),
            burst: self.burst.clone(),
        }
    }
}

/// A durable state mutation a strategy's `update_state` asks for: cut over the
/// realized config to a target and/or write or clear the reconfiguration/burst
/// records. The reconcile kernel pairs it with the [`ExpectedClusterState`] it
/// was derived from for the compare-and-append guard.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct StateWrite {
    /// New realized config to cut over to. `None` leaves it unchanged.
    pub new_size: Option<String>,
    pub new_replication_factor: Option<u32>,
    pub new_availability_zones: Option<Vec<String>>,
    pub new_logging: Option<ComputeReplicaLogging>,
    pub new_arrangement_compression: Option<bool>,
    /// Write or clear the reconfiguration record, together with its audit
    /// intent. `None` leaves the record unchanged.
    pub reconfiguration: Option<ReconfigurationWrite>,
    /// Write or clear the burst record, together with its audit intent.
    /// `None` leaves the record unchanged.
    pub burst: Option<BurstWrite>,
}

/// A write to the `reconfiguration` record, bundled with the audit intent
/// declaring which lifecycle transition the write represents.
///
/// Bundling means a writer cannot move the record without deciding, at the same
/// decision point, what the papertrail should say. The two travel together
/// through the apply path and are transacted atomically with the state.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ReconfigurationWrite {
    /// The record to write, or `None` to clear it.
    pub record: Option<ReconfigurationRecord>,
    /// The lifecycle transition to audit. `None` declares that this write is
    /// not a lifecycle transition and must not emit an event.
    pub audit: Option<ReconfigurationAudit>,
}

/// A write to the `burst` record, bundled with its audit intent. See
/// [`ReconfigurationWrite`]. A bookkeeping rewrite of an existing record (the
/// linger stamp and its reset) declares `audit: None`.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct BurstWrite {
    /// The record to write, or `None` to clear it.
    pub record: Option<BurstRecord>,
    /// The lifecycle transition to audit, or `None` for a bookkeeping rewrite.
    pub audit: Option<BurstAudit>,
}

impl StateWrite {
    /// Whether this write would actually mutate any durable field.
    pub fn is_empty(&self) -> bool {
        // Exhaustive destructure (no `..`): a field added to `StateWrite` is a
        // compile error here until it's accounted for.
        let StateWrite {
            new_size,
            new_replication_factor,
            new_availability_zones,
            new_logging,
            new_arrangement_compression,
            reconfiguration,
            burst,
        } = self;
        new_size.is_none()
            && new_replication_factor.is_none()
            && new_availability_zones.is_none()
            && new_logging.is_none()
            && new_arrangement_compression.is_none()
            && reconfiguration.is_none()
            && burst.is_none()
    }
}

/// A single command the controller emits for the environment to transact. The
/// apply path interprets these and turns them into catalog operations.
///
/// Every variant carries the [`ExpectedClusterState`] the decision was derived
/// from. The apply path re-reads each target cluster and rejects the whole batch
/// if any state has since diverged (compare-and-append), so a user `ALTER` that
/// lands mid-tick cannot let a stale create or drop reshape the replica set
/// against the new config; the controller recomputes from the new state next
/// tick.
#[derive(Clone, Debug)]
pub enum Decision {
    /// Create a replica of the given shape under a deterministic fresh name.
    /// `reason` is the audit attribution: the winning [`CreateReason`] among
    /// the strategies that desired the shape (see [`CreateReason::outranks`]).
    CreateReplica {
        cluster_id: ClusterId,
        name: String,
        shape: ReplicaShape,
        reason: CreateReason,
        expected: ExpectedClusterState,
    },
    /// Drop a specific existing replica. A drop happens exactly when no
    /// strategy desires the replica, so it carries no strategy attribution;
    /// the apply path audits every controller drop with the uniform `retired`
    /// reason.
    DropReplica {
        cluster_id: ClusterId,
        replica_id: ReplicaId,
        expected: ExpectedClusterState,
    },
    /// Apply a durable state write under a compare-and-append guard against
    /// `expected`.
    UpdateClusterState {
        cluster_id: ClusterId,
        expected: ExpectedClusterState,
        write: StateWrite,
    },
}

/// The outcome of applying one tick's batch of [`Decision`]s.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ApplyOutcome {
    /// Every decision in the batch was transacted.
    Applied,
    /// At least one decision failed its compare-and-append guard. The whole
    /// batch is rejected; the controller recomputes next tick.
    Rejected,
    /// The batch was rejected because it exceeded the environment's resource
    /// budget. Nothing was transacted. Unlike a guard rejection, retrying the
    /// same batch cannot succeed on its own: the controller decides what to
    /// shed to make room.
    ResourceExhausted,
}

/// The strategy-agnostic pull/apply interface between the controller and its
/// environment.
///
/// The controller depends on exactly this trait. Reads are batched and pulled
/// on demand; the single write applies a tick's batch under a compare-and-append
/// guard. Implementations marshal these to wherever the live signals live (for
/// v1, the Coordinator's catalog and compute/storage controllers, reached over a
/// channel from the controller's own task, hence the `Send` bound).
#[async_trait]
pub trait ClusterControllerCtx: Send {
    /// Current wall-clock time, as the controller's strategies should see it.
    fn now(&self) -> Timestamp;

    /// A consistent durable view of the given managed clusters and their
    /// replicas. Clusters that do not exist or are unmanaged are omitted from
    /// the result.
    async fn cluster_states(&mut self, clusters: &[ClusterId]) -> Vec<ClusterState>;

    /// The ids of all managed clusters the controller owns this tick.
    async fn managed_cluster_ids(&mut self) -> Vec<ClusterId>;

    /// Of `replicas` on `cluster`, which are online and have *all* current
    /// (non-transient) collections on the cluster hydrated. The returned set
    /// is a subset of `replicas`.
    ///
    /// Callers should request only replicas their strategy currently needs. This
    /// keeps live-signal dependencies local to the strategies that consume them.
    async fn hydrated_replicas(
        &mut self,
        cluster_id: ClusterId,
        replicas: &[ReplicaId],
    ) -> BTreeSet<ReplicaId>;

    /// Whether `cluster_id` has at least one hydratable (dataflow-backed) object
    /// bound to it: an index, materialized view, ingestion source, or sink.
    ///
    /// A catalog-level approximation of "the hydration check has something to
    /// count". Where the two disagree at the margin, the mismatch is
    /// self-healing: a replica with nothing to hydrate reads hydrated, and the
    /// burst winds down via its linger.
    async fn has_hydratable_objects(&mut self, cluster_id: ClusterId) -> bool;

    /// The refresh-window live signals for the given scheduled clusters.
    /// Returns one shared read timestamp plus the available per-cluster catalog
    /// and storage inputs. Omits a cluster when its inputs are unavailable,
    /// including when it is missing, unmanaged, or no longer scheduled `ON
    /// REFRESH` at pull time. Returns `None` when the batch fails or no requested
    /// cluster has valid inputs.
    ///
    /// Pulled on demand the same way as [`Self::hydrated_replicas`]: the
    /// controller includes a cluster only when the on-refresh strategy needs the
    /// signal (i.e. the cluster is scheduled), so a steady MANUAL cluster never
    /// pays for it. Implementations fetch the shared read timestamp once after
    /// gathering the per-cluster inputs, so oracle latency does not scale with
    /// the cluster count. The controller skips any omitted cluster for the
    /// reconciliation phase.
    async fn refresh_window_inputs(
        &mut self,
        cluster_ids: &[ClusterId],
    ) -> Option<RefreshWindowInputsBatch>;

    /// Apply a tick's batch of decisions under their compare-and-append guards.
    /// Each decision carries the [`ExpectedClusterState`] it was derived from;
    /// the implementation re-reads every target cluster and, if any has since
    /// diverged, returns [`ApplyOutcome::Rejected`] without transacting anything.
    /// Otherwise the batch's catalog operations are transacted together.
    async fn apply(&mut self, decisions: Vec<Decision>) -> ApplyOutcome;
}