mz_storage_controller/

instance.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.

//! A controller for a storage instance.

use crate::CollectionMetadata;
use std::collections::{BTreeMap, BTreeSet};
use std::sync::atomic::AtomicBool;
use std::sync::{Arc, atomic};
use std::time::{Duration, Instant};

use anyhow::bail;
use itertools::Itertools;
use mz_build_info::BuildInfo;
use mz_cluster_client::ReplicaId;
use mz_cluster_client::client::ClusterReplicaLocation;
use mz_dyncfg::ConfigUpdates;
use mz_ore::cast::CastFrom;
use mz_ore::now::NowFn;
use mz_ore::retry::{Retry, RetryState};
use mz_ore::task::AbortOnDropHandle;
use mz_repr::{GlobalId, Timestamp};
use mz_service::client::{GenericClient, Partitioned};
use mz_service::params::GrpcClientParameters;
use mz_service::transport;
use mz_storage_client::client::{
    RunIngestionCommand, RunSinkCommand, Status, StatusUpdate, StorageCommand, StorageResponse,
};
use mz_storage_client::metrics::{InstanceMetrics, ReplicaMetrics};
use mz_storage_types::sinks::StorageSinkDesc;
use mz_storage_types::sources::{IngestionDescription, SourceConnection};
use timely::progress::Antichain;
use tokio::select;
use tokio::sync::mpsc;
use tracing::{debug, info, warn};
use uuid::Uuid;

use crate::history::CommandHistory;

/// A controller for a storage instance.
///
/// Encapsulates communication with replicas in this instance, and their rehydration.
///
/// An instance can have any number of replicas. Objects that must run on one replica (ingestions
/// whose connection reports `SourceConnection::prefers_single_replica`, such as OLTP sources, and
/// all sinks) are scheduled on the replica with the lowest `ReplicaId` and move only when that
/// replica goes away; all other objects run on every replica. See `update_scheduling`.
#[derive(Debug)]
pub(crate) struct Instance {
    /// The workload class of this instance.
    ///
    /// This is currently only used to annotate metrics.
    pub workload_class: Option<String>,
    /// The replicas connected to this storage instance.
    replicas: BTreeMap<ReplicaId, Replica>,
    /// The ingestions currently running on this instance.
    ///
    /// While this is derivable from `history` on demand, keeping a denormalized
    /// list of running ingestions is quite a bit more convenient in the
    /// implementation of `StorageController::active_ingestions`.
    active_ingestions: BTreeMap<GlobalId, ActiveIngestion>,
    /// A map from ingestion export ID to the ingestion that is producing it.
    ingestion_exports: BTreeMap<GlobalId, GlobalId>,
    /// The exports currently running on this instance.
    ///
    /// While this is derivable from `history` on demand, keeping a denormalized
    /// list of running exports is quite a bit more convenient for the
    /// controller.
    active_exports: BTreeMap<GlobalId, ActiveExport>,
    /// The command history, used to replay past commands when introducing new replicas or
    /// reconnecting to existing replicas.
    history: CommandHistory,
    /// Per-replica dyncfg overrides, merged into `UpdateConfiguration` commands
    /// when they are sent or replayed to the overridden replica. The history
    /// records the base (un-specialized) commands, so overrides are re-applied
    /// at replay time rather than baked into the shared history.
    replica_dyncfg_overrides: BTreeMap<ReplicaId, ConfigUpdates>,
    /// Metrics tracked for this storage instance.
    metrics: InstanceMetrics,
    /// A function that returns the current time.
    now: NowFn,
    /// A sender for responses from replicas.
    ///
    /// Responses are tagged with the [`ReplicaId`] of the replica that sent the
    /// response. Responses that don't originate from a replica (e.g. a "paused"
    /// status update, when no replicas are connected) are tagged with `None`.
    response_tx: mpsc::UnboundedSender<(Option<ReplicaId>, StorageResponse)>,
}

#[derive(Debug)]
struct ActiveIngestion {
    /// The set of replicas that this ingestion is currently running on.
    active_replicas: BTreeSet<ReplicaId>,
}

#[derive(Debug)]
struct ActiveExport {
    /// The set of replicas that this export is currently running on.
    active_replicas: BTreeSet<ReplicaId>,
}

/// Which replicas actively run a given object, as resolved by
/// [`Instance::active_replica_ids`].
enum ActiveReplicas<'a> {
    /// The object has per-replica scheduling and runs on exactly these
    /// replicas. The set is empty when it currently runs nowhere, e.g. it has
    /// been compacted away or is not yet scheduled onto a replica.
    Scheduled(&'a BTreeSet<ReplicaId>),
    /// The object has no per-replica scheduling, so every replica runs it.
    All,
}

impl Instance {
    /// Creates a new [`Instance`].
    pub fn new(
        workload_class: Option<String>,
        metrics: InstanceMetrics,
        now: NowFn,
        instance_response_tx: mpsc::UnboundedSender<(Option<ReplicaId>, StorageResponse)>,
    ) -> Self {
        let history = CommandHistory::new(metrics.for_history());

        let mut instance = Self {
            workload_class,
            replicas: Default::default(),
            active_ingestions: Default::default(),
            ingestion_exports: Default::default(),
            active_exports: BTreeMap::new(),
            history,
            replica_dyncfg_overrides: Default::default(),
            metrics,
            now,
            response_tx: instance_response_tx,
        };

        instance.send(StorageCommand::Hello {
            // The nonce is protocol iteration-specific and will be set in
            // `ReplicaTask::specialize_command`.
            nonce: Default::default(),
        });

        instance
    }

    /// Returns the IDs of all replicas connected to this storage instance.
    pub fn replica_ids(&self) -> impl Iterator<Item = ReplicaId> + '_ {
        self.replicas.keys().copied()
    }

    /// Adds a new replica to this storage instance.
    pub fn add_replica(&mut self, id: ReplicaId, config: ReplicaConfig) {
        // Reduce the history to limit the amount of commands sent to the new replica.
        self.history.reduce();

        let metrics = self.metrics.for_replica(id);
        let replica = Replica::new(id, config, metrics, self.response_tx.clone());

        self.replicas.insert(id, replica);

        self.update_scheduling(false);

        self.replay_commands(id);
    }

    /// Replays commands to the specified replica.
    pub fn replay_commands(&mut self, replica_id: ReplicaId) {
        let commands = self.history.iter().cloned();

        let filtered_commands = commands
            .filter_map(|command| match command {
                StorageCommand::RunIngestion(ingestion) => {
                    if self.is_active_replica(&ingestion.id, &replica_id) {
                        Some(StorageCommand::RunIngestion(ingestion))
                    } else {
                        None
                    }
                }
                StorageCommand::RunSink(sink) => {
                    if self.is_active_replica(&sink.id, &replica_id) {
                        Some(StorageCommand::RunSink(sink))
                    } else {
                        None
                    }
                }
                StorageCommand::AllowCompaction(id, upper) => {
                    if self.is_active_replica(&id, &replica_id) {
                        Some(StorageCommand::AllowCompaction(id, upper))
                    } else {
                        None
                    }
                }
                command => Some(command),
            })
            .collect::<Vec<_>>();

        let replica = self
            .replicas
            .get_mut(&replica_id)
            .expect("replica must exist");

        // Replay the commands at the new replica, re-applying its dyncfg
        // override to configuration commands.
        for command in filtered_commands {
            let command = Self::specialize_command_for_replica(
                command,
                replica_id,
                &self.replica_dyncfg_overrides,
            );
            replica.send(command);
        }
    }

    /// Removes the identified replica from this storage instance.
    pub fn drop_replica(&mut self, id: ReplicaId) {
        let replica = self.replicas.remove(&id);

        // The coordinator only re-pushes the override map when the scoped configuration itself
        // changes, so a dropped replica's entry would otherwise be retained until the next such
        // change.
        self.replica_dyncfg_overrides.remove(&id);

        let mut needs_rescheduling = false;
        for (ingestion_id, ingestion) in self.active_ingestions.iter_mut() {
            let was_running = ingestion.active_replicas.remove(&id);
            if was_running {
                tracing::debug!(
                    %ingestion_id,
                    replica_id = %id,
                    "ingestion was running on dropped replica, updating scheduling decisions"
                );
                needs_rescheduling = true;
            }
        }
        for (export_id, export) in self.active_exports.iter_mut() {
            let was_running = export.active_replicas.remove(&id);
            if was_running {
                tracing::debug!(
                    %export_id,
                    replica_id = %id,
                    "export was running on dropped replica, updating scheduling decisions"
                );
                needs_rescheduling = true;
            }
        }

        tracing::info!(%id, %needs_rescheduling, "dropped replica");

        if needs_rescheduling {
            self.update_scheduling(true);
        }

        if replica.is_some() && self.replicas.is_empty() {
            self.update_paused_statuses();
        }
    }

    /// Rehydrates any failed replicas of this storage instance.
    pub fn rehydrate_failed_replicas(&mut self) {
        let replicas = self.replicas.iter();
        let failed_replicas: Vec<_> = replicas
            .filter_map(|(id, replica)| replica.failed().then_some(*id))
            .collect();

        for id in failed_replicas {
            let replica = self.replicas.remove(&id).expect("must exist");
            self.add_replica(id, replica.config);
        }
    }

    /// Returns ingestions running on this instance. This _only_ includes the
    /// "toplevel" ingestions, not any of their source tables (aka. subsources).
    pub fn active_ingestions(&self) -> impl Iterator<Item = &GlobalId> {
        self.active_ingestions.keys()
    }

    /// Returns ingestion exports running on this instance. This includes the
    /// ingestion itself, if any, and running source tables (aka. subsources).
    ///
    /// This does _not_ filter out exports whose write frontier is the empty
    /// frontier, which some might consider not active anymore. But for the
    /// purposes of the instance controller, these are still considered active
    /// because we don't know about frontiers.
    pub fn active_ingestion_exports(&self) -> impl Iterator<Item = &GlobalId> {
        let ingestion_exports = self.ingestion_exports.keys();
        self.active_ingestions.keys().chain(ingestion_exports)
    }

    /// Returns the exports running on this instance.
    pub fn active_exports(&self) -> impl Iterator<Item = &GlobalId> {
        self.active_exports.keys()
    }

    /// Sets the status to paused for all sources/sinks in the history.
    fn update_paused_statuses(&mut self) {
        let now = mz_ore::now::to_datetime((self.now)());
        let make_update = |id, object_type| StatusUpdate {
            id,
            status: Status::Paused,
            timestamp: now,
            error: None,
            hints: BTreeSet::from([format!(
                "There is currently no replica running this {object_type}"
            )]),
            namespaced_errors: Default::default(),
            replica_id: None,
        };

        self.history.reduce();

        let mut status_updates = Vec::new();
        for command in self.history.iter() {
            match command {
                StorageCommand::RunIngestion(ingestion) => {
                    let old_style_ingestion =
                        ingestion.id != ingestion.description.remap_collection_id;
                    let subsource_ids = ingestion.description.collection_ids().filter(|id| {
                        // NOTE(aljoscha): We filter out the remap collection for old style
                        // ingestions because it doesn't get any status updates about it from the
                        // replica side. So we don't want to synthesize a 'paused' status here.
                        // New style ingestion do, since the source itself contains the remap data.
                        let should_discard =
                            old_style_ingestion && id == &ingestion.description.remap_collection_id;
                        !should_discard
                    });
                    for id in subsource_ids {
                        status_updates.push(make_update(id, "source"));
                    }
                }
                StorageCommand::RunSink(sink) => {
                    status_updates.push(make_update(sink.id, "sink"));
                }
                _ => (),
            }
        }

        for update in status_updates {
            // NOTE: If we lift this "inject paused status" logic to the
            // controller, we could instead return ReplicaId instead of an
            // Option<ReplicaId>.
            let _ = self
                .response_tx
                .send((None, StorageResponse::StatusUpdate(update)));
        }
    }

    /// Replaces the per-replica dyncfg overrides. Callers should follow this
    /// with a configuration push (e.g. an `UpdateConfiguration` command) so
    /// that existing replicas observe the new overrides.
    pub fn update_replica_dyncfg_overrides(
        &mut self,
        overrides: BTreeMap<ReplicaId, ConfigUpdates>,
    ) {
        self.replica_dyncfg_overrides = overrides;
    }

    /// Specializes a command for a specific replica by merging that replica's
    /// dyncfg override into its configuration. For `UpdateConfiguration` the
    /// override is merged into the update. All other commands are returned
    /// unchanged.
    fn specialize_command_for_replica(
        mut command: StorageCommand,
        replica_id: ReplicaId,
        overrides: &BTreeMap<ReplicaId, ConfigUpdates>,
    ) -> StorageCommand {
        if let StorageCommand::UpdateConfiguration(params) = &mut command
            && let Some(over) = overrides.get(&replica_id)
            && !over.updates.is_empty()
        {
            params.dyncfg_updates.extend(over.clone());
        }
        command
    }

    /// Sends a command to this storage instance.
    pub fn send(&mut self, command: StorageCommand) {
        // Record the command so that new replicas can be brought up to speed.
        self.history.push(command.clone());

        match command.clone() {
            StorageCommand::RunIngestion(ingestion) => {
                // First absorb into our state, because this might change
                // scheduling decisions, which need to be respected just below
                // when sending commands.
                self.absorb_ingestion(*ingestion.clone());

                for replica in self.active_replicas(&ingestion.id) {
                    replica.send(StorageCommand::RunIngestion(ingestion.clone()));
                }
            }
            StorageCommand::RunSink(sink) => {
                // First absorb into our state, because this might change
                // scheduling decisions, which need to be respected just below
                // when sending commands.
                self.absorb_export(*sink.clone());

                for replica in self.active_replicas(&sink.id) {
                    replica.send(StorageCommand::RunSink(sink.clone()));
                }
            }
            StorageCommand::AllowCompaction(id, frontier) => {
                // First send out commands and then absorb into our state since
                // absorbing them might remove entries from active_ingestions.
                for replica in self.active_replicas(&id) {
                    replica.send(StorageCommand::AllowCompaction(
                        id.clone(),
                        frontier.clone(),
                    ));
                }

                self.absorb_compaction(id, frontier);
            }
            command => {
                let overrides = &self.replica_dyncfg_overrides;
                for (replica_id, replica) in self.replicas.iter_mut() {
                    let command = Self::specialize_command_for_replica(
                        command.clone(),
                        *replica_id,
                        overrides,
                    );
                    replica.send(command);
                }
            }
        }

        if command.installs_objects() && self.replicas.is_empty() {
            self.update_paused_statuses();
        }
    }

    /// Updates internal state based on incoming ingestion commands.
    ///
    /// This does _not_ send commands to replicas, we only record the ingestion
    /// in state and potentially update scheduling decisions.
    fn absorb_ingestion(&mut self, ingestion: RunIngestionCommand) {
        let existing_ingestion_state = self.active_ingestions.get_mut(&ingestion.id);

        // Always update our mapping from export to their ingestion.
        for id in ingestion.description.source_exports.keys() {
            self.ingestion_exports.insert(id.clone(), ingestion.id);
        }

        if let Some(ingestion_state) = existing_ingestion_state {
            // It's an update for an existing ingestion. We don't need to
            // change anything about our scheduling decisions, no need to
            // update active_ingestions.

            tracing::debug!(
                ingestion_id = %ingestion.id,
                active_replicas = %ingestion_state.active_replicas.iter().map(|id| id.to_string()).join(", "),
                "updating ingestion"
            );
        } else {
            // We create a new ingestion state for this ingestion.
            let ingestion_state = ActiveIngestion {
                active_replicas: BTreeSet::new(),
            };
            self.active_ingestions.insert(ingestion.id, ingestion_state);

            // Maybe update scheduling decisions.
            self.update_scheduling(false);
        }
    }

    /// Updates internal state based on incoming export commands.
    ///
    /// This does _not_ send commands to replicas, we only record the export
    /// in state and potentially update scheduling decisions.
    fn absorb_export(&mut self, export: RunSinkCommand) {
        let existing_export_state = self.active_exports.get_mut(&export.id);

        if let Some(export_state) = existing_export_state {
            // It's an update for an existing export. We don't need to
            // change anything about our scheduling decisions, no need to
            // update active_exports.

            tracing::debug!(
                export_id = %export.id,
                active_replicas = %export_state.active_replicas.iter().map(|id| id.to_string()).join(", "),
                "updating export"
            );
        } else {
            // We create a new export state for this export.
            let export_state = ActiveExport {
                active_replicas: BTreeSet::new(),
            };
            self.active_exports.insert(export.id, export_state);

            // Maybe update scheduling decisions.
            self.update_scheduling(false);
        }
    }

    /// Update scheduling decisions, that is what replicas should be running a
    /// given object, if needed.
    ///
    /// An important property of this scheduling algorithm is that we never
    /// change the scheduling decision for single-replica objects unless we
    /// have to, that is unless the replica that they are running on goes away.
    /// We do this, so that we don't send a mix of "run"/"allow
    /// compaction"/"run" messages to replicas, which wouldn't deal well with
    /// this. When we _do_ have to make a scheduling decision we schedule a
    /// single-replica ingestion on the first replica, according to the sort
    /// order of `ReplicaId`. We do this latter so that the scheduling decision
    /// is stable across restarts of `environmentd`/the controller.
    ///
    /// For multi-replica objects (e.g. Kafka ingestions), each active object is
    /// scheduled on all replicas.
    ///
    /// If `send_commands` is true, will send commands for newly-scheduled
    /// single-replica objects.
    fn update_scheduling(&mut self, send_commands: bool) {
        #[derive(Debug)]
        enum ObjectId {
            Ingestion(GlobalId),
            Export(GlobalId),
        }
        // We first collect scheduling preferences and then schedule below.
        // Applying the decision needs a mutable borrow but we also need to
        // borrow for determining `prefers_single_replica`, so we split this
        // into two loops.
        let mut scheduling_preferences: Vec<(ObjectId, bool)> = Vec::new();

        for ingestion_id in self.active_ingestions.keys() {
            let ingestion_description = self
                .get_ingestion_description(ingestion_id)
                .expect("missing ingestion description");

            let prefers_single_replica = ingestion_description
                .desc
                .connection
                .prefers_single_replica();

            scheduling_preferences
                .push((ObjectId::Ingestion(*ingestion_id), prefers_single_replica));
        }

        for export_id in self.active_exports.keys() {
            // All sinks prefer single replica
            scheduling_preferences.push((ObjectId::Export(*export_id), true));
        }

        // Collect all commands per replica and send them in one go.
        let mut commands_by_replica: BTreeMap<ReplicaId, Vec<ObjectId>> = BTreeMap::new();

        for (object_id, prefers_single_replica) in scheduling_preferences {
            let active_replicas = match object_id {
                ObjectId::Ingestion(ingestion_id) => {
                    &mut self
                        .active_ingestions
                        .get_mut(&ingestion_id)
                        .expect("missing ingestion state")
                        .active_replicas
                }
                ObjectId::Export(export_id) => {
                    &mut self
                        .active_exports
                        .get_mut(&export_id)
                        .expect("missing ingestion state")
                        .active_replicas
                }
            };

            if prefers_single_replica {
                // For single-replica ingestion, schedule only if it's not already running.
                if active_replicas.is_empty() {
                    let target_replica = self.replicas.keys().min().copied();
                    if let Some(first_replica_id) = target_replica {
                        tracing::info!(
                            object_id = ?object_id,
                            replica_id = %first_replica_id,
                            "scheduling single-replica object");
                        active_replicas.insert(first_replica_id);

                        commands_by_replica
                            .entry(first_replica_id)
                            .or_default()
                            .push(object_id);
                    }
                } else {
                    tracing::info!(
                        ?object_id,
                        active_replicas = %active_replicas.iter().map(|id| id.to_string()).join(", "),
                        "single-replica object already running, not scheduling again",
                    );
                }
            } else {
                let current_replica_ids: BTreeSet<_> = self.replicas.keys().copied().collect();
                let unscheduled_replicas: Vec<_> = current_replica_ids
                    .difference(active_replicas)
                    .copied()
                    .collect();
                for replica_id in unscheduled_replicas {
                    tracing::info!(
                        ?object_id,
                        %replica_id,
                        "scheduling multi-replica object"
                    );
                    active_replicas.insert(replica_id);
                }
            }
        }

        if send_commands {
            for (replica_id, object_ids) in commands_by_replica {
                let mut ingestion_commands = vec![];
                let mut export_commands = vec![];
                for object_id in object_ids {
                    match object_id {
                        ObjectId::Ingestion(id) => {
                            ingestion_commands.push(RunIngestionCommand {
                                id,
                                description: self
                                    .get_ingestion_description(&id)
                                    .expect("missing ingestion description")
                                    .clone(),
                            });
                        }
                        ObjectId::Export(id) => {
                            export_commands.push(RunSinkCommand {
                                id,
                                description: self
                                    .get_export_description(&id)
                                    .expect("missing export description")
                                    .clone(),
                            });
                        }
                    }
                }
                for ingestion in ingestion_commands {
                    let replica = self.replicas.get_mut(&replica_id).expect("missing replica");
                    let ingestion = Box::new(ingestion);
                    replica.send(StorageCommand::RunIngestion(ingestion));
                }
                for export in export_commands {
                    let replica = self.replicas.get_mut(&replica_id).expect("missing replica");
                    let export = Box::new(export);
                    replica.send(StorageCommand::RunSink(export));
                }
            }
        }
    }

    /// Returns the ingestion description for the given ingestion ID, if it
    /// exists.
    ///
    /// This function searches through the command history to find the most
    /// recent RunIngestionCommand for the specified ingestion ID and returns
    /// its description.  Returns None if no ingestion with the given ID is
    /// found.
    pub fn get_ingestion_description(
        &self,
        id: &GlobalId,
    ) -> Option<IngestionDescription<CollectionMetadata>> {
        if !self.active_ingestions.contains_key(id) {
            return None;
        }

        self.history.iter().rev().find_map(|command| {
            if let StorageCommand::RunIngestion(ingestion) = command {
                if &ingestion.id == id {
                    Some(ingestion.description.clone())
                } else {
                    None
                }
            } else {
                None
            }
        })
    }

    /// Returns the export description for the given export ID, if it
    /// exists.
    ///
    /// This function searches through the command history to find the most
    /// recent RunSinkCommand for the specified export ID and returns
    /// its description.  Returns None if no ingestion with the given ID is
    /// found.
    pub fn get_export_description(
        &self,
        id: &GlobalId,
    ) -> Option<StorageSinkDesc<CollectionMetadata>> {
        if !self.active_exports.contains_key(id) {
            return None;
        }

        self.history.iter().rev().find_map(|command| {
            if let StorageCommand::RunSink(sink) = command {
                if &sink.id == id {
                    Some(sink.description.clone())
                } else {
                    None
                }
            } else {
                None
            }
        })
    }

    /// Updates internal state based on incoming compaction commands.
    fn absorb_compaction(&mut self, id: GlobalId, frontier: Antichain<Timestamp>) {
        tracing::debug!(?self.active_ingestions, ?id, ?frontier, "allow_compaction");

        if frontier.is_empty() {
            self.active_ingestions.remove(&id);
            self.ingestion_exports.remove(&id);
            self.active_exports.remove(&id);
        }
    }

    /// Resolves an object to the replicas actively running it.
    ///
    /// Shared by [`Self::active_replicas`], [`Self::is_active_replica`], and
    /// [`Self::get_active_replicas_for_object`], which differ only in how they
    /// project the result.
    fn active_replica_ids(&self, id: &GlobalId) -> ActiveReplicas<'_> {
        // An empty set to borrow for objects whose scheduling target is gone.
        static EMPTY: BTreeSet<ReplicaId> = BTreeSet::new();

        if let Some(ingestion_id) = self.ingestion_exports.get(id) {
            match self.active_ingestions.get(ingestion_id) {
                Some(ingestion) => ActiveReplicas::Scheduled(&ingestion.active_replicas),
                // The ingestion has already been compacted away (aka. stopped).
                None => ActiveReplicas::Scheduled(&EMPTY),
            }
        } else if let Some(ingestion) = self.active_ingestions.get(id) {
            // A new-syntax source lists only its tables in `source_exports`, so
            // the primary ingestion id is not in `ingestion_exports`.
            ActiveReplicas::Scheduled(&ingestion.active_replicas)
        } else if let Some(export) = self.active_exports.get(id) {
            ActiveReplicas::Scheduled(&export.active_replicas)
        } else {
            // Objects that have no per-replica scheduling (e.g. tables and
            // webhooks) run on all replicas.
            ActiveReplicas::All
        }
    }

    /// Returns the replicas that are actively running the given object (ingestion or export).
    fn active_replicas(&mut self, id: &GlobalId) -> Box<dyn Iterator<Item = &mut Replica> + '_> {
        // Take an owned copy of the scheduled set so the immutable borrow of
        // `self` ends before we borrow `self.replicas` mutably below. The set
        // is tiny (one entry per replica) and this is not a per-row path.
        let scheduled = match self.active_replica_ids(id) {
            ActiveReplicas::All => None,
            ActiveReplicas::Scheduled(replicas) => Some(replicas.clone()),
        };
        match scheduled {
            None => Box::new(self.replicas.values_mut()),
            Some(scheduled) => Box::new(self.replicas.iter_mut().filter_map(
                move |(replica_id, replica)| scheduled.contains(replica_id).then_some(replica),
            )),
        }
    }

    /// Returns whether the given replica is actively running the given object (ingestion or export).
    fn is_active_replica(&self, id: &GlobalId, replica_id: &ReplicaId) -> bool {
        match self.active_replica_ids(id) {
            ActiveReplicas::All => true,
            ActiveReplicas::Scheduled(replicas) => replicas.contains(replica_id),
        }
    }

    /// Refresh the controller state metrics for this instance.
    ///
    /// We could also do state metric updates directly in response to state changes, but that would
    /// mean littering the code with metric update calls. Encapsulating state metric maintenance in
    /// a single method is less noisy.
    ///
    /// This method is invoked by `Controller::maintain`, which we expect to be called once per
    /// second during normal operation.
    pub(super) fn refresh_state_metrics(&self) {
        let connected_replica_count = self.replicas.values().filter(|r| r.is_connected()).count();

        self.metrics
            .connected_replica_count
            .set(u64::cast_from(connected_replica_count));
    }

    /// Returns the set of replica IDs that are actively running the given
    /// object (ingestion, ingestion export (aka. subsource), or export).
    pub fn get_active_replicas_for_object(&self, id: &GlobalId) -> BTreeSet<ReplicaId> {
        match self.active_replica_ids(id) {
            ActiveReplicas::All => self.replicas.keys().copied().collect(),
            ActiveReplicas::Scheduled(replicas) => replicas.clone(),
        }
    }
}

/// Replica-specific configuration.
#[derive(Clone, Debug)]
pub(super) struct ReplicaConfig {
    pub build_info: &'static BuildInfo,
    pub location: ClusterReplicaLocation,
    pub grpc_client: GrpcClientParameters,
}

/// State maintained about individual replicas.
#[derive(Debug)]
pub struct Replica {
    /// Replica configuration.
    config: ReplicaConfig,
    /// A sender for commands for the replica.
    ///
    /// If sending to this channel fails, the replica has failed and requires
    /// rehydration.
    command_tx: mpsc::UnboundedSender<StorageCommand>,
    /// A handle to the task that aborts it when the replica is dropped.
    task: AbortOnDropHandle<()>,
    /// Flag reporting whether the replica connection has been established.
    connected: Arc<AtomicBool>,
}

impl Replica {
    /// Creates a new [`Replica`].
    fn new(
        id: ReplicaId,
        config: ReplicaConfig,
        metrics: ReplicaMetrics,
        response_tx: mpsc::UnboundedSender<(Option<ReplicaId>, StorageResponse)>,
    ) -> Self {
        let (command_tx, command_rx) = mpsc::unbounded_channel();
        let connected = Arc::new(AtomicBool::new(false));

        let task = mz_ore::task::spawn(
            || "storage-replica-{id}",
            ReplicaTask {
                replica_id: id,
                config: config.clone(),
                metrics: metrics.clone(),
                connected: Arc::clone(&connected),
                command_rx,
                response_tx,
            }
            .run(),
        );

        Self {
            config,
            command_tx,
            task: task.abort_on_drop(),
            connected,
        }
    }

    /// Sends a command to the replica.
    fn send(&self, command: StorageCommand) {
        // Send failures ignored, we'll check for failed replicas separately.
        let _ = self.command_tx.send(command);
    }

    /// Determine if this replica has failed. This is true if the replica
    /// task has terminated.
    fn failed(&self) -> bool {
        self.task.is_finished()
    }

    /// Determine if the replica connection has been established.
    pub(super) fn is_connected(&self) -> bool {
        self.connected.load(atomic::Ordering::Relaxed)
    }
}

type StorageCtpClient = transport::Client<StorageCommand, StorageResponse>;
type ReplicaClient = Partitioned<StorageCtpClient, StorageCommand, StorageResponse>;

/// A task handling communication with a replica.
struct ReplicaTask {
    /// The ID of the replica.
    replica_id: ReplicaId,
    /// Replica configuration.
    config: ReplicaConfig,
    /// Replica metrics.
    metrics: ReplicaMetrics,
    /// Flag to report successful replica connection.
    connected: Arc<AtomicBool>,
    /// A channel upon which commands intended for the replica are delivered.
    command_rx: mpsc::UnboundedReceiver<StorageCommand>,
    /// A channel upon which responses from the replica are delivered.
    response_tx: mpsc::UnboundedSender<(Option<ReplicaId>, StorageResponse)>,
}

impl ReplicaTask {
    /// Runs the replica task.
    async fn run(self) {
        let replica_id = self.replica_id;
        info!(%replica_id, "starting replica task");

        let client = self.connect().await;
        match self.run_message_loop(client).await {
            Ok(()) => info!(%replica_id, "stopped replica task"),
            Err(error) => warn!(%replica_id, %error, "replica task failed"),
        }
    }

    /// Connects to the replica.
    ///
    /// The connection is retried forever (with backoff) and this method returns only after
    /// a connection was successfully established.
    async fn connect(&self) -> ReplicaClient {
        let try_connect = async move |retry: RetryState| {
            let version = self.config.build_info.semver_version();
            let client_params = &self.config.grpc_client;

            let connect_start = Instant::now();
            let connect_timeout = client_params.connect_timeout.unwrap_or(Duration::MAX);
            let keepalive_timeout = client_params
                .http2_keep_alive_timeout
                .unwrap_or(Duration::MAX);

            let connect_result = StorageCtpClient::connect_partitioned(
                self.config.location.ctl_addrs.clone(),
                version,
                connect_timeout,
                keepalive_timeout,
                self.metrics.clone(),
            )
            .await;

            self.metrics.observe_connect_time(connect_start.elapsed());

            connect_result.inspect_err(|error| {
                let next_backoff = retry.next_backoff.unwrap();
                if retry.i >= mz_service::retry::INFO_MIN_RETRIES {
                    info!(
                        replica_id = %self.replica_id, ?next_backoff,
                        "error connecting to replica: {error:#}",
                    );
                } else {
                    debug!(
                        replica_id = %self.replica_id, ?next_backoff,
                        "error connecting to replica: {error:#}",
                    );
                }
            })
        };

        let client = Retry::default()
            .clamp_backoff(Duration::from_secs(1))
            .retry_async(try_connect)
            .await
            .expect("retries forever");

        self.metrics.observe_connect();
        self.connected.store(true, atomic::Ordering::Relaxed);

        client
    }

    /// Runs the message loop.
    ///
    /// Returns (with an `Err`) if it encounters an error condition (e.g. the replica disconnects).
    /// If no error condition is encountered, the task runs until the controller disconnects from
    /// the command channel, or the task is dropped.
    async fn run_message_loop(mut self, mut client: ReplicaClient) -> Result<(), anyhow::Error> {
        loop {
            select! {
                // Command from controller to forward to replica.
                // `tokio::sync::mpsc::UnboundedReceiver::recv` is documented as cancel safe.
                command = self.command_rx.recv() => {
                    let Some(mut command) = command else {
                        tracing::debug!(%self.replica_id, "controller is no longer interested in this replica, shutting down message loop");
                        break;
                    };

                    self.specialize_command(&mut command);
                    client.send(command).await?;
                },
                // Response from replica to forward to controller.
                // `GenericClient::recv` implementations are required to be cancel safe.
                response = client.recv() => {
                    let Some(response) = response? else {
                        bail!("replica unexpectedly gracefully terminated connection");
                    };

                    if self.response_tx.send((Some(self.replica_id), response)).is_err() {
                        tracing::debug!(%self.replica_id, "controller (receiver) is no longer interested in this replica, shutting down message loop");
                        break;
                    }
                }
            }
        }

        Ok(())
    }

    /// Specialize a command for the given replica configuration.
    ///
    /// Most [`StorageCommand`]s are independent of the target replica, but some contain
    /// replica-specific fields that must be adjusted before sending.
    fn specialize_command(&self, command: &mut StorageCommand) {
        if let StorageCommand::Hello { nonce } = command {
            *nonce = Uuid::new_v4();
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use mz_dyncfg::{ConfigUpdates, ConfigVal};
    use mz_storage_types::dyncfgs::ENABLE_UPSERT_PAGED_SPILL;
    use mz_storage_types::parameters::StorageParameters;

    use super::{Instance, ReplicaId, StorageCommand};

    fn update_configuration_command() -> StorageCommand {
        StorageCommand::UpdateConfiguration(Box::new(StorageParameters::default()))
    }

    fn dyncfg_updates(command: &StorageCommand) -> &ConfigUpdates {
        match command {
            StorageCommand::UpdateConfiguration(params) => &params.dyncfg_updates,
            other => panic!("expected UpdateConfiguration, got {other:?}"),
        }
    }

    /// An overridden replica's `UpdateConfiguration` carries its override on
    /// top of the environment-wide updates, while replicas without an
    /// override receive the base command unchanged.
    #[mz_ore::test]
    fn update_configuration_applies_replica_override() {
        let mut over = ConfigUpdates::default();
        over.add(&ENABLE_UPSERT_PAGED_SPILL, true);
        let overrides = BTreeMap::from([(ReplicaId::User(1), over)]);

        let command = Instance::specialize_command_for_replica(
            update_configuration_command(),
            ReplicaId::User(1),
            &overrides,
        );
        assert_eq!(
            dyncfg_updates(&command)
                .updates
                .get(ENABLE_UPSERT_PAGED_SPILL.name()),
            Some(&ConfigVal::Bool(true)),
        );

        let command = Instance::specialize_command_for_replica(
            update_configuration_command(),
            ReplicaId::User(2),
            &overrides,
        );
        assert_eq!(
            dyncfg_updates(&command)
                .updates
                .get(ENABLE_UPSERT_PAGED_SPILL.name()),
            None,
        );
    }
}