Skip to main content

mz_compute_client/controller/
instance.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! A controller for a compute instance.
11
12use std::collections::{BTreeMap, BTreeSet};
13use std::fmt::Debug;
14use std::sync::{Arc, Mutex};
15use std::time::{Duration, Instant};
16
17use chrono::{DateTime, DurationRound, TimeDelta, Utc};
18use mz_build_info::BuildInfo;
19use mz_cluster_client::WallclockLagFn;
20use mz_compute_types::dataflows::{BuildDesc, DataflowDescription};
21use mz_compute_types::plan::render_plan::RenderPlan;
22use mz_compute_types::sinks::{
23    ComputeSinkConnection, ComputeSinkDesc, MaterializedViewSinkConnection,
24};
25use mz_compute_types::sources::SourceInstanceDesc;
26use mz_controller_types::dyncfgs::{
27    ENABLE_PAUSED_CLUSTER_READHOLD_DOWNGRADE, WALLCLOCK_LAG_RECORDING_INTERVAL,
28};
29use mz_dyncfg::{ConfigSet, ConfigUpdates};
30use mz_expr::RowSetFinishing;
31use mz_ore::cast::CastFrom;
32use mz_ore::channel::instrumented_unbounded_channel;
33use mz_ore::now::NowFn;
34use mz_ore::tracing::OpenTelemetryContext;
35use mz_ore::{soft_assert_or_log, soft_panic_or_log};
36use mz_persist_types::PersistLocation;
37use mz_repr::adt::timestamp::CheckedTimestamp;
38use mz_repr::refresh_schedule::RefreshSchedule;
39use mz_repr::{Datum, Diff, GlobalId, RelationDesc, Row, Timestamp};
40use mz_storage_client::controller::{IntrospectionType, WallclockLag, WallclockLagHistogramPeriod};
41use mz_storage_types::read_holds::{self, ReadHold};
42use mz_storage_types::read_policy::ReadPolicy;
43use thiserror::Error;
44use timely::PartialOrder;
45use timely::progress::frontier::MutableAntichain;
46use timely::progress::{Antichain, ChangeBatch};
47use tokio::sync::{mpsc, oneshot};
48use uuid::Uuid;
49
50use crate::controller::error::{
51    CollectionMissing, ERROR_TARGET_REPLICA_FAILED, HydrationCheckBadTarget,
52};
53use crate::controller::instance_client::PeekError;
54use crate::controller::replica::{ReplicaClient, ReplicaConfig};
55use crate::controller::{
56    ComputeControllerResponse, IntrospectionUpdates, PeekNotification, ReplicaId,
57    StorageCollections,
58};
59use crate::logging::LogVariant;
60use crate::metrics::IntCounter;
61use crate::metrics::{InstanceMetrics, ReplicaCollectionMetrics, ReplicaMetrics, UIntGauge};
62use crate::protocol::command::{
63    ComputeCommand, ComputeParameters, InstanceConfig, Peek, PeekTarget,
64};
65use crate::protocol::history::ComputeCommandHistory;
66use crate::protocol::response::{
67    ComputeResponse, CopyToResponse, FrontiersResponse, PeekResponse, StatusResponse,
68    SubscribeBatch, SubscribeResponse,
69};
70
71#[derive(Error, Debug)]
72#[error("replica exists already: {0}")]
73pub(super) struct ReplicaExists(pub ReplicaId);
74
75#[derive(Error, Debug)]
76#[error("replica does not exist: {0}")]
77pub(super) struct ReplicaMissing(pub ReplicaId);
78
79#[derive(Error, Debug)]
80pub(super) enum DataflowCreationError {
81    #[error("collection does not exist: {0}")]
82    CollectionMissing(GlobalId),
83    #[error("replica does not exist: {0}")]
84    ReplicaMissing(ReplicaId),
85    #[error("dataflow definition lacks an as_of value")]
86    MissingAsOf,
87    #[error("subscribe dataflow has an empty as_of")]
88    EmptyAsOfForSubscribe,
89    #[error("copy to dataflow has an empty as_of")]
90    EmptyAsOfForCopyTo,
91    #[error("no read hold provided for dataflow import: {0}")]
92    ReadHoldMissing(GlobalId),
93    #[error("insufficient read hold provided for dataflow import: {0}")]
94    ReadHoldInsufficient(GlobalId),
95}
96
97impl From<CollectionMissing> for DataflowCreationError {
98    fn from(error: CollectionMissing) -> Self {
99        Self::CollectionMissing(error.0)
100    }
101}
102
103#[derive(Error, Debug)]
104pub(super) enum ReadPolicyError {
105    #[error("collection does not exist: {0}")]
106    CollectionMissing(GlobalId),
107    #[error("collection is write-only: {0}")]
108    WriteOnlyCollection(GlobalId),
109}
110
111impl From<CollectionMissing> for ReadPolicyError {
112    fn from(error: CollectionMissing) -> Self {
113        Self::CollectionMissing(error.0)
114    }
115}
116
117/// A command sent to an [`Instance`] task.
118pub(super) type Command = Box<dyn FnOnce(&mut Instance) + Send>;
119
120/// A response from a replica, composed of a replica ID, the replica's current epoch, and the
121/// compute response itself.
122pub(super) type ReplicaResponse = (ReplicaId, u64, ComputeResponse);
123
124/// The state we keep for a compute instance.
125pub(super) struct Instance {
126    /// Build info for spawning replicas
127    build_info: &'static BuildInfo,
128    /// A handle providing access to storage collections.
129    storage_collections: StorageCollections,
130    /// Whether instance initialization has been completed.
131    initialized: bool,
132    /// Whether this instance is in read-only mode.
133    ///
134    /// When in read-only mode, this instance will not update persistent state, such as
135    /// wallclock lag introspection.
136    read_only: bool,
137    /// The workload class of this instance.
138    ///
139    /// This is currently only used to annotate metrics.
140    workload_class: Option<String>,
141    /// The replicas of this compute instance.
142    replicas: BTreeMap<ReplicaId, ReplicaState>,
143    /// Per-replica dyncfg overrides, merged into the `UpdateConfiguration`
144    /// command sent to each replica (and into the command-history replay used
145    /// to hydrate new replicas). Populated from the scoped feature flags
146    /// (replica-local) layer; empty by default, in which case every replica
147    /// receives the unmodified environment-wide configuration. Stores only the
148    /// values that differ from the environment-wide value, so the map is sparse.
149    replica_dyncfg_overrides: BTreeMap<ReplicaId, ConfigUpdates>,
150    /// Currently installed compute collections.
151    ///
152    /// New entries are added for all collections exported from dataflows created through
153    /// [`Instance::create_dataflow`].
154    ///
155    /// Entries are removed by [`Instance::cleanup_collections`]. See that method's documentation
156    /// about the conditions for removing collection state.
157    collections: BTreeMap<GlobalId, CollectionState>,
158    /// IDs of log sources maintained by this compute instance.
159    log_sources: BTreeMap<LogVariant, GlobalId>,
160    /// Currently outstanding peeks.
161    ///
162    /// New entries are added for all peeks initiated through [`Instance::peek`].
163    ///
164    /// The entry for a peek is only removed once all replicas have responded to the peek. This is
165    /// currently required to ensure all replicas have stopped reading from the peeked collection's
166    /// inputs before we allow them to compact. database-issues#4822 tracks changing this so we only have to wait
167    /// for the first peek response.
168    peeks: BTreeMap<Uuid, PendingPeek>,
169    /// Currently in-progress subscribes.
170    ///
171    /// New entries are added for all subscribes exported from dataflows created through
172    /// [`Instance::create_dataflow`].
173    ///
174    /// The entry for a subscribe is removed once at least one replica has reported the subscribe
175    /// to have advanced to the empty frontier or to have been dropped, implying that no further
176    /// updates will be emitted for this subscribe.
177    ///
178    /// Note that subscribes are tracked both in `collections` and `subscribes`. `collections`
179    /// keeps track of the subscribe's upper and since frontiers and ensures appropriate read holds
180    /// on the subscribe's input. `subscribes` is only used to track which updates have been
181    /// emitted, to decide if new ones should be emitted or suppressed.
182    subscribes: BTreeMap<GlobalId, ActiveSubscribe>,
183    /// Tracks all in-progress COPY TOs.
184    ///
185    /// New entries are added for all s3 oneshot sinks (corresponding to a COPY TO) exported from
186    /// dataflows created through [`Instance::create_dataflow`].
187    ///
188    /// The entry for a copy to is removed once at least one replica has finished
189    /// or the exporting collection is dropped.
190    copy_tos: BTreeSet<GlobalId>,
191    /// The command history, used when introducing new replicas or restarting existing replicas.
192    history: ComputeCommandHistory<UIntGauge>,
193    /// Receiver for commands to be executed.
194    command_rx: mpsc::UnboundedReceiver<Command>,
195    /// Sender for responses to be delivered.
196    response_tx: mpsc::UnboundedSender<ComputeControllerResponse>,
197    /// Sender for introspection updates to be recorded.
198    introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
199    /// The registry the controller uses to report metrics.
200    metrics: InstanceMetrics,
201    /// Dynamic system configuration.
202    dyncfg: Arc<ConfigSet>,
203
204    /// The persist location where we can stash large peek results.
205    peek_stash_persist_location: PersistLocation,
206
207    /// A function that produces the current wallclock time.
208    now: NowFn,
209    /// A function that computes the lag between the given time and wallclock time.
210    wallclock_lag: WallclockLagFn<Timestamp>,
211    /// The last time wallclock lag introspection was recorded.
212    wallclock_lag_last_recorded: DateTime<Utc>,
213
214    /// Sender for updates to collection read holds.
215    ///
216    /// Copies of this sender are given to [`ReadHold`]s that are created in
217    /// [`CollectionState::new`].
218    read_hold_tx: read_holds::ChangeTx,
219    /// A sender for responses from replicas.
220    replica_tx: mz_ore::channel::InstrumentedUnboundedSender<ReplicaResponse, IntCounter>,
221    /// A receiver for responses from replicas.
222    replica_rx: mz_ore::channel::InstrumentedUnboundedReceiver<ReplicaResponse, IntCounter>,
223}
224
225impl Instance {
226    /// Acquire a handle to the collection state associated with `id`.
227    fn collection(&self, id: GlobalId) -> Result<&CollectionState, CollectionMissing> {
228        self.collections.get(&id).ok_or(CollectionMissing(id))
229    }
230
231    /// Acquire a mutable handle to the collection state associated with `id`.
232    fn collection_mut(&mut self, id: GlobalId) -> Result<&mut CollectionState, CollectionMissing> {
233        self.collections.get_mut(&id).ok_or(CollectionMissing(id))
234    }
235
236    /// Acquire a handle to the collection state associated with `id`.
237    ///
238    /// # Panics
239    ///
240    /// Panics if the identified collection does not exist.
241    fn expect_collection(&self, id: GlobalId) -> &CollectionState {
242        self.collections.get(&id).expect("collection must exist")
243    }
244
245    /// Acquire a mutable handle to the collection state associated with `id`.
246    ///
247    /// # Panics
248    ///
249    /// Panics if the identified collection does not exist.
250    fn expect_collection_mut(&mut self, id: GlobalId) -> &mut CollectionState {
251        self.collections
252            .get_mut(&id)
253            .expect("collection must exist")
254    }
255
256    fn collections_iter(&self) -> impl Iterator<Item = (GlobalId, &CollectionState)> {
257        self.collections.iter().map(|(id, coll)| (*id, coll))
258    }
259
260    /// Returns an iterator over replicas that host the given collection.
261    ///
262    /// For replica-targeted collections, this returns only the target replica.
263    /// For non-targeted collections, this returns all replicas.
264    ///
265    /// Returns `Err` if the collection does not exist.
266    fn replicas_hosting(
267        &self,
268        id: GlobalId,
269    ) -> Result<impl Iterator<Item = &ReplicaState>, CollectionMissing> {
270        let target = self.collection(id)?.target_replica;
271        Ok(self
272            .replicas
273            .values()
274            .filter(move |r| target.map_or(true, |t| t == r.id)))
275    }
276
277    /// Add a collection to the instance state.
278    ///
279    /// # Panics
280    ///
281    /// Panics if a collection with the same ID exists already.
282    fn add_collection(
283        &mut self,
284        id: GlobalId,
285        as_of: Antichain<Timestamp>,
286        shared: SharedCollectionState,
287        storage_dependencies: BTreeMap<GlobalId, ReadHold>,
288        compute_dependencies: BTreeMap<GlobalId, ReadHold>,
289        replica_input_read_holds: Vec<ReadHold>,
290        write_only: bool,
291        storage_sink: bool,
292        initial_as_of: Option<Antichain<Timestamp>>,
293        refresh_schedule: Option<RefreshSchedule>,
294        target_replica: Option<ReplicaId>,
295    ) {
296        // Add global collection state.
297        let dependency_ids: Vec<GlobalId> = compute_dependencies
298            .keys()
299            .chain(storage_dependencies.keys())
300            .copied()
301            .collect();
302        let introspection = CollectionIntrospection::new(
303            id,
304            self.introspection_tx.clone(),
305            as_of.clone(),
306            storage_sink,
307            initial_as_of,
308            refresh_schedule,
309            dependency_ids,
310        );
311        let mut state = CollectionState::new(
312            id,
313            as_of.clone(),
314            shared,
315            storage_dependencies,
316            compute_dependencies,
317            Arc::clone(&self.read_hold_tx),
318            introspection,
319        );
320        state.target_replica = target_replica;
321        // If the collection is write-only, clear its read policy to reflect that.
322        if write_only {
323            state.read_policy = None;
324        }
325
326        if let Some(previous) = self.collections.insert(id, state) {
327            panic!("attempt to add a collection with existing ID {id} (previous={previous:?}");
328        }
329
330        // Add per-replica collection state.
331        for replica in self.replicas.values_mut() {
332            if target_replica.is_some_and(|id| id != replica.id) {
333                continue;
334            }
335            replica.add_collection(id, as_of.clone(), replica_input_read_holds.clone());
336        }
337    }
338
339    fn remove_collection(&mut self, id: GlobalId) {
340        // Remove per-replica collection state.
341        for replica in self.replicas.values_mut() {
342            replica.remove_collection(id);
343        }
344
345        // Remove global collection state.
346        self.collections.remove(&id);
347    }
348
349    fn add_replica_state(
350        &mut self,
351        id: ReplicaId,
352        client: ReplicaClient,
353        config: ReplicaConfig,
354        epoch: u64,
355    ) -> Result<(), read_holds::ReadHoldIssuerHungUp> {
356        let log_ids: BTreeSet<_> = config.logging.index_logs.values().copied().collect();
357
358        let metrics = self.metrics.for_replica(id);
359        let mut replica = ReplicaState::new(
360            id,
361            client,
362            config,
363            metrics,
364            self.introspection_tx.clone(),
365            epoch,
366        );
367
368        // Add per-replica collection state.
369        let mut shutdown_input = None;
370        for (collection_id, collection) in &self.collections {
371            // Skip log collections not maintained by this replica,
372            // and collections targeted at a different replica.
373            if (collection.log_collection && !log_ids.contains(collection_id))
374                || collection.target_replica.is_some_and(|rid| rid != id)
375            {
376                continue;
377            }
378
379            let as_of = if collection.log_collection {
380                // For log collections, we don't send a `CreateDataflow` command to the replica, so
381                // it doesn't know which as-of the controler chose and defaults to the minimum
382                // frontier instead. We need to initialize the controller-side tracking with the
383                // same frontier, to avoid observing regressions in the reported frontiers.
384                Antichain::from_elem(Timestamp::MIN)
385            } else {
386                collection.read_frontier().to_owned()
387            };
388
389            // Cloning a `ReadHold` fails when its issuer has hung up. For these holds the issuer
390            // is the `StorageCollections`, which doesn't hang up as long as the `Instance` exists,
391            // except during process shutdown, when the tokio runtime drops tasks in arbitrary
392            // order. In that case there is no way of correctly initializing the per-replica
393            // collection state, so we give up. We still add the replica itself, to keep the
394            // bookkeeping consistent with the controller's, and then signal the unrecoverable
395            // error to the caller, which shuts the instance down.
396            let mut input_read_holds = Vec::with_capacity(collection.storage_dependencies.len());
397            let mut hung_up = Vec::new();
398            for hold in collection.storage_dependencies.values() {
399                match hold.try_clone() {
400                    Ok(hold) => input_read_holds.push(hold),
401                    Err(read_holds::ReadHoldIssuerHungUp(input_id)) => hung_up.push(input_id),
402                }
403            }
404            if !hung_up.is_empty() {
405                tracing::error!(
406                    replica_id = %id,
407                    %collection_id,
408                    ?hung_up,
409                    "giving up on adding replica collections: storage read hold issuers hung \
410                     up, the process is shutting down",
411                );
412                shutdown_input = hung_up.into_iter().next();
413                break;
414            }
415
416            replica.add_collection(*collection_id, as_of, input_read_holds);
417        }
418
419        self.replicas.insert(id, replica);
420
421        match shutdown_input {
422            Some(input_id) => Err(read_holds::ReadHoldIssuerHungUp(input_id)),
423            None => Ok(()),
424        }
425    }
426
427    /// Enqueue the given response for delivery to the controller clients.
428    fn deliver_response(&self, response: ComputeControllerResponse) {
429        // Failure to send means the `ComputeController` has been dropped and doesn't care about
430        // responses anymore.
431        let _ = self.response_tx.send(response);
432    }
433
434    /// Enqueue the given introspection updates for recording.
435    fn deliver_introspection_updates(&self, type_: IntrospectionType, updates: Vec<(Row, Diff)>) {
436        // Failure to send means the `ComputeController` has been dropped and doesn't care about
437        // introspection updates anymore.
438        let _ = self.introspection_tx.send((type_, updates));
439    }
440
441    /// Returns whether the identified replica exists.
442    fn replica_exists(&self, id: ReplicaId) -> bool {
443        self.replicas.contains_key(&id)
444    }
445
446    /// Return the IDs of pending peeks targeting the specified replica.
447    fn peeks_targeting(&self, replica_id: ReplicaId) -> impl Iterator<Item = (Uuid, &PendingPeek)> {
448        self.peeks.iter().filter_map(move |(uuid, peek)| {
449            if peek.target_replica == Some(replica_id) {
450                Some((*uuid, peek))
451            } else {
452                None
453            }
454        })
455    }
456
457    /// Return the IDs of in-progress subscribes targeting the specified replica.
458    fn subscribes_targeting(&self, replica_id: ReplicaId) -> impl Iterator<Item = GlobalId> + '_ {
459        self.subscribes.keys().copied().filter(move |id| {
460            let collection = self.expect_collection(*id);
461            collection.target_replica == Some(replica_id)
462        })
463    }
464
465    /// Update introspection with the current collection frontiers.
466    ///
467    /// We could also do this directly in response to frontier changes, but doing it periodically
468    /// lets us avoid emitting some introspection updates that can be consolidated (e.g. a write
469    /// frontier updated immediately followed by a read frontier update).
470    ///
471    /// This method is invoked by `ComputeController::maintain`, which we expect to be called once
472    /// per second during normal operation.
473    fn update_frontier_introspection(&mut self) {
474        for collection in self.collections.values_mut() {
475            collection
476                .introspection
477                .observe_frontiers(&collection.read_frontier(), &collection.write_frontier());
478        }
479
480        for replica in self.replicas.values_mut() {
481            for collection in replica.collections.values_mut() {
482                collection
483                    .introspection
484                    .observe_frontier(&collection.write_frontier);
485            }
486        }
487    }
488
489    /// Refresh the controller state metrics for this instance.
490    ///
491    /// We could also do state metric updates directly in response to state changes, but that would
492    /// mean littering the code with metric update calls. Encapsulating state metric maintenance in
493    /// a single method is less noisy.
494    ///
495    /// This method is invoked by `ComputeController::maintain`, which we expect to be called once
496    /// per second during normal operation.
497    fn refresh_state_metrics(&self) {
498        let unscheduled_collections_count =
499            self.collections.values().filter(|c| !c.scheduled).count();
500        let connected_replica_count = self
501            .replicas
502            .values()
503            .filter(|r| r.client.is_connected())
504            .count();
505
506        self.metrics
507            .replica_count
508            .set(u64::cast_from(self.replicas.len()));
509        self.metrics
510            .collection_count
511            .set(u64::cast_from(self.collections.len()));
512        self.metrics
513            .collection_unscheduled_count
514            .set(u64::cast_from(unscheduled_collections_count));
515        self.metrics
516            .peek_count
517            .set(u64::cast_from(self.peeks.len()));
518        self.metrics
519            .subscribe_count
520            .set(u64::cast_from(self.subscribes.len()));
521        self.metrics
522            .copy_to_count
523            .set(u64::cast_from(self.copy_tos.len()));
524        self.metrics
525            .connected_replica_count
526            .set(u64::cast_from(connected_replica_count));
527    }
528
529    /// Refresh the wallclock lag introspection and metrics with the current lag values.
530    ///
531    /// This method produces wallclock lag metrics of two different shapes:
532    ///
533    /// * Histories: For each replica and each collection, we measure the lag of the write frontier
534    ///   behind the wallclock time every second. Every minute we emit the maximum lag observed
535    ///   over the last minute, together with the current time.
536    /// * Histograms: For each collection, we measure the lag of the write frontier behind
537    ///   wallclock time every second. Every minute we emit all lags observed over the last minute,
538    ///   together with the current histogram period.
539    ///
540    /// Histories are emitted to both Mz introspection and Prometheus, histograms only to
541    /// introspection. We treat lags of unreadable collections (i.e. collections that contain no
542    /// readable times) as undefined and set them to NULL in introspection and `u64::MAX` in
543    /// Prometheus.
544    ///
545    /// This method is invoked by `ComputeController::maintain`, which we expect to be called once
546    /// per second during normal operation.
547    fn refresh_wallclock_lag(&mut self) {
548        let frontier_lag = |frontier: &Antichain<Timestamp>| match frontier.as_option() {
549            Some(ts) => (self.wallclock_lag)(ts.clone()),
550            None => Duration::ZERO,
551        };
552
553        let now_ms = (self.now)();
554        let histogram_period = WallclockLagHistogramPeriod::from_epoch_millis(now_ms, &self.dyncfg);
555        let histogram_labels = match &self.workload_class {
556            Some(wc) => [("workload_class", wc.clone())].into(),
557            None => BTreeMap::new(),
558        };
559
560        // For collections that sink into storage, we need to ask the storage controller to know
561        // whether they're currently readable.
562        let readable_storage_collections: BTreeSet<_> = self
563            .collections
564            .keys()
565            .filter_map(|id| {
566                let frontiers = self.storage_collections.collection_frontiers(*id).ok()?;
567                PartialOrder::less_than(&frontiers.read_capabilities, &frontiers.write_frontier)
568                    .then_some(*id)
569            })
570            .collect();
571
572        // First, iterate over all collections and collect histogram measurements.
573        for (id, collection) in &mut self.collections {
574            let write_frontier = collection.write_frontier();
575            let readable = if self.storage_collections.check_exists(*id).is_ok() {
576                readable_storage_collections.contains(id)
577            } else {
578                PartialOrder::less_than(&collection.read_frontier(), &write_frontier)
579            };
580
581            if let Some(stash) = &mut collection.wallclock_lag_histogram_stash {
582                let bucket = if readable {
583                    let lag = frontier_lag(&write_frontier);
584                    let lag = lag.as_secs().next_power_of_two();
585                    WallclockLag::Seconds(lag)
586                } else {
587                    WallclockLag::Undefined
588                };
589
590                let key = (histogram_period, bucket, histogram_labels.clone());
591                *stash.entry(key).or_default() += Diff::ONE;
592            }
593        }
594
595        // Second, iterate over all per-replica collections and collect history measurements.
596        for replica in self.replicas.values_mut() {
597            for (id, collection) in &mut replica.collections {
598                // A per-replica collection is considered readable in the context of lag
599                // measurement if either:
600                //  (a) it sinks into a storage collection that is readable
601                //  (b) it is hydrated
602                let readable = readable_storage_collections.contains(id) || collection.hydrated();
603
604                let lag = if readable {
605                    let lag = frontier_lag(&collection.write_frontier);
606                    WallclockLag::Seconds(lag.as_secs())
607                } else {
608                    WallclockLag::Undefined
609                };
610
611                if let Some(wallclock_lag_max) = &mut collection.wallclock_lag_max {
612                    *wallclock_lag_max = (*wallclock_lag_max).max(lag);
613                }
614
615                if let Some(metrics) = &mut collection.metrics {
616                    // No way to specify values as undefined in Prometheus metrics, so we use the
617                    // maximum value instead.
618                    let secs = lag.unwrap_seconds_or(u64::MAX);
619                    metrics.wallclock_lag.observe(secs);
620                };
621            }
622        }
623
624        // Record lags to persist, if it's time.
625        self.maybe_record_wallclock_lag();
626    }
627
628    /// Produce new wallclock lag introspection updates, provided enough time has passed since the
629    /// last recording.
630    //
631    /// We emit new introspection updates if the system time has passed into a new multiple of the
632    /// recording interval (typically 1 minute) since the last refresh. The storage controller uses
633    /// the same approach, ensuring that both controllers commit their lags at roughly the same
634    /// time, avoiding confusion caused by inconsistencies.
635    fn maybe_record_wallclock_lag(&mut self) {
636        if self.read_only {
637            return;
638        }
639
640        let duration_trunc = |datetime: DateTime<_>, interval| {
641            let td = TimeDelta::from_std(interval).ok()?;
642            datetime.duration_trunc(td).ok()
643        };
644
645        let interval = WALLCLOCK_LAG_RECORDING_INTERVAL.get(&self.dyncfg);
646        let now_dt = mz_ore::now::to_datetime((self.now)());
647        let now_trunc = duration_trunc(now_dt, interval).unwrap_or_else(|| {
648            soft_panic_or_log!("excessive wallclock lag recording interval: {interval:?}");
649            let default = WALLCLOCK_LAG_RECORDING_INTERVAL.default();
650            duration_trunc(now_dt, *default).unwrap()
651        });
652        if now_trunc <= self.wallclock_lag_last_recorded {
653            return;
654        }
655
656        let now_ts: CheckedTimestamp<_> = now_trunc.try_into().expect("must fit");
657
658        let mut history_updates = Vec::new();
659        for (replica_id, replica) in &mut self.replicas {
660            for (collection_id, collection) in &mut replica.collections {
661                let Some(wallclock_lag_max) = &mut collection.wallclock_lag_max else {
662                    continue;
663                };
664
665                let max_lag = std::mem::replace(wallclock_lag_max, WallclockLag::MIN);
666                let row = Row::pack_slice(&[
667                    Datum::String(&collection_id.to_string()),
668                    Datum::String(&replica_id.to_string()),
669                    max_lag.into_interval_datum(),
670                    Datum::TimestampTz(now_ts),
671                ]);
672                history_updates.push((row, Diff::ONE));
673            }
674        }
675        if !history_updates.is_empty() {
676            self.deliver_introspection_updates(
677                IntrospectionType::WallclockLagHistory,
678                history_updates,
679            );
680        }
681
682        let mut histogram_updates = Vec::new();
683        let mut row_buf = Row::default();
684        for (collection_id, collection) in &mut self.collections {
685            let Some(stash) = &mut collection.wallclock_lag_histogram_stash else {
686                continue;
687            };
688
689            for ((period, lag, labels), count) in std::mem::take(stash) {
690                let mut packer = row_buf.packer();
691                packer.extend([
692                    Datum::TimestampTz(period.start),
693                    Datum::TimestampTz(period.end),
694                    Datum::String(&collection_id.to_string()),
695                    lag.into_uint64_datum(),
696                ]);
697                let labels = labels.iter().map(|(k, v)| (*k, Datum::String(v)));
698                packer.push_dict(labels);
699
700                histogram_updates.push((row_buf.clone(), count));
701            }
702        }
703        if !histogram_updates.is_empty() {
704            self.deliver_introspection_updates(
705                IntrospectionType::WallclockLagHistogram,
706                histogram_updates,
707            );
708        }
709
710        self.wallclock_lag_last_recorded = now_trunc;
711    }
712
713    /// Returns `true` if the given collection is hydrated on at least one
714    /// replica.
715    ///
716    /// This also returns `true` in case this cluster does not have any
717    /// replicas that host the given collection.
718    #[mz_ore::instrument(level = "debug")]
719    pub fn collection_hydrated(&self, collection_id: GlobalId) -> Result<bool, CollectionMissing> {
720        let mut hosting_replicas = self.replicas_hosting(collection_id)?.peekable();
721        if hosting_replicas.peek().is_none() {
722            return Ok(true);
723        }
724        for replica_state in hosting_replicas {
725            let collection_state = replica_state
726                .collections
727                .get(&collection_id)
728                .expect("hosting replica must have per-replica collection state");
729
730            if collection_state.hydrated() {
731                return Ok(true);
732            }
733        }
734
735        Ok(false)
736    }
737
738    /// Returns `true` if each non-transient, non-excluded collection is hydrated on at
739    /// least one replica.
740    ///
741    /// This also returns `true` in case this cluster does not have any
742    /// replicas.
743    #[mz_ore::instrument(level = "debug")]
744    pub fn collections_hydrated_on_replicas(
745        &self,
746        target_replica_ids: Option<Vec<ReplicaId>>,
747        exclude_collections: &BTreeSet<GlobalId>,
748    ) -> Result<bool, HydrationCheckBadTarget> {
749        if self.replicas.is_empty() {
750            return Ok(true);
751        }
752        let target_replicas: BTreeSet<ReplicaId> = self
753            .replicas
754            .keys()
755            .filter_map(|id| match target_replica_ids {
756                None => Some(id.clone()),
757                Some(ref ids) if ids.contains(id) => Some(id.clone()),
758                Some(_) => None,
759            })
760            .collect();
761        if let Some(targets) = target_replica_ids {
762            if target_replicas.is_empty() {
763                return Err(HydrationCheckBadTarget(targets));
764            }
765        }
766
767        let mut unhydrated = BTreeSet::new();
768        for (id, _collection) in self.collections_iter() {
769            if id.is_transient() || exclude_collections.contains(&id) {
770                continue;
771            }
772
773            let mut collection_hydrated = false;
774            // `replicas_hosting` cannot fail here because `collections_iter`
775            // only yields collections that exist.
776            for replica_state in self.replicas_hosting(id).expect("collection must exist") {
777                if !target_replicas.contains(&replica_state.id) {
778                    continue;
779                }
780                let collection_state = replica_state
781                    .collections
782                    .get(&id)
783                    .expect("hosting replica must have per-replica collection state");
784
785                if collection_state.hydrated() {
786                    collection_hydrated = true;
787                    break;
788                }
789            }
790
791            if !collection_hydrated {
792                // We collect all non-hydrated collections instead of breaking
793                // out early, so that the log below names every collection the
794                // caller is waiting on.
795                unhydrated.insert(id);
796            }
797        }
798
799        if !unhydrated.is_empty() {
800            // Callers poll this on the cluster controller's reconcile tick,
801            // which tests turn down to milliseconds, so this is deliberately
802            // one line per call rather than one per collection.
803            tracing::info!(
804                replicas = ?target_replicas,
805                collections = ?unhydrated,
806                "collections are not hydrated on any target replica",
807            );
808        }
809
810        Ok(unhydrated.is_empty())
811    }
812
813    /// Clean up collection state that is not needed anymore.
814    ///
815    /// Three conditions need to be true before we can remove state for a collection:
816    ///
817    ///  1. A client must have explicitly dropped the collection. If that is not the case, clients
818    ///     can still reasonably assume that the controller knows about the collection and can
819    ///     answer queries about it.
820    ///  2. There must be no outstanding read capabilities on the collection. As long as someone
821    ///     still holds read capabilities on a collection, we need to keep it around to be able
822    ///     to properly handle downgrading of said capabilities.
823    ///  3. All replica frontiers for the collection must have advanced to the empty frontier.
824    ///     Advancement to the empty frontiers signals that replicas are done computing the
825    ///     collection and that they won't send more `ComputeResponse`s for it. As long as we might
826    ///     receive responses for a collection we want to keep it around to be able to validate and
827    ///     handle these responses.
828    fn cleanup_collections(&mut self) {
829        let to_remove: Vec<_> = self
830            .collections_iter()
831            .filter(|(id, collection)| {
832                collection.dropped
833                    && collection.shared.lock_read_capabilities(|c| c.is_empty())
834                    && self
835                        .replicas
836                        .values()
837                        .all(|r| r.collection_frontiers_empty(*id))
838            })
839            .map(|(id, _collection)| id)
840            .collect();
841
842        for id in to_remove {
843            self.remove_collection(id);
844        }
845    }
846
847    /// Returns the state of the [`Instance`] formatted as JSON.
848    ///
849    /// The returned value is not guaranteed to be stable and may change at any point in time.
850    #[mz_ore::instrument(level = "debug")]
851    pub fn dump(&self) -> Result<serde_json::Value, anyhow::Error> {
852        // Note: We purposefully use the `Debug` formatting for the value of all fields in the
853        // returned object as a tradeoff between usability and stability. `serde_json` will fail
854        // to serialize an object if the keys aren't strings, so `Debug` formatting the values
855        // prevents a future unrelated change from silently breaking this method.
856
857        // Destructure `self` here so we don't forget to consider dumping newly added fields.
858        let Self {
859            build_info: _,
860            storage_collections: _,
861            peek_stash_persist_location: _,
862            initialized,
863            read_only,
864            workload_class,
865            replicas,
866            replica_dyncfg_overrides: _,
867            collections,
868            log_sources: _,
869            peeks,
870            subscribes,
871            copy_tos,
872            history: _,
873            command_rx: _,
874            response_tx: _,
875            introspection_tx: _,
876            metrics: _,
877            dyncfg: _,
878            now: _,
879            wallclock_lag: _,
880            wallclock_lag_last_recorded,
881            read_hold_tx: _,
882            replica_tx: _,
883            replica_rx: _,
884        } = self;
885
886        let replicas: BTreeMap<_, _> = replicas
887            .iter()
888            .map(|(id, replica)| Ok((id.to_string(), replica.dump()?)))
889            .collect::<Result<_, anyhow::Error>>()?;
890        let collections: BTreeMap<_, _> = collections
891            .iter()
892            .map(|(id, collection)| (id.to_string(), format!("{collection:?}")))
893            .collect();
894        let peeks: BTreeMap<_, _> = peeks
895            .iter()
896            .map(|(uuid, peek)| (uuid.to_string(), format!("{peek:?}")))
897            .collect();
898        let subscribes: BTreeMap<_, _> = subscribes
899            .iter()
900            .map(|(id, subscribe)| (id.to_string(), format!("{subscribe:?}")))
901            .collect();
902        let copy_tos: Vec<_> = copy_tos.iter().map(|id| id.to_string()).collect();
903        let wallclock_lag_last_recorded = format!("{wallclock_lag_last_recorded:?}");
904
905        Ok(serde_json::json!({
906            "initialized": initialized,
907            "read_only": read_only,
908            "workload_class": workload_class,
909            "replicas": replicas,
910            "collections": collections,
911            "peeks": peeks,
912            "subscribes": subscribes,
913            "copy_tos": copy_tos,
914            "wallclock_lag_last_recorded": wallclock_lag_last_recorded,
915        }))
916    }
917
918    /// Reports the current write frontier for the identified compute collection.
919    pub(super) fn collection_write_frontier(
920        &self,
921        id: GlobalId,
922    ) -> Result<Antichain<Timestamp>, CollectionMissing> {
923        Ok(self.collection(id)?.write_frontier())
924    }
925}
926
927impl Instance {
928    pub(super) fn new(
929        build_info: &'static BuildInfo,
930        storage: StorageCollections,
931        peek_stash_persist_location: PersistLocation,
932        arranged_logs: Vec<(LogVariant, GlobalId, SharedCollectionState)>,
933        metrics: InstanceMetrics,
934        now: NowFn,
935        wallclock_lag: WallclockLagFn<Timestamp>,
936        dyncfg: Arc<ConfigSet>,
937        command_rx: mpsc::UnboundedReceiver<Command>,
938        response_tx: mpsc::UnboundedSender<ComputeControllerResponse>,
939        read_hold_tx: read_holds::ChangeTx,
940        introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
941        read_only: bool,
942    ) -> Self {
943        let mut collections = BTreeMap::new();
944        let mut log_sources = BTreeMap::new();
945        for (log, id, shared) in arranged_logs {
946            let collection = CollectionState::new_log_collection(
947                id,
948                shared,
949                Arc::clone(&read_hold_tx),
950                introspection_tx.clone(),
951            );
952            collections.insert(id, collection);
953            log_sources.insert(log, id);
954        }
955
956        let history = ComputeCommandHistory::new(metrics.for_history());
957
958        let send_count = metrics.response_send_count.clone();
959        let recv_count = metrics.response_recv_count.clone();
960        let (replica_tx, replica_rx) = instrumented_unbounded_channel(send_count, recv_count);
961
962        let now_dt = mz_ore::now::to_datetime(now());
963
964        Self {
965            build_info,
966            storage_collections: storage,
967            peek_stash_persist_location,
968            initialized: false,
969            read_only,
970            workload_class: None,
971            replicas: Default::default(),
972            replica_dyncfg_overrides: Default::default(),
973            collections,
974            log_sources,
975            peeks: Default::default(),
976            subscribes: Default::default(),
977            copy_tos: Default::default(),
978            history,
979            command_rx,
980            response_tx,
981            introspection_tx,
982            metrics,
983            dyncfg,
984            now,
985            wallclock_lag,
986            wallclock_lag_last_recorded: now_dt,
987            read_hold_tx,
988            replica_tx,
989            replica_rx,
990        }
991    }
992
993    pub(super) async fn run(mut self) {
994        self.send(ComputeCommand::Hello {
995            // The nonce is protocol iteration-specific and will be set in
996            // `ReplicaTask::specialize_command`.
997            nonce: Uuid::default(),
998        });
999
1000        let instance_config = InstanceConfig {
1001            peek_stash_persist_location: self.peek_stash_persist_location.clone(),
1002            // The remaining fields are replica-specific and will be set in
1003            // `ReplicaTask::specialize_command` (logging, expiration, dictionary compression) and
1004            // `Instance::specialize_command_for_replica` (the initial config snapshot).
1005            logging: Default::default(),
1006            expiration_offset: Default::default(),
1007            arrangement_dictionary_compression: Default::default(),
1008            initial_config: Default::default(),
1009        };
1010
1011        self.send(ComputeCommand::CreateInstance(Box::new(instance_config)));
1012
1013        loop {
1014            tokio::select! {
1015                command = self.command_rx.recv() => match command {
1016                    Some(cmd) => cmd(&mut self),
1017                    None => break,
1018                },
1019                response = self.replica_rx.recv() => match response {
1020                    Some(response) => self.handle_response(response),
1021                    None => unreachable!("self owns a sender side of the channel"),
1022                }
1023            }
1024        }
1025    }
1026
1027    /// Update instance configuration.
1028    #[mz_ore::instrument(level = "debug")]
1029    pub fn update_configuration(&mut self, config_params: ComputeParameters) {
1030        if let Some(workload_class) = &config_params.workload_class {
1031            self.workload_class = workload_class.clone();
1032        }
1033
1034        let command = ComputeCommand::UpdateConfiguration(Box::new(config_params));
1035        self.send(command);
1036    }
1037
1038    /// Marks the end of any initialization commands.
1039    ///
1040    /// Intended to be called by `Controller`, rather than by other code.
1041    /// Calling this method repeatedly has no effect.
1042    #[mz_ore::instrument(level = "debug")]
1043    pub fn initialization_complete(&mut self) {
1044        // The compute protocol requires that `InitializationComplete` is sent only once.
1045        if !self.initialized {
1046            self.send(ComputeCommand::InitializationComplete);
1047            self.initialized = true;
1048        }
1049    }
1050
1051    /// Allows collections to affect writes to external systems (persist).
1052    ///
1053    /// Calling this method repeatedly has no effect.
1054    #[mz_ore::instrument(level = "debug")]
1055    pub fn allow_writes(&mut self, collection_id: GlobalId) -> Result<(), CollectionMissing> {
1056        let collection = self.collection_mut(collection_id)?;
1057
1058        // Do not send redundant allow-writes commands.
1059        if !collection.read_only {
1060            return Ok(());
1061        }
1062
1063        // Don't send allow-writes for collections that are not installed.
1064        let as_of = collection.read_frontier();
1065
1066        // If the collection has an empty `as_of`, it was either never installed on the replica or
1067        // has since been dropped. In either case the replica does not expect any commands for it.
1068        if as_of.is_empty() {
1069            return Ok(());
1070        }
1071
1072        collection.read_only = false;
1073        self.send(ComputeCommand::AllowWrites(collection_id));
1074
1075        Ok(())
1076    }
1077
1078    /// Shut down this instance.
1079    ///
1080    /// This method asserts that the instance has no replicas left. It exists to help
1081    /// us find bugs where the client drops a compute instance that still has replicas
1082    /// installed, and later assumes that said replicas still exist.
1083    ///
1084    /// # Panics
1085    ///
1086    /// Soft-panics if the compute instance still has active replicas.
1087    #[mz_ore::instrument(level = "debug")]
1088    pub fn shutdown(&mut self) {
1089        // Taking the `command_rx` ensures that the [`Instance::run`] loop terminates.
1090        let (_tx, rx) = mpsc::unbounded_channel();
1091        self.command_rx = rx;
1092
1093        let stray_replicas: Vec<_> = self.replicas.keys().collect();
1094        soft_assert_or_log!(
1095            stray_replicas.is_empty(),
1096            "dropped instance still has provisioned replicas: {stray_replicas:?}",
1097        );
1098    }
1099
1100    /// Terminate the [`Instance::run`] loop, causing the instance task to shut down.
1101    ///
1102    /// Unlike [`Instance::shutdown`], this does not assert that the instance has no replicas
1103    /// left. We use it to react to unrecoverable errors that can only occur during process
1104    /// shutdown, such as a storage read hold issuer hanging up while we rehydrate a replica.
1105    fn initiate_shutdown(&mut self) {
1106        // Replacing `command_rx` with a fresh, sender-less channel makes the next `recv` in
1107        // [`Instance::run`] return `None`, terminating the loop.
1108        let (_tx, rx) = mpsc::unbounded_channel();
1109        self.command_rx = rx;
1110    }
1111
1112    /// Sends a command to replicas of this instance.
1113    #[mz_ore::instrument(level = "debug")]
1114    fn send(&mut self, cmd: ComputeCommand) {
1115        // Record the command so that new replicas can be brought up to speed.
1116        // We record the *base* (un-specialized) command, so that the per-replica
1117        // dyncfg overrides are re-applied at replay time in `add_replica` rather
1118        // than baked into the shared history.
1119        self.history.push(cmd.clone());
1120
1121        let target_replica = self.target_replica(&cmd);
1122
1123        // Borrow the overrides and dyncfg separately from `self.replicas` so the per-replica
1124        // specialization below does not conflict with the mutable replica borrow.
1125        let overrides = &self.replica_dyncfg_overrides;
1126        let dyncfg = &self.dyncfg;
1127
1128        if let Some(rid) = target_replica {
1129            if let Some(replica) = self.replicas.get_mut(&rid) {
1130                let cmd = Self::specialize_command_for_replica(cmd, rid, overrides, dyncfg);
1131                let _ = replica.client.send(cmd);
1132            }
1133        } else {
1134            for (rid, replica) in self.replicas.iter_mut() {
1135                let cmd =
1136                    Self::specialize_command_for_replica(cmd.clone(), *rid, overrides, dyncfg);
1137                let _ = replica.client.send(cmd);
1138            }
1139        }
1140    }
1141
1142    /// Specializes a command for a specific replica by merging that replica's dyncfg override into
1143    /// its configuration. For `UpdateConfiguration` the override is merged into the update. For
1144    /// `CreateInstance` the current dyncfg, with the override applied on top, is captured as the
1145    /// initial config snapshot. All other commands are returned unchanged.
1146    ///
1147    /// The snapshot is built here, rather than baked into the history, so it reflects the dyncfg
1148    /// and override values current at the time the command is sent or replayed to the replica.
1149    fn specialize_command_for_replica(
1150        mut cmd: ComputeCommand,
1151        replica_id: ReplicaId,
1152        overrides: &BTreeMap<ReplicaId, ConfigUpdates>,
1153        dyncfg: &ConfigSet,
1154    ) -> ComputeCommand {
1155        let over = overrides.get(&replica_id);
1156        match &mut cmd {
1157            ComputeCommand::UpdateConfiguration(params) => {
1158                if let Some(over) = over
1159                    && !over.updates.is_empty()
1160                {
1161                    params.dyncfg_updates.extend(over.clone());
1162                }
1163            }
1164            ComputeCommand::CreateInstance(config) => {
1165                let mut initial = ConfigUpdates::from(dyncfg);
1166                if let Some(over) = over {
1167                    initial.extend(over.clone());
1168                }
1169                config.initial_config = initial;
1170            }
1171            _ => {}
1172        }
1173        cmd
1174    }
1175
1176    /// Replaces the per-replica dyncfg overrides. Callers should follow this
1177    /// with a configuration push (e.g. `update_configuration`) so that existing
1178    /// replicas observe the new overrides.
1179    pub(super) fn update_replica_dyncfg_overrides(
1180        &mut self,
1181        overrides: BTreeMap<ReplicaId, ConfigUpdates>,
1182    ) {
1183        self.replica_dyncfg_overrides = overrides;
1184    }
1185
1186    /// Determine the target replica for a compute command. Retrieves the
1187    /// collection named by the command, and returns the target replica if
1188    /// it is set, and None if not set, or the command doesn't name a collection.
1189    ///
1190    /// Panics if a create-dataflow command names collections that have different
1191    /// target replicas. It is an error to construct such an object and would
1192    /// indicate a bug in [`Self::create_dataflow`].
1193    fn target_replica(&self, cmd: &ComputeCommand) -> Option<ReplicaId> {
1194        match &cmd {
1195            ComputeCommand::Schedule(id)
1196            | ComputeCommand::AllowWrites(id)
1197            | ComputeCommand::AllowCompaction { id, .. } => {
1198                self.expect_collection(*id).target_replica
1199            }
1200            ComputeCommand::CreateDataflow(desc) => {
1201                let mut target_replica = None;
1202                for id in desc.export_ids() {
1203                    if let Some(replica) = self.expect_collection(id).target_replica {
1204                        if target_replica.is_some() {
1205                            assert_eq!(target_replica, Some(replica));
1206                        }
1207                        target_replica = Some(replica);
1208                    }
1209                }
1210                target_replica
1211            }
1212            // Skip Peek as we don't allow replica-targeted indexes.
1213            ComputeCommand::Peek(_)
1214            | ComputeCommand::Hello { .. }
1215            | ComputeCommand::CreateInstance(_)
1216            | ComputeCommand::InitializationComplete
1217            | ComputeCommand::UpdateConfiguration(_)
1218            | ComputeCommand::CancelPeek { .. } => None,
1219        }
1220    }
1221
1222    /// Add a new instance replica, by ID.
1223    #[mz_ore::instrument(level = "debug")]
1224    pub fn add_replica(
1225        &mut self,
1226        id: ReplicaId,
1227        mut config: ReplicaConfig,
1228        epoch: Option<u64>,
1229    ) -> Result<(), ReplicaExists> {
1230        if self.replica_exists(id) {
1231            return Err(ReplicaExists(id));
1232        }
1233
1234        config.logging.index_logs = self.log_sources.clone();
1235
1236        let epoch = epoch.unwrap_or(1);
1237        let metrics = self.metrics.for_replica(id);
1238        let client = ReplicaClient::spawn(
1239            id,
1240            self.build_info,
1241            config.clone(),
1242            epoch,
1243            metrics.clone(),
1244            Arc::clone(&self.dyncfg),
1245            self.replica_tx.clone(),
1246        );
1247
1248        // Take this opportunity to clean up the history we should present.
1249        self.history.reduce();
1250
1251        // Advance the uppers of source imports
1252        self.history.update_source_uppers(&self.storage_collections);
1253
1254        // Replay the commands at the client, creating new dataflow identifiers.
1255        for command in self.history.iter() {
1256            // Skip `CreateDataflow` commands targeted at different replicas.
1257            if let Some(target_replica) = self.target_replica(command)
1258                && target_replica != id
1259            {
1260                continue;
1261            }
1262
1263            // Re-apply this replica's dyncfg override to replayed config commands, and rebuild the
1264            // create-instance snapshot from the current dyncfg.
1265            let command = Self::specialize_command_for_replica(
1266                command.clone(),
1267                id,
1268                &self.replica_dyncfg_overrides,
1269                &self.dyncfg,
1270            );
1271            if client.send(command).is_err() {
1272                // We swallow the error here. On the next send, we will fail again, and
1273                // restart the connection as well as this rehydration.
1274                tracing::warn!("Replica {:?} connection terminated during hydration", id);
1275                break;
1276            }
1277        }
1278
1279        // Add replica to tracked state.
1280        if self.add_replica_state(id, client, config, epoch).is_err() {
1281            // A storage read hold issuer hung up, which only happens during process shutdown.
1282            // There is no way to correctly bring up the replica anymore, so we shut the instance
1283            // down instead of running on with half-initialized replica state. `add_replica_state`
1284            // has already logged the details and inserted the replica to keep our bookkeeping
1285            // consistent with the controller's.
1286            self.initiate_shutdown();
1287        }
1288
1289        Ok(())
1290    }
1291
1292    /// Remove an existing instance replica, by ID.
1293    #[mz_ore::instrument(level = "debug")]
1294    pub fn remove_replica(&mut self, id: ReplicaId) -> Result<(), ReplicaMissing> {
1295        let replica = self.replicas.remove(&id).ok_or(ReplicaMissing(id))?;
1296
1297        // The coordinator only re-pushes the override map when the scoped configuration itself
1298        // changes, so a dropped replica's entry would otherwise be retained until the next such
1299        // change.
1300        self.replica_dyncfg_overrides.remove(&id);
1301
1302        // Before dropping the replica state (and the contained input read holds), log read holds
1303        // that are the last line of defense against compaction of a dataflow's storage inputs. If
1304        // the corresponding global read hold has already been released, dropping the per-replica
1305        // read hold will allow compaction, which can cause the replica to panic trying to install
1306        // the dataflow.
1307        //
1308        // This exists primarily to help diagnose incidents-and-escalations#39.
1309        for (collection_id, replica_collection) in &replica.collections {
1310            let collection = self.collections.get(collection_id);
1311            for replica_hold in &replica_collection.input_read_holds {
1312                let input_id = replica_hold.id();
1313                let global_hold = collection.and_then(|c| c.storage_dependencies.get(&input_id));
1314                let unprotected = global_hold
1315                    .is_none_or(|h| PartialOrder::less_than(replica_hold.since(), h.since()));
1316                if unprotected {
1317                    tracing::warn!(
1318                        replica_id = %id,
1319                        %collection_id,
1320                        %input_id,
1321                        replica_hold_since = ?replica_hold.since(),
1322                        global_hold_since = ?global_hold.map(|h| h.since()),
1323                        "dropping per-replica read hold without equivalent global read hold",
1324                    );
1325                }
1326            }
1327        }
1328        drop(replica);
1329
1330        // Subscribes targeting this replica either won't be served anymore (if the replica is
1331        // dropped) or might produce inconsistent output (if the target collection is an
1332        // introspection index). We produce an error to inform upstream.
1333        let to_drop: Vec<_> = self.subscribes_targeting(id).collect();
1334        for subscribe_id in to_drop {
1335            let subscribe = self.subscribes.remove(&subscribe_id).unwrap();
1336            let response = ComputeControllerResponse::SubscribeResponse(
1337                subscribe_id,
1338                SubscribeBatch {
1339                    lower: subscribe.frontier.clone(),
1340                    upper: subscribe.frontier,
1341                    updates: Err(ERROR_TARGET_REPLICA_FAILED.into()),
1342                },
1343            );
1344            self.deliver_response(response);
1345        }
1346
1347        // Peeks targeting this replica might not be served anymore (if the replica is dropped).
1348        // If the replica has failed it might come back and respond to the peek later, but it still
1349        // seems like a good idea to cancel the peek to inform the caller about the failure. This
1350        // is consistent with how we handle targeted subscribes above.
1351        let mut peek_responses = Vec::new();
1352        let mut to_drop = Vec::new();
1353        for (uuid, peek) in self.peeks_targeting(id) {
1354            peek_responses.push(ComputeControllerResponse::PeekNotification(
1355                uuid,
1356                PeekNotification::Error(ERROR_TARGET_REPLICA_FAILED.into()),
1357                peek.otel_ctx.clone(),
1358            ));
1359            to_drop.push(uuid);
1360        }
1361        for response in peek_responses {
1362            self.deliver_response(response);
1363        }
1364        for uuid in to_drop {
1365            let response = PeekResponse::Error(ERROR_TARGET_REPLICA_FAILED.into());
1366            self.finish_peek(uuid, response);
1367        }
1368
1369        // We might have a chance to forward implied capabilities and reduce the cost of bringing
1370        // up the next replica, if the dropped replica was the only one in the cluster.
1371        self.forward_implied_capabilities();
1372
1373        Ok(())
1374    }
1375
1376    /// Rehydrate the given instance replica.
1377    ///
1378    /// # Panics
1379    ///
1380    /// Panics if the specified replica does not exist.
1381    fn rehydrate_replica(&mut self, id: ReplicaId) {
1382        let config = self.replicas[&id].config.clone();
1383        let epoch = self.replicas[&id].epoch + 1;
1384
1385        self.remove_replica(id).expect("replica must exist");
1386        let result = self.add_replica(id, config, Some(epoch));
1387
1388        match result {
1389            Ok(()) => (),
1390            Err(ReplicaExists(_)) => unreachable!("replica was removed"),
1391        }
1392    }
1393
1394    /// Rehydrate any failed replicas of this instance.
1395    fn rehydrate_failed_replicas(&mut self) {
1396        let replicas = self.replicas.iter();
1397        let failed_replicas: Vec<_> = replicas
1398            .filter_map(|(id, replica)| replica.client.is_failed().then_some(*id))
1399            .collect();
1400
1401        for replica_id in failed_replicas {
1402            self.rehydrate_replica(replica_id);
1403        }
1404    }
1405
1406    /// Creates the described dataflow and initializes state for its output.
1407    ///
1408    /// This method expects a `DataflowDescription` with an `as_of` frontier specified, as well as
1409    /// for each imported collection a read hold in `import_read_holds` at at least the `as_of`.
1410    #[mz_ore::instrument(level = "debug")]
1411    pub fn create_dataflow(
1412        &mut self,
1413        dataflow: DataflowDescription<mz_compute_types::plan::LirRelationExpr, ()>,
1414        import_read_holds: Vec<ReadHold>,
1415        mut shared_collection_state: BTreeMap<GlobalId, SharedCollectionState>,
1416        target_replica: Option<ReplicaId>,
1417    ) -> Result<(), DataflowCreationError> {
1418        use DataflowCreationError::*;
1419
1420        // Validate that the target replica, if specified, exists.
1421        // A targeted dataflow is only installed on a single replica; if that
1422        // replica doesn't exist, we can't create the dataflow.
1423        if let Some(replica_id) = target_replica {
1424            if !self.replica_exists(replica_id) {
1425                return Err(ReplicaMissing(replica_id));
1426            }
1427        }
1428
1429        // Simple sanity checks around `as_of`
1430        let as_of = dataflow.as_of.as_ref().ok_or(MissingAsOf)?;
1431        if as_of.is_empty() && dataflow.subscribe_ids().next().is_some() {
1432            return Err(EmptyAsOfForSubscribe);
1433        }
1434        if as_of.is_empty() && dataflow.copy_to_ids().next().is_some() {
1435            return Err(EmptyAsOfForCopyTo);
1436        }
1437
1438        // Collect all dependencies of the dataflow, and read holds on them at the `as_of`.
1439        let mut storage_dependencies = BTreeMap::new();
1440        let mut compute_dependencies = BTreeMap::new();
1441
1442        // When we install per-replica input read holds, we cannot use the `as_of` because of
1443        // reconciliation: Existing slow replicas might be reading from the inputs at times before
1444        // the `as_of` and we would rather not crash them by allowing their inputs to compact too
1445        // far. So instead we take read holds at the least time available.
1446        let mut replica_input_read_holds = Vec::new();
1447
1448        let mut import_read_holds: BTreeMap<_, _> =
1449            import_read_holds.into_iter().map(|r| (r.id(), r)).collect();
1450
1451        for &id in dataflow.source_imports.keys() {
1452            let mut read_hold = import_read_holds.remove(&id).ok_or(ReadHoldMissing(id))?;
1453            replica_input_read_holds.push(read_hold.clone());
1454
1455            read_hold
1456                .try_downgrade(as_of.clone())
1457                .map_err(|_| ReadHoldInsufficient(id))?;
1458            storage_dependencies.insert(id, read_hold);
1459        }
1460
1461        for &id in dataflow.index_imports.keys() {
1462            let mut read_hold = import_read_holds.remove(&id).ok_or(ReadHoldMissing(id))?;
1463            read_hold
1464                .try_downgrade(as_of.clone())
1465                .map_err(|_| ReadHoldInsufficient(id))?;
1466            compute_dependencies.insert(id, read_hold);
1467        }
1468
1469        // If the `as_of` is empty, we are not going to create a dataflow, so replicas won't read
1470        // from the inputs.
1471        if as_of.is_empty() {
1472            replica_input_read_holds = Default::default();
1473        }
1474
1475        // Install collection state for each of the exports.
1476        for export_id in dataflow.export_ids() {
1477            let shared = shared_collection_state
1478                .remove(&export_id)
1479                .unwrap_or_else(|| SharedCollectionState::new(as_of.clone()));
1480            let write_only = dataflow.sink_exports.contains_key(&export_id);
1481            let storage_sink = dataflow.persist_sink_ids().any(|id| id == export_id);
1482
1483            self.add_collection(
1484                export_id,
1485                as_of.clone(),
1486                shared,
1487                storage_dependencies.clone(),
1488                compute_dependencies.clone(),
1489                replica_input_read_holds.clone(),
1490                write_only,
1491                storage_sink,
1492                dataflow.initial_storage_as_of.clone(),
1493                dataflow.refresh_schedule.clone(),
1494                target_replica,
1495            );
1496
1497            // If the export is a storage sink, we can advance its write frontier to the write
1498            // frontier of the target storage collection.
1499            if let Ok(frontiers) = self.storage_collections.collection_frontiers(export_id) {
1500                self.maybe_update_global_write_frontier(export_id, frontiers.write_frontier);
1501            }
1502        }
1503
1504        // Initialize tracking of subscribes.
1505        for subscribe_id in dataflow.subscribe_ids() {
1506            self.subscribes
1507                .insert(subscribe_id, ActiveSubscribe::default());
1508        }
1509
1510        // Initialize tracking of copy tos.
1511        for copy_to_id in dataflow.copy_to_ids() {
1512            self.copy_tos.insert(copy_to_id);
1513        }
1514
1515        // Here we augment all imported sources and all exported sinks with the appropriate
1516        // storage metadata needed by the compute instance.
1517        let mut source_imports = BTreeMap::new();
1518        for (id, import) in dataflow.source_imports {
1519            let frontiers = self
1520                .storage_collections
1521                .collection_frontiers(id)
1522                .expect("collection exists");
1523
1524            let collection_metadata = self
1525                .storage_collections
1526                .collection_metadata(id)
1527                .expect("we have a read hold on this collection");
1528
1529            let desc = SourceInstanceDesc {
1530                storage_metadata: collection_metadata.clone(),
1531                arguments: import.desc.arguments,
1532                typ: import.desc.typ.clone(),
1533            };
1534            source_imports.insert(
1535                id,
1536                mz_compute_types::dataflows::SourceImport {
1537                    desc,
1538                    monotonic: import.monotonic,
1539                    with_snapshot: import.with_snapshot,
1540                    upper: frontiers.write_frontier,
1541                },
1542            );
1543        }
1544
1545        let mut sink_exports = BTreeMap::new();
1546        for (id, se) in dataflow.sink_exports {
1547            let connection = match se.connection {
1548                ComputeSinkConnection::MaterializedView(conn) => {
1549                    let metadata = self
1550                        .storage_collections
1551                        .collection_metadata(id)
1552                        .map_err(|_| CollectionMissing(id))?
1553                        .clone();
1554                    let conn = MaterializedViewSinkConnection {
1555                        value_desc: conn.value_desc,
1556                        storage_metadata: metadata,
1557                    };
1558                    ComputeSinkConnection::MaterializedView(conn)
1559                }
1560                ComputeSinkConnection::Subscribe(conn) => ComputeSinkConnection::Subscribe(conn),
1561                ComputeSinkConnection::CopyToS3Oneshot(conn) => {
1562                    ComputeSinkConnection::CopyToS3Oneshot(conn)
1563                }
1564                ComputeSinkConnection::MetricSink(conn) => ComputeSinkConnection::MetricSink(conn),
1565            };
1566            let desc = ComputeSinkDesc {
1567                from: se.from,
1568                from_desc: se.from_desc,
1569                connection,
1570                with_snapshot: se.with_snapshot,
1571                up_to: se.up_to,
1572                non_null_assertions: se.non_null_assertions,
1573                refresh_schedule: se.refresh_schedule,
1574            };
1575            sink_exports.insert(id, desc);
1576        }
1577
1578        // Flatten the dataflow plans into the representation expected by replicas.
1579        let objects_to_build = dataflow
1580            .objects_to_build
1581            .into_iter()
1582            .map(|object| BuildDesc {
1583                id: object.id,
1584                plan: RenderPlan::try_from(object.plan).expect("valid plan"),
1585            })
1586            .collect();
1587
1588        let augmented_dataflow = DataflowDescription {
1589            source_imports,
1590            sink_exports,
1591            objects_to_build,
1592            // The rest of the fields are identical
1593            index_imports: dataflow.index_imports,
1594            index_exports: dataflow.index_exports,
1595            as_of: dataflow.as_of.clone(),
1596            until: dataflow.until,
1597            initial_storage_as_of: dataflow.initial_storage_as_of,
1598            refresh_schedule: dataflow.refresh_schedule,
1599            debug_name: dataflow.debug_name,
1600            time_dependence: dataflow.time_dependence,
1601        };
1602
1603        if augmented_dataflow.is_transient() {
1604            tracing::debug!(
1605                name = %augmented_dataflow.debug_name,
1606                import_ids = %augmented_dataflow.display_import_ids(),
1607                export_ids = %augmented_dataflow.display_export_ids(),
1608                as_of = ?augmented_dataflow.as_of.as_ref().unwrap().elements(),
1609                until = ?augmented_dataflow.until.elements(),
1610                "creating dataflow",
1611            );
1612        } else {
1613            tracing::info!(
1614                name = %augmented_dataflow.debug_name,
1615                import_ids = %augmented_dataflow.display_import_ids(),
1616                export_ids = %augmented_dataflow.display_export_ids(),
1617                as_of = ?augmented_dataflow.as_of.as_ref().unwrap().elements(),
1618                until = ?augmented_dataflow.until.elements(),
1619                "creating dataflow",
1620            );
1621        }
1622
1623        // Skip the actual dataflow creation for an empty `as_of`. (Happens e.g. for the
1624        // bootstrapping of a REFRESH AT mat view that is past its last refresh.)
1625        if as_of.is_empty() {
1626            tracing::info!(
1627                name = %augmented_dataflow.debug_name,
1628                "not sending `CreateDataflow`, because of empty `as_of`",
1629            );
1630        } else {
1631            let collections: Vec<_> = augmented_dataflow.export_ids().collect();
1632            self.send(ComputeCommand::CreateDataflow(Box::new(augmented_dataflow)));
1633
1634            for id in collections {
1635                self.maybe_schedule_collection(id);
1636            }
1637        }
1638
1639        Ok(())
1640    }
1641
1642    /// Schedule the identified collection if all its inputs are available.
1643    ///
1644    /// # Panics
1645    ///
1646    /// Panics if the identified collection does not exist.
1647    fn maybe_schedule_collection(&mut self, id: GlobalId) {
1648        let collection = self.expect_collection(id);
1649
1650        // Don't schedule collections twice.
1651        if collection.scheduled {
1652            return;
1653        }
1654
1655        let as_of = collection.read_frontier();
1656
1657        // If the collection has an empty `as_of`, it was either never installed on the replica or
1658        // has since been dropped. In either case the replica does not expect any commands for it.
1659        if as_of.is_empty() {
1660            return;
1661        }
1662
1663        let ready = if id.is_transient() {
1664            // Always schedule transient collections immediately. The assumption is that those are
1665            // created by interactive user commands and we want to schedule them as quickly as
1666            // possible. Inputs might not yet be available, but when they become available, we
1667            // don't need to wait for the controller to become aware and for the scheduling check
1668            // to run again.
1669            true
1670        } else {
1671            // Ignore self-dependencies. Any self-dependencies do not need to be
1672            // available at the as_of for the dataflow to make progress, so we
1673            // can ignore them here. At the moment, only continual tasks have
1674            // self-dependencies, but this logic is correct for any dataflow, so
1675            // we don't special case it to CTs.
1676            let not_self_dep = |x: &GlobalId| *x != id;
1677
1678            // Make sure we never schedule a collection before its input compute collections have
1679            // been scheduled. Scheduling in the wrong order can lead to deadlocks.
1680            let mut deps_scheduled = true;
1681
1682            // Check dependency frontiers to determine if all inputs are
1683            // available. An input is available when its frontier is greater
1684            // than the `as_of`, i.e., all input data up to and including the
1685            // `as_of` has been sealed.
1686            let compute_deps = collection.compute_dependency_ids().filter(not_self_dep);
1687            let mut compute_frontiers = Vec::new();
1688            for id in compute_deps {
1689                let dep = &self.expect_collection(id);
1690                deps_scheduled &= dep.scheduled;
1691                compute_frontiers.push(dep.write_frontier());
1692            }
1693
1694            let storage_deps = collection.storage_dependency_ids().filter(not_self_dep);
1695            let storage_frontiers = self
1696                .storage_collections
1697                .collections_frontiers(storage_deps.collect())
1698                .expect("must exist");
1699            let storage_frontiers = storage_frontiers.into_iter().map(|f| f.write_frontier);
1700
1701            let mut frontiers = compute_frontiers.into_iter().chain(storage_frontiers);
1702            let frontiers_ready =
1703                frontiers.all(|frontier| PartialOrder::less_than(&as_of, &frontier));
1704
1705            deps_scheduled && frontiers_ready
1706        };
1707
1708        if ready {
1709            self.send(ComputeCommand::Schedule(id));
1710            let collection = self.expect_collection_mut(id);
1711            collection.scheduled = true;
1712        }
1713    }
1714
1715    /// Schedule any unscheduled collections that are ready.
1716    fn schedule_collections(&mut self) {
1717        let ids: Vec<_> = self.collections.keys().copied().collect();
1718        for id in ids {
1719            self.maybe_schedule_collection(id);
1720        }
1721    }
1722
1723    /// Drops the read capability for the given collections and allows their resources to be
1724    /// reclaimed.
1725    #[mz_ore::instrument(level = "debug")]
1726    pub fn drop_collections(&mut self, ids: Vec<GlobalId>) -> Result<(), CollectionMissing> {
1727        for id in &ids {
1728            let collection = self.collection_mut(*id)?;
1729
1730            // Mark the collection as dropped to allow it to be removed from the controller state.
1731            collection.dropped = true;
1732
1733            // Drop the implied and warmup read holds to announce that clients are not
1734            // interested in the collection anymore.
1735            collection.implied_read_hold.release();
1736            collection.warmup_read_hold.release();
1737
1738            // If the collection is a subscribe, stop tracking it. This ensures that the controller
1739            // ceases to produce `SubscribeResponse`s for this subscribe.
1740            self.subscribes.remove(id);
1741            // If the collection is a copy to, stop tracking it. This ensures that the controller
1742            // ceases to produce `CopyToResponse`s` for this copy to.
1743            self.copy_tos.remove(id);
1744        }
1745
1746        Ok(())
1747    }
1748
1749    /// Initiate a peek request for the contents of `id` at `timestamp`.
1750    ///
1751    /// If this returns an error, then it didn't modify any `Instance` state.
1752    #[mz_ore::instrument(level = "debug")]
1753    pub fn peek(
1754        &mut self,
1755        peek_target: PeekTarget,
1756        literal_constraints: Option<Vec<Row>>,
1757        uuid: Uuid,
1758        timestamp: Timestamp,
1759        result_desc: RelationDesc,
1760        finishing: RowSetFinishing,
1761        map_filter_project: mz_expr::SafeMfpPlan,
1762        mut read_hold: ReadHold,
1763        target_replica: Option<ReplicaId>,
1764        peek_response_tx: oneshot::Sender<PeekResponse>,
1765    ) -> Result<(), PeekError> {
1766        use PeekError::*;
1767
1768        let target_id = peek_target.id();
1769
1770        // Downgrade the provided read hold to the peek time.
1771        if read_hold.id() != target_id {
1772            return Err(ReadHoldIdMismatch(read_hold.id()));
1773        }
1774        read_hold
1775            .try_downgrade(Antichain::from_elem(timestamp.clone()))
1776            .map_err(|_| ReadHoldInsufficient(target_id))?;
1777
1778        if let Some(target) = target_replica {
1779            if !self.replica_exists(target) {
1780                return Err(ReplicaMissing(target));
1781            }
1782        }
1783
1784        let otel_ctx = OpenTelemetryContext::obtain();
1785
1786        self.peeks.insert(
1787            uuid,
1788            PendingPeek {
1789                target_replica,
1790                // TODO(guswynn): can we just hold the `tracing::Span` here instead?
1791                otel_ctx: otel_ctx.clone(),
1792                requested_at: Instant::now(),
1793                read_hold,
1794                peek_response_tx,
1795                limit: finishing.limit.map(usize::cast_from),
1796                offset: finishing.offset,
1797            },
1798        );
1799
1800        let peek = Peek {
1801            literal_constraints,
1802            uuid,
1803            timestamp,
1804            finishing,
1805            map_filter_project,
1806            // Obtain an `OpenTelemetryContext` from the thread-local tracing
1807            // tree to forward it on to the compute worker.
1808            otel_ctx,
1809            target: peek_target,
1810            result_desc,
1811        };
1812        self.send(ComputeCommand::Peek(Box::new(peek)));
1813
1814        Ok(())
1815    }
1816
1817    /// Cancels an existing peek request.
1818    #[mz_ore::instrument(level = "debug")]
1819    pub fn cancel_peek(&mut self, uuid: Uuid, reason: PeekResponse) {
1820        let Some(peek) = self.peeks.get_mut(&uuid) else {
1821            tracing::warn!("did not find pending peek for {uuid}");
1822            return;
1823        };
1824
1825        let duration = peek.requested_at.elapsed();
1826        self.metrics
1827            .observe_peek_response(&PeekResponse::Canceled, duration);
1828
1829        // Enqueue a notification for the cancellation.
1830        let otel_ctx = peek.otel_ctx.clone();
1831        otel_ctx.attach_as_parent();
1832
1833        self.deliver_response(ComputeControllerResponse::PeekNotification(
1834            uuid,
1835            PeekNotification::Canceled,
1836            otel_ctx,
1837        ));
1838
1839        // Finish the peek.
1840        // This will also propagate the cancellation to the replicas.
1841        self.finish_peek(uuid, reason);
1842    }
1843
1844    /// Assigns a read policy to specific identifiers.
1845    ///
1846    /// The policies are assigned in the order presented, and repeated identifiers should
1847    /// conclude with the last policy. Changing a policy will immediately downgrade the read
1848    /// capability if appropriate, but it will not "recover" the read capability if the prior
1849    /// capability is already ahead of it.
1850    ///
1851    /// Identifiers not present in `policies` retain their existing read policies.
1852    ///
1853    /// It is an error to attempt to set a read policy for a collection that is not readable in the
1854    /// context of compute. At this time, only indexes are readable compute collections.
1855    #[mz_ore::instrument(level = "debug")]
1856    pub fn set_read_policy(
1857        &mut self,
1858        policies: Vec<(GlobalId, ReadPolicy)>,
1859    ) -> Result<(), ReadPolicyError> {
1860        // Do error checking upfront, to avoid introducing inconsistencies between a collection's
1861        // `implied_capability` and `read_capabilities`.
1862        for (id, _policy) in &policies {
1863            let collection = self.collection(*id)?;
1864            if collection.read_policy.is_none() {
1865                return Err(ReadPolicyError::WriteOnlyCollection(*id));
1866            }
1867        }
1868
1869        for (id, new_policy) in policies {
1870            let collection = self.expect_collection_mut(id);
1871            let new_since = new_policy.frontier(collection.write_frontier().borrow());
1872            let _ = collection.implied_read_hold.try_downgrade(new_since);
1873            collection.read_policy = Some(new_policy);
1874        }
1875
1876        Ok(())
1877    }
1878
1879    /// Advance the global write frontier of the given collection.
1880    ///
1881    /// Frontier regressions are gracefully ignored.
1882    ///
1883    /// # Panics
1884    ///
1885    /// Panics if the identified collection does not exist.
1886    #[mz_ore::instrument(level = "debug")]
1887    fn maybe_update_global_write_frontier(
1888        &mut self,
1889        id: GlobalId,
1890        new_frontier: Antichain<Timestamp>,
1891    ) {
1892        let collection = self.expect_collection_mut(id);
1893
1894        let advanced = collection.shared.lock_write_frontier(|f| {
1895            let advanced = PartialOrder::less_than(f, &new_frontier);
1896            if advanced {
1897                f.clone_from(&new_frontier);
1898            }
1899            advanced
1900        });
1901
1902        if !advanced {
1903            return;
1904        }
1905
1906        // Relax the implied read hold according to the read policy.
1907        let new_since = match &collection.read_policy {
1908            Some(read_policy) => {
1909                // For readable collections the read frontier is determined by applying the
1910                // client-provided read policy to the write frontier.
1911                read_policy.frontier(new_frontier.borrow())
1912            }
1913            None => {
1914                // Write-only collections cannot be read within the context of the compute
1915                // controller, so their read frontier only controls the read holds taken on their
1916                // inputs. We can safely downgrade the input read holds to any time less than the
1917                // write frontier.
1918                //
1919                // Note that some write-only collections (continual tasks) need to observe changes
1920                // at their current write frontier during hydration. Thus, we cannot downgrade the
1921                // read frontier to the write frontier and instead step it back by one.
1922                Antichain::from_iter(
1923                    new_frontier
1924                        .iter()
1925                        .map(|t| t.step_back().unwrap_or(Timestamp::MIN)),
1926                )
1927            }
1928        };
1929        let _ = collection.implied_read_hold.try_downgrade(new_since);
1930
1931        // Report the frontier advancement.
1932        self.deliver_response(ComputeControllerResponse::FrontierUpper {
1933            id,
1934            upper: new_frontier,
1935        });
1936    }
1937
1938    /// Apply a collection read hold change.
1939    pub(super) fn apply_read_hold_change(
1940        &mut self,
1941        id: GlobalId,
1942        mut update: ChangeBatch<Timestamp>,
1943    ) {
1944        let Some(collection) = self.collections.get_mut(&id) else {
1945            soft_panic_or_log!(
1946                "read hold change for absent collection (id={id}, changes={update:?})"
1947            );
1948            return;
1949        };
1950
1951        let new_since = collection.shared.lock_read_capabilities(|caps| {
1952            // Sanity check to prevent corrupted `read_capabilities`, which can cause hard-to-debug
1953            // issues (usually stuck read frontiers).
1954            let read_frontier = caps.frontier();
1955            for (time, diff) in update.iter() {
1956                let count = caps.count_for(time) + diff;
1957                assert!(
1958                    count >= 0,
1959                    "invalid read capabilities update: negative capability \
1960             (id={id:?}, read_capabilities={caps:?}, update={update:?})",
1961                );
1962                assert!(
1963                    count == 0 || read_frontier.less_equal(time),
1964                    "invalid read capabilities update: frontier regression \
1965             (id={id:?}, read_capabilities={caps:?}, update={update:?})",
1966                );
1967            }
1968
1969            // Apply read capability updates and learn about resulting changes to the read
1970            // frontier.
1971            let changes = caps.update_iter(update.drain());
1972
1973            let changed = changes.count() > 0;
1974            changed.then(|| caps.frontier().to_owned())
1975        });
1976
1977        let Some(new_since) = new_since else {
1978            return; // read frontier did not change
1979        };
1980
1981        // Propagate read frontier update to dependencies.
1982        for read_hold in collection.compute_dependencies.values_mut() {
1983            read_hold
1984                .try_downgrade(new_since.clone())
1985                .expect("frontiers don't regress");
1986        }
1987        for read_hold in collection.storage_dependencies.values_mut() {
1988            read_hold
1989                .try_downgrade(new_since.clone())
1990                .expect("frontiers don't regress");
1991        }
1992
1993        // Produce `AllowCompaction` command.
1994        self.send(ComputeCommand::AllowCompaction {
1995            id,
1996            frontier: new_since,
1997        });
1998    }
1999
2000    /// Fulfills a registered peek and cleans up associated state.
2001    ///
2002    /// As part of this we:
2003    ///  * Send a `PeekResponse` through the peek's response channel.
2004    ///  * Emit a `CancelPeek` command to instruct replicas to stop spending resources on this
2005    ///    peek, and to allow the `ComputeCommandHistory` to reduce away the corresponding `Peek`
2006    ///    command.
2007    ///  * Remove the read hold for this peek, unblocking compaction that might have waited on it.
2008    fn finish_peek(&mut self, uuid: Uuid, response: PeekResponse) {
2009        let Some(peek) = self.peeks.remove(&uuid) else {
2010            return;
2011        };
2012
2013        // The recipient might not be interested in the peek response anymore, which is fine.
2014        let _ = peek.peek_response_tx.send(response);
2015
2016        // NOTE: We need to send the `CancelPeek` command _before_ we release the peek's read hold
2017        // (by dropping it), to avoid the edge case that caused database-issues#4812.
2018        self.send(ComputeCommand::CancelPeek { uuid });
2019
2020        drop(peek.read_hold);
2021    }
2022
2023    /// Handles a response from a replica. Replica IDs are re-used across replica restarts, so we
2024    /// use the replica epoch to drop stale responses.
2025    fn handle_response(&mut self, (replica_id, epoch, response): ReplicaResponse) {
2026        // Filter responses from non-existing or stale replicas.
2027        if self
2028            .replicas
2029            .get(&replica_id)
2030            .filter(|replica| replica.epoch == epoch)
2031            .is_none()
2032        {
2033            return;
2034        }
2035
2036        // Invariant: the replica exists and has the expected epoch.
2037
2038        match response {
2039            ComputeResponse::Frontiers(id, frontiers) => {
2040                self.handle_frontiers_response(id, frontiers, replica_id);
2041            }
2042            ComputeResponse::PeekResponse(uuid, peek_response, otel_ctx) => {
2043                self.handle_peek_response(uuid, peek_response, otel_ctx, replica_id);
2044            }
2045            ComputeResponse::CopyToResponse(id, response) => {
2046                self.handle_copy_to_response(id, response, replica_id);
2047            }
2048            ComputeResponse::SubscribeResponse(id, response) => {
2049                self.handle_subscribe_response(id, response, replica_id);
2050            }
2051            ComputeResponse::Status(response) => {
2052                self.handle_status_response(response, replica_id);
2053            }
2054        }
2055    }
2056
2057    /// Handle new frontiers, returning any compute response that needs to
2058    /// be sent to the client.
2059    fn handle_frontiers_response(
2060        &mut self,
2061        id: GlobalId,
2062        frontiers: FrontiersResponse,
2063        replica_id: ReplicaId,
2064    ) {
2065        if !self.collections.contains_key(&id) {
2066            soft_panic_or_log!(
2067                "frontiers update for an unknown collection \
2068                 (id={id}, replica_id={replica_id}, frontiers={frontiers:?})"
2069            );
2070            return;
2071        }
2072        let Some(replica) = self.replicas.get_mut(&replica_id) else {
2073            soft_panic_or_log!(
2074                "frontiers update for an unknown replica \
2075                 (replica_id={replica_id}, frontiers={frontiers:?})"
2076            );
2077            return;
2078        };
2079        let Some(replica_collection) = replica.collections.get_mut(&id) else {
2080            soft_panic_or_log!(
2081                "frontiers update for an unknown replica collection \
2082                 (id={id}, replica_id={replica_id}, frontiers={frontiers:?})"
2083            );
2084            return;
2085        };
2086
2087        if let Some(new_frontier) = frontiers.input_frontier {
2088            replica_collection.update_input_frontier(new_frontier.clone());
2089        }
2090        if let Some(new_frontier) = frontiers.output_frontier {
2091            replica_collection.update_output_frontier(new_frontier.clone());
2092        }
2093        if let Some(new_frontier) = frontiers.write_frontier {
2094            replica_collection.update_write_frontier(new_frontier.clone());
2095            self.maybe_update_global_write_frontier(id, new_frontier);
2096        }
2097    }
2098
2099    #[mz_ore::instrument(level = "debug")]
2100    fn handle_peek_response(
2101        &mut self,
2102        uuid: Uuid,
2103        response: PeekResponse,
2104        otel_ctx: OpenTelemetryContext,
2105        replica_id: ReplicaId,
2106    ) {
2107        otel_ctx.attach_as_parent();
2108
2109        // We might not be tracking this peek anymore, because we have served a response already or
2110        // because it was canceled. If this is the case, we ignore the response.
2111        let Some(peek) = self.peeks.get(&uuid) else {
2112            return;
2113        };
2114
2115        // If the peek is targeting a replica, ignore responses from other replicas.
2116        let target_replica = peek.target_replica.unwrap_or(replica_id);
2117        if target_replica != replica_id {
2118            return;
2119        }
2120
2121        let duration = peek.requested_at.elapsed();
2122        self.metrics.observe_peek_response(&response, duration);
2123
2124        let notification = PeekNotification::new(&response, peek.offset, peek.limit);
2125        // NOTE: We use the `otel_ctx` from the response, not the pending peek, because we
2126        // currently want the parent to be whatever the compute worker did with this peek.
2127        self.deliver_response(ComputeControllerResponse::PeekNotification(
2128            uuid,
2129            notification,
2130            otel_ctx,
2131        ));
2132
2133        self.finish_peek(uuid, response)
2134    }
2135
2136    fn handle_copy_to_response(
2137        &mut self,
2138        sink_id: GlobalId,
2139        response: CopyToResponse,
2140        replica_id: ReplicaId,
2141    ) {
2142        if !self.collections.contains_key(&sink_id) {
2143            soft_panic_or_log!(
2144                "received response for an unknown copy-to \
2145                 (sink_id={sink_id}, replica_id={replica_id})",
2146            );
2147            return;
2148        }
2149        let Some(replica) = self.replicas.get_mut(&replica_id) else {
2150            soft_panic_or_log!("copy-to response for an unknown replica (replica_id={replica_id})");
2151            return;
2152        };
2153        let Some(replica_collection) = replica.collections.get_mut(&sink_id) else {
2154            soft_panic_or_log!(
2155                "copy-to response for an unknown replica collection \
2156                 (sink_id={sink_id}, replica_id={replica_id})"
2157            );
2158            return;
2159        };
2160
2161        // Downgrade the replica frontiers, to enable dropping of input read holds and clean up of
2162        // collection state.
2163        // TODO(database-issues#4701): report copy-to frontiers through `Frontiers` responses
2164        replica_collection.update_write_frontier(Antichain::new());
2165        replica_collection.update_input_frontier(Antichain::new());
2166        replica_collection.update_output_frontier(Antichain::new());
2167
2168        // We might not be tracking this COPY TO because we have already returned a response
2169        // from one of the replicas. In that case, we ignore the response.
2170        if !self.copy_tos.remove(&sink_id) {
2171            return;
2172        }
2173
2174        let result = match response {
2175            CopyToResponse::RowCount(count) => Ok(count),
2176            CopyToResponse::Error(error) => Err(anyhow::anyhow!(error)),
2177            // We should never get here: Replicas only drop copy to collections in response
2178            // to the controller allowing them to do so, and when the controller drops a
2179            // copy to it also removes it from the list of tracked copy_tos (see
2180            // [`Instance::drop_collections`]).
2181            CopyToResponse::Dropped => {
2182                tracing::error!(
2183                    %sink_id, %replica_id,
2184                    "received `Dropped` response for a tracked copy to",
2185                );
2186                return;
2187            }
2188        };
2189
2190        self.deliver_response(ComputeControllerResponse::CopyToResponse(sink_id, result));
2191    }
2192
2193    fn handle_subscribe_response(
2194        &mut self,
2195        subscribe_id: GlobalId,
2196        response: SubscribeResponse,
2197        replica_id: ReplicaId,
2198    ) {
2199        if !self.collections.contains_key(&subscribe_id) {
2200            soft_panic_or_log!(
2201                "received response for an unknown subscribe \
2202                 (subscribe_id={subscribe_id}, replica_id={replica_id})",
2203            );
2204            return;
2205        }
2206        let Some(replica) = self.replicas.get_mut(&replica_id) else {
2207            soft_panic_or_log!(
2208                "subscribe response for an unknown replica (replica_id={replica_id})"
2209            );
2210            return;
2211        };
2212        let Some(replica_collection) = replica.collections.get_mut(&subscribe_id) else {
2213            soft_panic_or_log!(
2214                "subscribe response for an unknown replica collection \
2215                 (subscribe_id={subscribe_id}, replica_id={replica_id})"
2216            );
2217            return;
2218        };
2219
2220        // Always apply replica write frontier updates. Even if the subscribe is not tracked
2221        // anymore, there might still be replicas reading from its inputs, so we need to track the
2222        // frontiers until all replicas have advanced to the empty one.
2223        let write_frontier = match &response {
2224            SubscribeResponse::Batch(batch) => batch.upper.clone(),
2225            SubscribeResponse::DroppedAt(_) => Antichain::new(),
2226        };
2227
2228        // For subscribes we downgrade all replica frontiers based on write frontiers. This should
2229        // be fine because the input and output frontier of a subscribe track its write frontier.
2230        // TODO(database-issues#4701): report subscribe frontiers through `Frontiers` responses
2231        replica_collection.update_write_frontier(write_frontier.clone());
2232        replica_collection.update_input_frontier(write_frontier.clone());
2233        replica_collection.update_output_frontier(write_frontier.clone());
2234
2235        // If the subscribe is not tracked, or targets a different replica, there is nothing to do.
2236        let Some(mut subscribe) = self.subscribes.get(&subscribe_id).cloned() else {
2237            return;
2238        };
2239
2240        // Apply a global frontier update.
2241        // If this is a replica-targeted subscribe, it is important that we advance the global
2242        // frontier only based on responses from the targeted replica. Otherwise, another replica
2243        // could advance to the empty frontier, making us drop the subscribe on the targeted
2244        // replica prematurely.
2245        self.maybe_update_global_write_frontier(subscribe_id, write_frontier);
2246
2247        match response {
2248            SubscribeResponse::Batch(batch) => {
2249                let upper = batch.upper;
2250                let mut updates = batch.updates;
2251
2252                // If this batch advances the subscribe's frontier, we emit all updates at times
2253                // greater or equal to the last frontier (to avoid emitting duplicate updates).
2254                if PartialOrder::less_than(&subscribe.frontier, &upper) {
2255                    let lower = std::mem::replace(&mut subscribe.frontier, upper.clone());
2256
2257                    if upper.is_empty() {
2258                        // This subscribe cannot produce more data. Stop tracking it.
2259                        self.subscribes.remove(&subscribe_id);
2260                    } else {
2261                        // This subscribe can produce more data. Update our tracking of it.
2262                        self.subscribes.insert(subscribe_id, subscribe);
2263                    }
2264
2265                    if let Ok(updates) = updates.as_mut() {
2266                        updates.retain_mut(|updates| {
2267                            let offset = updates.times().partition_point(|t| {
2268                                // True for times that are strictly less than lower (and should be skipped)
2269                                // and false otherwise.
2270                                !lower.less_equal(t)
2271                            });
2272                            let (_, past_lower) = std::mem::take(updates).split_at(offset);
2273                            *updates = past_lower;
2274                            updates.len() > 0
2275                        });
2276                    }
2277                    self.deliver_response(ComputeControllerResponse::SubscribeResponse(
2278                        subscribe_id,
2279                        SubscribeBatch {
2280                            lower,
2281                            upper,
2282                            updates,
2283                        },
2284                    ));
2285                }
2286            }
2287            SubscribeResponse::DroppedAt(frontier) => {
2288                // We should never get here: Replicas only drop subscribe collections in response
2289                // to the controller allowing them to do so, and when the controller drops a
2290                // subscribe it also removes it from the list of tracked subscribes (see
2291                // [`Instance::drop_collections`]).
2292                tracing::error!(
2293                    %subscribe_id,
2294                    %replica_id,
2295                    frontier = ?frontier.elements(),
2296                    "received `DroppedAt` response for a tracked subscribe",
2297                );
2298                self.subscribes.remove(&subscribe_id);
2299            }
2300        }
2301    }
2302
2303    fn handle_status_response(&self, response: StatusResponse, _replica_id: ReplicaId) {
2304        match response {
2305            StatusResponse::Placeholder => {}
2306        }
2307    }
2308
2309    /// Return the write frontiers of the dependencies of the given collection.
2310    fn dependency_write_frontiers<'b>(
2311        &'b self,
2312        collection: &'b CollectionState,
2313    ) -> impl Iterator<Item = Antichain<Timestamp>> + 'b {
2314        let compute_frontiers = collection.compute_dependency_ids().filter_map(|dep_id| {
2315            let collection = self.collections.get(&dep_id);
2316            collection.map(|c| c.write_frontier())
2317        });
2318        let storage_frontiers = collection.storage_dependency_ids().filter_map(|dep_id| {
2319            let frontiers = self.storage_collections.collection_frontiers(dep_id).ok();
2320            frontiers.map(|f| f.write_frontier)
2321        });
2322
2323        compute_frontiers.chain(storage_frontiers)
2324    }
2325
2326    /// Return the write frontiers of transitive storage dependencies of the given collection.
2327    fn transitive_storage_dependency_write_frontiers<'b>(
2328        &'b self,
2329        collection: &'b CollectionState,
2330    ) -> impl Iterator<Item = Antichain<Timestamp>> + 'b {
2331        let mut storage_ids: BTreeSet<_> = collection.storage_dependency_ids().collect();
2332        let mut todo: Vec<_> = collection.compute_dependency_ids().collect();
2333        let mut done = BTreeSet::new();
2334
2335        while let Some(id) = todo.pop() {
2336            if done.contains(&id) {
2337                continue;
2338            }
2339            if let Some(dep) = self.collections.get(&id) {
2340                storage_ids.extend(dep.storage_dependency_ids());
2341                todo.extend(dep.compute_dependency_ids())
2342            }
2343            done.insert(id);
2344        }
2345
2346        let storage_frontiers = storage_ids.into_iter().filter_map(|id| {
2347            let frontiers = self.storage_collections.collection_frontiers(id).ok();
2348            frontiers.map(|f| f.write_frontier)
2349        });
2350
2351        storage_frontiers
2352    }
2353
2354    /// Downgrade the warmup capabilities of collections as much as possible.
2355    ///
2356    /// The only requirement we have for a collection's warmup capability is that it is for a time
2357    /// that is available in all of the collection's inputs. For each input the latest time that is
2358    /// the case for is `write_frontier - 1`. So the farthest we can downgrade a collection's
2359    /// warmup capability is the minimum of `write_frontier - 1` of all its inputs.
2360    ///
2361    /// This method expects to be periodically called as part of instance maintenance work.
2362    /// We would like to instead update the warmup capabilities synchronously in response to
2363    /// frontier updates of dependency collections, but that is not generally possible because we
2364    /// don't learn about frontier updates of storage collections synchronously. We could do
2365    /// synchronous updates for compute dependencies, but we refrain from doing for simplicity.
2366    fn downgrade_warmup_capabilities(&mut self) {
2367        let mut new_capabilities = BTreeMap::new();
2368        for (id, collection) in &self.collections {
2369            // For write-only collections that have advanced to the empty frontier, we can drop the
2370            // warmup capability entirely. There is no reason why we would need to hydrate those
2371            // collections again, so being able to warm them up is not useful.
2372            if collection.read_policy.is_none()
2373                && collection.shared.lock_write_frontier(|f| f.is_empty())
2374            {
2375                new_capabilities.insert(*id, Antichain::new());
2376                continue;
2377            }
2378
2379            let mut new_capability = Antichain::new();
2380            for frontier in self.dependency_write_frontiers(collection) {
2381                for time in frontier {
2382                    new_capability.insert(time.step_back().unwrap_or(time));
2383                }
2384            }
2385
2386            new_capabilities.insert(*id, new_capability);
2387        }
2388
2389        for (id, new_capability) in new_capabilities {
2390            let collection = self.expect_collection_mut(id);
2391            let _ = collection.warmup_read_hold.try_downgrade(new_capability);
2392        }
2393    }
2394
2395    /// Forward the implied capabilities of collections, if possible.
2396    ///
2397    /// The implied capability of a collection controls (a) which times are still readable (for
2398    /// indexes) and (b) with which as-of the collection gets installed on a new replica. We are
2399    /// usually not allowed to advance an implied capability beyond the frontier that follows from
2400    /// the collection's read policy applied to its write frontier:
2401    ///
2402    ///  * For sink collections, some external consumer might rely on seeing all distinct times in
2403    ///    the input reflected in the output. If we'd forward the implied capability of a sink,
2404    ///    we'd risk skipping times in the output across replica restarts.
2405    ///  * For index collections, we might make the index unreadable by advancing its read frontier
2406    ///    beyond its write frontier.
2407    ///
2408    /// There is one case where forwarding an implied capability is fine though: an index installed
2409    /// on a cluster that has no replicas. Such indexes are not readable anyway until a new replica
2410    /// is added, so advancing its read frontier can't make it unreadable. We can thus advance the
2411    /// implied capability as long as we make sure that when a new replica is added, the expected
2412    /// relationship between write frontier, read policy, and implied capability can be restored
2413    /// immediately (modulo computation time).
2414    ///
2415    /// Forwarding implied capabilities is not necessary for the correct functioning of the
2416    /// controller but an optimization that is beneficial in two ways:
2417    ///
2418    ///  * It relaxes read holds on inputs to forwarded collections, allowing their compaction.
2419    ///  * It reduces the amount of historical detail new replicas need to process when computing
2420    ///    forwarded collections, as forwarding the implied capability also forwards the corresponding
2421    ///    dataflow as-of.
2422    fn forward_implied_capabilities(&mut self) {
2423        if !ENABLE_PAUSED_CLUSTER_READHOLD_DOWNGRADE.get(&self.dyncfg) {
2424            return;
2425        }
2426        if !self.replicas.is_empty() {
2427            return;
2428        }
2429
2430        let mut new_capabilities = BTreeMap::new();
2431        for (id, collection) in &self.collections {
2432            let Some(read_policy) = &collection.read_policy else {
2433                // Collection is write-only, i.e. a sink.
2434                continue;
2435            };
2436
2437            // When a new replica is started, it will immediately be able to compute all collection
2438            // output up to the write frontier of its transitive storage inputs. So the new implied
2439            // read capability should be the read policy applied to that frontier.
2440            let mut dep_frontier = Antichain::new();
2441            for frontier in self.transitive_storage_dependency_write_frontiers(collection) {
2442                dep_frontier.extend(frontier);
2443            }
2444
2445            let new_capability = read_policy.frontier(dep_frontier.borrow());
2446            if PartialOrder::less_than(collection.implied_read_hold.since(), &new_capability) {
2447                new_capabilities.insert(*id, new_capability);
2448            }
2449        }
2450
2451        for (id, new_capability) in new_capabilities {
2452            let collection = self.expect_collection_mut(id);
2453            let _ = collection.implied_read_hold.try_downgrade(new_capability);
2454        }
2455    }
2456
2457    /// Acquires a `ReadHold` for the identified compute collection.
2458    ///
2459    /// This mirrors the logic used by the controller-side `InstanceState::acquire_read_hold`,
2460    /// but executes on the instance task itself.
2461    pub(super) fn acquire_read_hold(&self, id: GlobalId) -> Result<ReadHold, CollectionMissing> {
2462        // Similarly to InstanceState::acquire_read_hold and StorageCollections::acquire_read_holds,
2463        // we acquire read holds at the earliest possible time rather than returning a copy
2464        // of the implied read hold. This is so that dependents can acquire read holds on
2465        // compute dependencies at frontiers that are held back by other read holds the caller
2466        // has previously taken.
2467        let collection = self.collection(id)?;
2468        let since = collection.shared.lock_read_capabilities(|caps| {
2469            let since = caps.frontier().to_owned();
2470            caps.update_iter(since.iter().map(|t| (t.clone(), 1)));
2471            since
2472        });
2473        let hold = ReadHold::new(id, since, Arc::clone(&self.read_hold_tx));
2474        Ok(hold)
2475    }
2476
2477    /// Process pending maintenance work.
2478    ///
2479    /// This method is invoked periodically by the global controller.
2480    /// It is a good place to perform maintenance work that arises from various controller state
2481    /// changes and that cannot conveniently be handled synchronously with those state changes.
2482    #[mz_ore::instrument(level = "debug")]
2483    pub fn maintain(&mut self) {
2484        self.rehydrate_failed_replicas();
2485        self.downgrade_warmup_capabilities();
2486        self.forward_implied_capabilities();
2487        self.schedule_collections();
2488        self.cleanup_collections();
2489        self.update_frontier_introspection();
2490        self.refresh_state_metrics();
2491        self.refresh_wallclock_lag();
2492    }
2493}
2494
2495/// State maintained about individual compute collections.
2496///
2497/// A compute collection is either an index, or a storage sink, or a subscribe, exported by a
2498/// compute dataflow.
2499#[derive(Debug)]
2500struct CollectionState {
2501    /// If set, this collection is only maintained by the specified replica.
2502    target_replica: Option<ReplicaId>,
2503    /// Whether this collection is a log collection.
2504    ///
2505    /// Log collections are special in that they are only maintained by a subset of all replicas.
2506    log_collection: bool,
2507    /// Whether this collection has been dropped by a controller client.
2508    ///
2509    /// The controller is allowed to remove the `CollectionState` for a collection only when
2510    /// `dropped == true`. Otherwise, clients might still expect to be able to query information
2511    /// about this collection.
2512    dropped: bool,
2513    /// Whether this collection has been scheduled, i.e., the controller has sent a `Schedule`
2514    /// command for it.
2515    scheduled: bool,
2516
2517    /// Whether this collection is in read-only mode.
2518    ///
2519    /// When in read-only mode, the dataflow is not allowed to affect external state (largely persist).
2520    read_only: bool,
2521
2522    /// State shared with the `ComputeController`.
2523    shared: SharedCollectionState,
2524
2525    /// A read hold maintaining the implicit capability of the collection.
2526    ///
2527    /// This capability is kept to ensure that the collection remains readable according to its
2528    /// `read_policy`. It also ensures that read holds on the collection's dependencies are kept at
2529    /// some time not greater than the collection's `write_frontier`, guaranteeing that the
2530    /// collection's next outputs can always be computed without skipping times.
2531    implied_read_hold: ReadHold,
2532    /// A read hold held to enable dataflow warmup.
2533    ///
2534    /// Dataflow warmup is an optimization that allows dataflows to immediately start hydrating
2535    /// even when their next output time (as implied by the `write_frontier`) is in the future.
2536    /// By installing a read capability derived from the write frontiers of the collection's
2537    /// inputs, we ensure that the as-of of new dataflows installed for the collection is at a time
2538    /// that is immediately available, so hydration can begin immediately too.
2539    warmup_read_hold: ReadHold,
2540    /// The policy to use to downgrade `self.implied_read_hold`.
2541    ///
2542    /// If `None`, the collection is a write-only collection (i.e. a sink). For write-only
2543    /// collections, the `implied_read_hold` is only required for maintaining read holds on the
2544    /// inputs, so we can immediately downgrade it to the `write_frontier`.
2545    read_policy: Option<ReadPolicy>,
2546
2547    /// Storage identifiers on which this collection depends, and read holds this collection
2548    /// requires on them.
2549    storage_dependencies: BTreeMap<GlobalId, ReadHold>,
2550    /// Compute identifiers on which this collection depends, and read holds this collection
2551    /// requires on them.
2552    compute_dependencies: BTreeMap<GlobalId, ReadHold>,
2553
2554    /// Introspection state associated with this collection.
2555    introspection: CollectionIntrospection,
2556
2557    /// Frontier wallclock lag measurements stashed until the next `WallclockLagHistogram`
2558    /// introspection update.
2559    ///
2560    /// Keys are `(period, lag, labels)` triples, values are counts.
2561    ///
2562    /// If this is `None`, wallclock lag is not tracked for this collection.
2563    wallclock_lag_histogram_stash: Option<
2564        BTreeMap<
2565            (
2566                WallclockLagHistogramPeriod,
2567                WallclockLag,
2568                BTreeMap<&'static str, String>,
2569            ),
2570            Diff,
2571        >,
2572    >,
2573}
2574
2575impl CollectionState {
2576    /// Creates a new collection state, with an initial read policy valid from `since`.
2577    fn new(
2578        collection_id: GlobalId,
2579        as_of: Antichain<Timestamp>,
2580        shared: SharedCollectionState,
2581        storage_dependencies: BTreeMap<GlobalId, ReadHold>,
2582        compute_dependencies: BTreeMap<GlobalId, ReadHold>,
2583        read_hold_tx: read_holds::ChangeTx,
2584        introspection: CollectionIntrospection,
2585    ) -> Self {
2586        // A collection is not readable before the `as_of`.
2587        let since = as_of.clone();
2588        // A collection won't produce updates for times before the `as_of`.
2589        let upper = as_of;
2590
2591        // Ensure that the provided `shared` is valid for the given `as_of`.
2592        assert!(shared.lock_read_capabilities(|c| c.frontier() == since.borrow()));
2593        assert!(shared.lock_write_frontier(|f| f == &upper));
2594
2595        // Initialize collection read holds.
2596        // Note that the implied read hold was already added to the `read_capabilities` when
2597        // `shared` was created, so we only need to add the warmup read hold here.
2598        let implied_read_hold =
2599            ReadHold::new(collection_id, since.clone(), Arc::clone(&read_hold_tx));
2600        let warmup_read_hold = ReadHold::new(collection_id, since.clone(), read_hold_tx);
2601
2602        let updates = warmup_read_hold.since().iter().map(|t| (t.clone(), 1));
2603        shared.lock_read_capabilities(|c| {
2604            c.update_iter(updates);
2605        });
2606
2607        // In an effort to keep the produced wallclock lag introspection data small and
2608        // predictable, we disable wallclock lag tracking for transient collections, i.e. slow-path
2609        // select indexes and subscribes.
2610        let wallclock_lag_histogram_stash = match collection_id.is_transient() {
2611            true => None,
2612            false => Some(Default::default()),
2613        };
2614
2615        Self {
2616            target_replica: None,
2617            log_collection: false,
2618            dropped: false,
2619            scheduled: false,
2620            read_only: true,
2621            shared,
2622            implied_read_hold,
2623            warmup_read_hold,
2624            read_policy: Some(ReadPolicy::ValidFrom(since)),
2625            storage_dependencies,
2626            compute_dependencies,
2627            introspection,
2628            wallclock_lag_histogram_stash,
2629        }
2630    }
2631
2632    /// Creates a new collection state for a log collection.
2633    fn new_log_collection(
2634        id: GlobalId,
2635        shared: SharedCollectionState,
2636        read_hold_tx: read_holds::ChangeTx,
2637        introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
2638    ) -> Self {
2639        let since = Antichain::from_elem(Timestamp::MIN);
2640        let introspection = CollectionIntrospection::new(
2641            id,
2642            introspection_tx,
2643            since.clone(),
2644            false,
2645            None,
2646            None,
2647            Vec::new(),
2648        );
2649        let mut state = Self::new(
2650            id,
2651            since,
2652            shared,
2653            Default::default(),
2654            Default::default(),
2655            read_hold_tx,
2656            introspection,
2657        );
2658        state.log_collection = true;
2659        // Log collections are created and scheduled implicitly as part of replica initialization.
2660        state.scheduled = true;
2661        state
2662    }
2663
2664    /// Reports the current read frontier.
2665    fn read_frontier(&self) -> Antichain<Timestamp> {
2666        self.shared
2667            .lock_read_capabilities(|c| c.frontier().to_owned())
2668    }
2669
2670    /// Reports the current write frontier.
2671    fn write_frontier(&self) -> Antichain<Timestamp> {
2672        self.shared.lock_write_frontier(|f| f.clone())
2673    }
2674
2675    fn storage_dependency_ids(&self) -> impl Iterator<Item = GlobalId> + '_ {
2676        self.storage_dependencies.keys().copied()
2677    }
2678
2679    fn compute_dependency_ids(&self) -> impl Iterator<Item = GlobalId> + '_ {
2680        self.compute_dependencies.keys().copied()
2681    }
2682}
2683
2684/// Collection state shared with the `ComputeController`.
2685///
2686/// Having this allows certain controller APIs, such as `ComputeController::collection_frontiers`
2687/// and `ComputeController::acquire_read_hold` to be non-`async`. This comes at the cost of
2688/// complexity (by introducing shared mutable state) and performance (by introducing locking). We
2689/// should aim to reduce the amount of shared state over time, rather than expand it.
2690///
2691/// Note that [`SharedCollectionState`]s are initialized by the `ComputeController` prior to the
2692/// collection's creation in the [`Instance`]. This is to allow compute clients to query frontiers
2693/// and take new read holds immediately, without having to wait for the [`Instance`] to update.
2694#[derive(Clone, Debug)]
2695pub(super) struct SharedCollectionState {
2696    /// Accumulation of read capabilities for the collection.
2697    ///
2698    /// This accumulation contains the capabilities held by all [`ReadHold`]s given out for the
2699    /// collection, including `implied_read_hold` and `warmup_read_hold`.
2700    ///
2701    /// NOTE: This field may only be modified by [`Instance::apply_read_hold_change`],
2702    /// [`Instance::acquire_read_hold`], and `ComputeController::acquire_read_hold`.
2703    /// Nobody else should modify read capabilities directly. Instead, collection users should
2704    /// manage read holds through [`ReadHold`] objects acquired through
2705    /// `ComputeController::acquire_read_hold`.
2706    ///
2707    /// TODO(teskje): Restructure the code to enforce the above in the type system.
2708    read_capabilities: Arc<Mutex<MutableAntichain<Timestamp>>>,
2709    /// The write frontier of this collection.
2710    write_frontier: Arc<Mutex<Antichain<Timestamp>>>,
2711}
2712
2713impl SharedCollectionState {
2714    pub fn new(as_of: Antichain<Timestamp>) -> Self {
2715        // A collection is not readable before the `as_of`.
2716        let since = as_of.clone();
2717        // A collection won't produce updates for times before the `as_of`.
2718        let upper = as_of;
2719
2720        // Initialize read capabilities to the `since`.
2721        // The is the implied read capability. The corresponding [`ReadHold`] is created in
2722        // [`CollectionState::new`].
2723        let mut read_capabilities = MutableAntichain::new();
2724        read_capabilities.update_iter(since.iter().map(|time| (time.clone(), 1)));
2725
2726        Self {
2727            read_capabilities: Arc::new(Mutex::new(read_capabilities)),
2728            write_frontier: Arc::new(Mutex::new(upper)),
2729        }
2730    }
2731
2732    pub fn lock_read_capabilities<F, R>(&self, f: F) -> R
2733    where
2734        F: FnOnce(&mut MutableAntichain<Timestamp>) -> R,
2735    {
2736        let mut caps = self.read_capabilities.lock().expect("poisoned");
2737        f(&mut *caps)
2738    }
2739
2740    pub fn lock_write_frontier<F, R>(&self, f: F) -> R
2741    where
2742        F: FnOnce(&mut Antichain<Timestamp>) -> R,
2743    {
2744        let mut frontier = self.write_frontier.lock().expect("poisoned");
2745        f(&mut *frontier)
2746    }
2747}
2748
2749/// Manages certain introspection relations associated with a collection. Upon creation, it adds
2750/// rows to introspection relations. When dropped, it retracts its managed rows.
2751#[derive(Debug)]
2752struct CollectionIntrospection {
2753    /// The ID of the compute collection.
2754    collection_id: GlobalId,
2755    /// A channel through which introspection updates are delivered.
2756    introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
2757    /// Introspection state for `IntrospectionType::Frontiers`.
2758    ///
2759    /// `Some` if the collection does _not_ sink into a storage collection (i.e. is not an MV). If
2760    /// the collection sinks into storage, the storage controller reports its frontiers instead.
2761    frontiers: Option<FrontiersIntrospectionState>,
2762    /// Introspection state for `IntrospectionType::ComputeMaterializedViewRefreshes`.
2763    ///
2764    /// `Some` if the collection is a REFRESH MV.
2765    refresh: Option<RefreshIntrospectionState>,
2766    /// The IDs of the collection's dependencies, for `IntrospectionType::ComputeDependencies`.
2767    dependency_ids: Vec<GlobalId>,
2768}
2769
2770impl CollectionIntrospection {
2771    fn new(
2772        collection_id: GlobalId,
2773        introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
2774        as_of: Antichain<Timestamp>,
2775        storage_sink: bool,
2776        initial_as_of: Option<Antichain<Timestamp>>,
2777        refresh_schedule: Option<RefreshSchedule>,
2778        dependency_ids: Vec<GlobalId>,
2779    ) -> Self {
2780        let refresh =
2781            match (refresh_schedule, initial_as_of) {
2782                (Some(refresh_schedule), Some(initial_as_of)) => Some(
2783                    RefreshIntrospectionState::new(refresh_schedule, initial_as_of, &as_of),
2784                ),
2785                (refresh_schedule, _) => {
2786                    // If we have a `refresh_schedule`, then the collection is a MV, so we should also have
2787                    // an `initial_as_of`.
2788                    soft_assert_or_log!(
2789                        refresh_schedule.is_none(),
2790                        "`refresh_schedule` without an `initial_as_of`: {collection_id}"
2791                    );
2792                    None
2793                }
2794            };
2795        let frontiers = (!storage_sink).then(|| FrontiersIntrospectionState::new(as_of));
2796
2797        let self_ = Self {
2798            collection_id,
2799            introspection_tx,
2800            frontiers,
2801            refresh,
2802            dependency_ids,
2803        };
2804
2805        self_.report_initial_state();
2806        self_
2807    }
2808
2809    /// Reports the initial introspection state.
2810    fn report_initial_state(&self) {
2811        if let Some(frontiers) = &self.frontiers {
2812            let row = frontiers.row_for_collection(self.collection_id);
2813            let updates = vec![(row, Diff::ONE)];
2814            self.send(IntrospectionType::Frontiers, updates);
2815        }
2816
2817        if let Some(refresh) = &self.refresh {
2818            let row = refresh.row_for_collection(self.collection_id);
2819            let updates = vec![(row, Diff::ONE)];
2820            self.send(IntrospectionType::ComputeMaterializedViewRefreshes, updates);
2821        }
2822
2823        if !self.dependency_ids.is_empty() {
2824            let updates = self.dependency_rows(Diff::ONE);
2825            self.send(IntrospectionType::ComputeDependencies, updates);
2826        }
2827    }
2828
2829    /// Produces rows for the `ComputeDependencies` introspection relation.
2830    fn dependency_rows(&self, diff: Diff) -> Vec<(Row, Diff)> {
2831        self.dependency_ids
2832            .iter()
2833            .map(|dependency_id| {
2834                let row = Row::pack_slice(&[
2835                    Datum::String(&self.collection_id.to_string()),
2836                    Datum::String(&dependency_id.to_string()),
2837                ]);
2838                (row, diff)
2839            })
2840            .collect()
2841    }
2842
2843    /// Observe the given current collection frontiers and update the introspection state as
2844    /// necessary.
2845    fn observe_frontiers(
2846        &mut self,
2847        read_frontier: &Antichain<Timestamp>,
2848        write_frontier: &Antichain<Timestamp>,
2849    ) {
2850        self.update_frontier_introspection(read_frontier, write_frontier);
2851        self.update_refresh_introspection(write_frontier);
2852    }
2853
2854    fn update_frontier_introspection(
2855        &mut self,
2856        read_frontier: &Antichain<Timestamp>,
2857        write_frontier: &Antichain<Timestamp>,
2858    ) {
2859        let Some(frontiers) = &mut self.frontiers else {
2860            return;
2861        };
2862
2863        if &frontiers.read_frontier == read_frontier && &frontiers.write_frontier == write_frontier
2864        {
2865            return; // no change
2866        };
2867
2868        let retraction = frontiers.row_for_collection(self.collection_id);
2869        frontiers.update(read_frontier, write_frontier);
2870        let insertion = frontiers.row_for_collection(self.collection_id);
2871        let updates = vec![(retraction, Diff::MINUS_ONE), (insertion, Diff::ONE)];
2872        self.send(IntrospectionType::Frontiers, updates);
2873    }
2874
2875    fn update_refresh_introspection(&mut self, write_frontier: &Antichain<Timestamp>) {
2876        let Some(refresh) = &mut self.refresh else {
2877            return;
2878        };
2879
2880        let retraction = refresh.row_for_collection(self.collection_id);
2881        refresh.frontier_update(write_frontier);
2882        let insertion = refresh.row_for_collection(self.collection_id);
2883
2884        if retraction == insertion {
2885            return; // no change
2886        }
2887
2888        let updates = vec![(retraction, Diff::MINUS_ONE), (insertion, Diff::ONE)];
2889        self.send(IntrospectionType::ComputeMaterializedViewRefreshes, updates);
2890    }
2891
2892    fn send(&self, introspection_type: IntrospectionType, updates: Vec<(Row, Diff)>) {
2893        // Failure to send means the `ComputeController` has been dropped and doesn't care about
2894        // introspection updates anymore.
2895        let _ = self.introspection_tx.send((introspection_type, updates));
2896    }
2897}
2898
2899impl Drop for CollectionIntrospection {
2900    fn drop(&mut self) {
2901        // Retract collection frontiers.
2902        if let Some(frontiers) = &self.frontiers {
2903            let row = frontiers.row_for_collection(self.collection_id);
2904            let updates = vec![(row, Diff::MINUS_ONE)];
2905            self.send(IntrospectionType::Frontiers, updates);
2906        }
2907
2908        // Retract MV refresh state.
2909        if let Some(refresh) = &self.refresh {
2910            let retraction = refresh.row_for_collection(self.collection_id);
2911            let updates = vec![(retraction, Diff::MINUS_ONE)];
2912            self.send(IntrospectionType::ComputeMaterializedViewRefreshes, updates);
2913        }
2914
2915        // Retract collection dependencies.
2916        if !self.dependency_ids.is_empty() {
2917            let updates = self.dependency_rows(Diff::MINUS_ONE);
2918            self.send(IntrospectionType::ComputeDependencies, updates);
2919        }
2920    }
2921}
2922
2923#[derive(Debug)]
2924struct FrontiersIntrospectionState {
2925    read_frontier: Antichain<Timestamp>,
2926    write_frontier: Antichain<Timestamp>,
2927}
2928
2929impl FrontiersIntrospectionState {
2930    fn new(as_of: Antichain<Timestamp>) -> Self {
2931        Self {
2932            read_frontier: as_of.clone(),
2933            write_frontier: as_of,
2934        }
2935    }
2936
2937    /// Return a `Row` reflecting the current collection frontiers.
2938    fn row_for_collection(&self, collection_id: GlobalId) -> Row {
2939        let read_frontier = self
2940            .read_frontier
2941            .as_option()
2942            .map_or(Datum::Null, |ts| ts.clone().into());
2943        let write_frontier = self
2944            .write_frontier
2945            .as_option()
2946            .map_or(Datum::Null, |ts| ts.clone().into());
2947        Row::pack_slice(&[
2948            Datum::String(&collection_id.to_string()),
2949            read_frontier,
2950            write_frontier,
2951        ])
2952    }
2953
2954    /// Update the introspection state with the given new frontiers.
2955    fn update(
2956        &mut self,
2957        read_frontier: &Antichain<Timestamp>,
2958        write_frontier: &Antichain<Timestamp>,
2959    ) {
2960        if read_frontier != &self.read_frontier {
2961            self.read_frontier.clone_from(read_frontier);
2962        }
2963        if write_frontier != &self.write_frontier {
2964            self.write_frontier.clone_from(write_frontier);
2965        }
2966    }
2967}
2968
2969/// Information needed to compute introspection updates for a REFRESH materialized view when the
2970/// write frontier advances.
2971#[derive(Debug)]
2972struct RefreshIntrospectionState {
2973    // Immutable properties of the MV
2974    refresh_schedule: RefreshSchedule,
2975    initial_as_of: Antichain<Timestamp>,
2976    // Refresh state
2977    next_refresh: Datum<'static>,           // Null or an MzTimestamp
2978    last_completed_refresh: Datum<'static>, // Null or an MzTimestamp
2979}
2980
2981impl RefreshIntrospectionState {
2982    /// Return a `Row` reflecting the current refresh introspection state.
2983    fn row_for_collection(&self, collection_id: GlobalId) -> Row {
2984        Row::pack_slice(&[
2985            Datum::String(&collection_id.to_string()),
2986            self.last_completed_refresh,
2987            self.next_refresh,
2988        ])
2989    }
2990}
2991
2992impl RefreshIntrospectionState {
2993    /// Construct a new [`RefreshIntrospectionState`], and apply an initial `frontier_update()` at
2994    /// the `upper`.
2995    fn new(
2996        refresh_schedule: RefreshSchedule,
2997        initial_as_of: Antichain<Timestamp>,
2998        upper: &Antichain<Timestamp>,
2999    ) -> Self {
3000        let mut self_ = Self {
3001            refresh_schedule: refresh_schedule.clone(),
3002            initial_as_of: initial_as_of.clone(),
3003            next_refresh: Datum::Null,
3004            last_completed_refresh: Datum::Null,
3005        };
3006        self_.frontier_update(upper);
3007        self_
3008    }
3009
3010    /// Should be called whenever the write frontier of the collection advances. It updates the
3011    /// state that should be recorded in introspection relations, but doesn't send the updates yet.
3012    fn frontier_update(&mut self, write_frontier: &Antichain<Timestamp>) {
3013        if write_frontier.is_empty() {
3014            self.last_completed_refresh =
3015                if let Some(last_refresh) = self.refresh_schedule.last_refresh() {
3016                    last_refresh.into()
3017                } else {
3018                    // If there is no last refresh, then we have a `REFRESH EVERY`, in which case
3019                    // the saturating roundup puts a refresh at the maximum possible timestamp.
3020                    Timestamp::MAX.into()
3021                };
3022            self.next_refresh = Datum::Null;
3023        } else {
3024            if PartialOrder::less_equal(write_frontier, &self.initial_as_of) {
3025                // We are before the first refresh.
3026                self.last_completed_refresh = Datum::Null;
3027                let initial_as_of = self.initial_as_of.as_option().expect(
3028                    "initial_as_of can't be [], because then there would be no refreshes at all",
3029                );
3030                let first_refresh = self
3031                    .refresh_schedule
3032                    .round_up_timestamp(*initial_as_of)
3033                    .expect("sequencing makes sure that REFRESH MVs always have a first refresh");
3034                soft_assert_or_log!(
3035                    first_refresh == *initial_as_of,
3036                    "initial_as_of should be set to the first refresh"
3037                );
3038                self.next_refresh = first_refresh.into();
3039            } else {
3040                // The first refresh has already happened.
3041                let write_frontier = write_frontier.as_option().expect("checked above");
3042                self.last_completed_refresh = self
3043                    .refresh_schedule
3044                    .round_down_timestamp_m1(*write_frontier)
3045                    .map_or_else(
3046                        || {
3047                            soft_panic_or_log!(
3048                                "rounding down should have returned the first refresh or later"
3049                            );
3050                            Datum::Null
3051                        },
3052                        |last_completed_refresh| last_completed_refresh.into(),
3053                    );
3054                self.next_refresh = write_frontier.clone().into();
3055            }
3056        }
3057    }
3058}
3059
3060/// A note of an outstanding peek response.
3061#[derive(Debug)]
3062struct PendingPeek {
3063    /// For replica-targeted peeks, this specifies the replica whose response we should pass on.
3064    ///
3065    /// If this value is `None`, we pass on the first response.
3066    target_replica: Option<ReplicaId>,
3067    /// The OpenTelemetry context for this peek.
3068    otel_ctx: OpenTelemetryContext,
3069    /// The time at which the peek was requested.
3070    ///
3071    /// Used to track peek durations.
3072    requested_at: Instant,
3073    /// The read hold installed to serve this peek.
3074    read_hold: ReadHold,
3075    /// The channel to send peek results.
3076    peek_response_tx: oneshot::Sender<PeekResponse>,
3077    /// An optional limit of the peek's result size.
3078    limit: Option<usize>,
3079    /// The offset into the peek's result.
3080    offset: usize,
3081}
3082
3083#[derive(Debug, Clone)]
3084struct ActiveSubscribe {
3085    /// Current upper frontier of this subscribe.
3086    frontier: Antichain<Timestamp>,
3087}
3088
3089impl Default for ActiveSubscribe {
3090    fn default() -> Self {
3091        Self {
3092            frontier: Antichain::from_elem(Timestamp::MIN),
3093        }
3094    }
3095}
3096
3097/// State maintained about individual replicas.
3098#[derive(Debug)]
3099struct ReplicaState {
3100    /// The ID of the replica.
3101    id: ReplicaId,
3102    /// Client for the running replica task.
3103    client: ReplicaClient,
3104    /// The replica configuration.
3105    config: ReplicaConfig,
3106    /// Replica metrics.
3107    metrics: ReplicaMetrics,
3108    /// A channel through which introspection updates are delivered.
3109    introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
3110    /// Per-replica collection state.
3111    collections: BTreeMap<GlobalId, ReplicaCollectionState>,
3112    /// The epoch of the replica.
3113    epoch: u64,
3114}
3115
3116impl ReplicaState {
3117    fn new(
3118        id: ReplicaId,
3119        client: ReplicaClient,
3120        config: ReplicaConfig,
3121        metrics: ReplicaMetrics,
3122        introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
3123        epoch: u64,
3124    ) -> Self {
3125        Self {
3126            id,
3127            client,
3128            config,
3129            metrics,
3130            introspection_tx,
3131            epoch,
3132            collections: Default::default(),
3133        }
3134    }
3135
3136    /// Add a collection to the replica state.
3137    ///
3138    /// # Panics
3139    ///
3140    /// Panics if a collection with the same ID exists already.
3141    fn add_collection(
3142        &mut self,
3143        id: GlobalId,
3144        as_of: Antichain<Timestamp>,
3145        input_read_holds: Vec<ReadHold>,
3146    ) {
3147        let metrics = self.metrics.for_collection(id);
3148        let introspection = ReplicaCollectionIntrospection::new(
3149            self.id,
3150            id,
3151            self.introspection_tx.clone(),
3152            as_of.clone(),
3153        );
3154        let mut state =
3155            ReplicaCollectionState::new(metrics, as_of, introspection, input_read_holds);
3156
3157        // In an effort to keep the produced wallclock lag introspection data small and
3158        // predictable, we disable wallclock lag tracking for transient collections, i.e. slow-path
3159        // select indexes and subscribes.
3160        if id.is_transient() {
3161            state.wallclock_lag_max = None;
3162        }
3163
3164        if let Some(previous) = self.collections.insert(id, state) {
3165            panic!("attempt to add a collection with existing ID {id} (previous={previous:?}");
3166        }
3167    }
3168
3169    /// Remove state for a collection.
3170    fn remove_collection(&mut self, id: GlobalId) -> Option<ReplicaCollectionState> {
3171        self.collections.remove(&id)
3172    }
3173
3174    /// Returns whether all replica frontiers of the given collection are empty.
3175    fn collection_frontiers_empty(&self, id: GlobalId) -> bool {
3176        self.collections.get(&id).map_or(true, |c| {
3177            c.write_frontier.is_empty()
3178                && c.input_frontier.is_empty()
3179                && c.output_frontier.is_empty()
3180        })
3181    }
3182
3183    /// Returns the state of the [`ReplicaState`] formatted as JSON.
3184    ///
3185    /// The returned value is not guaranteed to be stable and may change at any point in time.
3186    #[mz_ore::instrument(level = "debug")]
3187    pub fn dump(&self) -> Result<serde_json::Value, anyhow::Error> {
3188        // Note: We purposefully use the `Debug` formatting for the value of all fields in the
3189        // returned object as a tradeoff between usability and stability. `serde_json` will fail
3190        // to serialize an object if the keys aren't strings, so `Debug` formatting the values
3191        // prevents a future unrelated change from silently breaking this method.
3192
3193        // Destructure `self` here so we don't forget to consider dumping newly added fields.
3194        let Self {
3195            id,
3196            client: _,
3197            config: _,
3198            metrics: _,
3199            introspection_tx: _,
3200            epoch,
3201            collections,
3202        } = self;
3203
3204        let collections: BTreeMap<_, _> = collections
3205            .iter()
3206            .map(|(id, collection)| (id.to_string(), format!("{collection:?}")))
3207            .collect();
3208
3209        Ok(serde_json::json!({
3210            "id": id.to_string(),
3211            "collections": collections,
3212            "epoch": epoch,
3213        }))
3214    }
3215}
3216
3217#[derive(Debug)]
3218struct ReplicaCollectionState {
3219    /// The replica write frontier of this collection.
3220    ///
3221    /// See [`FrontiersResponse::write_frontier`].
3222    write_frontier: Antichain<Timestamp>,
3223    /// The replica input frontier of this collection.
3224    ///
3225    /// See [`FrontiersResponse::input_frontier`].
3226    input_frontier: Antichain<Timestamp>,
3227    /// The replica output frontier of this collection.
3228    ///
3229    /// See [`FrontiersResponse::output_frontier`].
3230    output_frontier: Antichain<Timestamp>,
3231
3232    /// Metrics tracked for this collection.
3233    ///
3234    /// If this is `None`, no metrics are collected.
3235    metrics: Option<ReplicaCollectionMetrics>,
3236    /// As-of frontier with which this collection was installed on the replica.
3237    as_of: Antichain<Timestamp>,
3238    /// Tracks introspection state for this collection.
3239    introspection: ReplicaCollectionIntrospection,
3240    /// Read holds on storage inputs to this collection.
3241    ///
3242    /// These read holds are kept to ensure that the replica is able to read from storage inputs at
3243    /// all times it hasn't read yet. We only need to install read holds for storage inputs since
3244    /// compaction of compute inputs is implicitly held back by Timely/DD.
3245    input_read_holds: Vec<ReadHold>,
3246
3247    /// Maximum frontier wallclock lag since the last `WallclockLagHistory` introspection update.
3248    ///
3249    /// If this is `None`, wallclock lag is not tracked for this collection.
3250    wallclock_lag_max: Option<WallclockLag>,
3251}
3252
3253impl ReplicaCollectionState {
3254    fn new(
3255        metrics: Option<ReplicaCollectionMetrics>,
3256        as_of: Antichain<Timestamp>,
3257        introspection: ReplicaCollectionIntrospection,
3258        input_read_holds: Vec<ReadHold>,
3259    ) -> Self {
3260        Self {
3261            write_frontier: as_of.clone(),
3262            input_frontier: as_of.clone(),
3263            output_frontier: as_of.clone(),
3264            metrics,
3265            as_of,
3266            introspection,
3267            input_read_holds,
3268            wallclock_lag_max: Some(WallclockLag::MIN),
3269        }
3270    }
3271
3272    /// Returns whether this collection is hydrated.
3273    fn hydrated(&self) -> bool {
3274        // If the observed frontier is greater than the collection's as-of, the collection has
3275        // produced some output and is therefore hydrated.
3276        //
3277        // We need to consider the edge case where the as-of is the empty frontier. Such an as-of
3278        // is not useful for indexes, because they wouldn't be readable. For write-only
3279        // collections, an empty as-of means that the collection has been fully written and no new
3280        // dataflow needs to be created for it. Consequently, no hydration will happen either.
3281        //
3282        // Based on this, we could respond in two ways:
3283        //  * `false`, as in "the dataflow was never created"
3284        //  * `true`, as in "the dataflow completed immediately"
3285        //
3286        // Since hydration is often used as a measure of dataflow progress and we don't want to
3287        // give the impression that certain dataflows are somehow stuck when they are not, we go
3288        // with the second interpretation here.
3289        self.as_of.is_empty() || PartialOrder::less_than(&self.as_of, &self.output_frontier)
3290    }
3291
3292    /// Updates the replica write frontier of this collection.
3293    fn update_write_frontier(&mut self, new_frontier: Antichain<Timestamp>) {
3294        if PartialOrder::less_than(&new_frontier, &self.write_frontier) {
3295            soft_panic_or_log!(
3296                "replica collection write frontier regression (old={:?}, new={new_frontier:?})",
3297                self.write_frontier,
3298            );
3299            return;
3300        } else if new_frontier == self.write_frontier {
3301            return;
3302        }
3303
3304        self.write_frontier = new_frontier;
3305    }
3306
3307    /// Updates the replica input frontier of this collection.
3308    fn update_input_frontier(&mut self, new_frontier: Antichain<Timestamp>) {
3309        if PartialOrder::less_than(&new_frontier, &self.input_frontier) {
3310            soft_panic_or_log!(
3311                "replica collection input frontier regression (old={:?}, new={new_frontier:?})",
3312                self.input_frontier,
3313            );
3314            return;
3315        } else if new_frontier == self.input_frontier {
3316            return;
3317        }
3318
3319        self.input_frontier = new_frontier;
3320
3321        // Relax our read holds on collection inputs.
3322        for read_hold in &mut self.input_read_holds {
3323            let result = read_hold.try_downgrade(self.input_frontier.clone());
3324            soft_assert_or_log!(
3325                result.is_ok(),
3326                "read hold downgrade failed (read_hold={read_hold:?}, new_since={:?})",
3327                self.input_frontier,
3328            );
3329        }
3330    }
3331
3332    /// Updates the replica output frontier of this collection.
3333    fn update_output_frontier(&mut self, new_frontier: Antichain<Timestamp>) {
3334        if PartialOrder::less_than(&new_frontier, &self.output_frontier) {
3335            soft_panic_or_log!(
3336                "replica collection output frontier regression (old={:?}, new={new_frontier:?})",
3337                self.output_frontier,
3338            );
3339            return;
3340        } else if new_frontier == self.output_frontier {
3341            return;
3342        }
3343
3344        self.output_frontier = new_frontier;
3345    }
3346}
3347
3348/// Maintains the introspection state for a given replica and collection, and ensures that reported
3349/// introspection data is retracted when the collection is dropped.
3350#[derive(Debug)]
3351struct ReplicaCollectionIntrospection {
3352    /// The ID of the replica.
3353    replica_id: ReplicaId,
3354    /// The ID of the compute collection.
3355    collection_id: GlobalId,
3356    /// The collection's reported replica write frontier.
3357    write_frontier: Antichain<Timestamp>,
3358    /// A channel through which introspection updates are delivered.
3359    introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
3360}
3361
3362impl ReplicaCollectionIntrospection {
3363    /// Create a new `HydrationState` and initialize introspection.
3364    fn new(
3365        replica_id: ReplicaId,
3366        collection_id: GlobalId,
3367        introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
3368        as_of: Antichain<Timestamp>,
3369    ) -> Self {
3370        let self_ = Self {
3371            replica_id,
3372            collection_id,
3373            write_frontier: as_of,
3374            introspection_tx,
3375        };
3376
3377        self_.report_initial_state();
3378        self_
3379    }
3380
3381    /// Reports the initial introspection state.
3382    fn report_initial_state(&self) {
3383        let row = self.write_frontier_row();
3384        let updates = vec![(row, Diff::ONE)];
3385        self.send(IntrospectionType::ReplicaFrontiers, updates);
3386    }
3387
3388    /// Observe the given current write frontier and update the introspection state as necessary.
3389    fn observe_frontier(&mut self, write_frontier: &Antichain<Timestamp>) {
3390        if self.write_frontier == *write_frontier {
3391            return; // no change
3392        }
3393
3394        let retraction = self.write_frontier_row();
3395        self.write_frontier.clone_from(write_frontier);
3396        let insertion = self.write_frontier_row();
3397
3398        let updates = vec![(retraction, Diff::MINUS_ONE), (insertion, Diff::ONE)];
3399        self.send(IntrospectionType::ReplicaFrontiers, updates);
3400    }
3401
3402    /// Return a `Row` reflecting the current replica write frontier.
3403    fn write_frontier_row(&self) -> Row {
3404        let write_frontier = self
3405            .write_frontier
3406            .as_option()
3407            .map_or(Datum::Null, |ts| ts.clone().into());
3408        Row::pack_slice(&[
3409            Datum::String(&self.collection_id.to_string()),
3410            Datum::String(&self.replica_id.to_string()),
3411            write_frontier,
3412        ])
3413    }
3414
3415    fn send(&self, introspection_type: IntrospectionType, updates: Vec<(Row, Diff)>) {
3416        // Failure to send means the `ComputeController` has been dropped and doesn't care about
3417        // introspection updates anymore.
3418        let _ = self.introspection_tx.send((introspection_type, updates));
3419    }
3420}
3421
3422impl Drop for ReplicaCollectionIntrospection {
3423    fn drop(&mut self) {
3424        // Retract the write frontier.
3425        let row = self.write_frontier_row();
3426        let updates = vec![(row, Diff::MINUS_ONE)];
3427        self.send(IntrospectionType::ReplicaFrontiers, updates);
3428    }
3429}
3430
3431#[cfg(test)]
3432mod tests {
3433    use std::collections::BTreeMap;
3434
3435    use mz_compute_types::dyncfgs::{ENABLE_COLUMN_PAGED_BATCHER, ENABLE_MZ_JOIN_CORE};
3436    use mz_dyncfg::{ConfigSet, ConfigUpdates, ConfigVal};
3437    use mz_persist_types::PersistLocation;
3438
3439    use crate::protocol::command::{ComputeCommand, InstanceConfig};
3440
3441    use super::{Instance, ReplicaId};
3442
3443    fn create_instance_command() -> ComputeCommand {
3444        ComputeCommand::CreateInstance(Box::new(InstanceConfig {
3445            logging: Default::default(),
3446            expiration_offset: None,
3447            peek_stash_persist_location: PersistLocation::new_in_mem(),
3448            arrangement_dictionary_compression: false,
3449            initial_config: Default::default(),
3450        }))
3451    }
3452
3453    fn initial_config(cmd: &ComputeCommand) -> &ConfigUpdates {
3454        match cmd {
3455            ComputeCommand::CreateInstance(config) => &config.initial_config,
3456            other => panic!("expected CreateInstance, got {other:?}"),
3457        }
3458    }
3459
3460    /// `CreateInstance` is specialized with a full snapshot of the instance-wide dyncfg, so the
3461    /// replica seeds its worker config at create time rather than waiting for the first
3462    /// `UpdateConfiguration`. This is the regression guard for create-time setup observing dyncfg
3463    /// defaults.
3464    #[mz_ore::test]
3465    fn create_instance_snapshots_instance_wide_dyncfg() {
3466        let dyncfg = ConfigSet::default()
3467            .add(&ENABLE_COLUMN_PAGED_BATCHER)
3468            .add(&ENABLE_MZ_JOIN_CORE);
3469        let mut updates = ConfigUpdates::default();
3470        updates.add(&ENABLE_COLUMN_PAGED_BATCHER, true);
3471        updates.add(&ENABLE_MZ_JOIN_CORE, false);
3472        updates.apply(&dyncfg);
3473
3474        // A replica without an override sees exactly the instance-wide values.
3475        let overrides = BTreeMap::new();
3476        let cmd = Instance::specialize_command_for_replica(
3477            create_instance_command(),
3478            ReplicaId::User(1),
3479            &overrides,
3480            &dyncfg,
3481        );
3482        let snapshot = initial_config(&cmd);
3483        assert_eq!(
3484            snapshot.updates.get(ENABLE_COLUMN_PAGED_BATCHER.name()),
3485            Some(&ConfigVal::Bool(true)),
3486        );
3487        assert_eq!(
3488            snapshot.updates.get(ENABLE_MZ_JOIN_CORE.name()),
3489            Some(&ConfigVal::Bool(false)),
3490        );
3491    }
3492
3493    /// A replica-scoped override beats the instance-wide value in the create-time snapshot, so a
3494    /// create-time-frozen scoped flag reaches the replica with its override applied.
3495    #[mz_ore::test]
3496    fn create_instance_snapshot_applies_replica_override() {
3497        let dyncfg = ConfigSet::default().add(&ENABLE_COLUMN_PAGED_BATCHER);
3498        let mut updates = ConfigUpdates::default();
3499        updates.add(&ENABLE_COLUMN_PAGED_BATCHER, true);
3500        updates.apply(&dyncfg);
3501
3502        let replica = ReplicaId::User(1);
3503        let mut override_updates = ConfigUpdates::default();
3504        override_updates.add(&ENABLE_COLUMN_PAGED_BATCHER, false);
3505        let overrides = BTreeMap::from([(replica, override_updates)]);
3506
3507        let cmd = Instance::specialize_command_for_replica(
3508            create_instance_command(),
3509            replica,
3510            &overrides,
3511            &dyncfg,
3512        );
3513        assert_eq!(
3514            initial_config(&cmd)
3515                .updates
3516                .get(ENABLE_COLUMN_PAGED_BATCHER.name()),
3517            Some(&ConfigVal::Bool(false)),
3518            "replica override should win over the instance-wide value",
3519        );
3520    }
3521
3522    /// `UpdateConfiguration` continues to merge the replica's override into the update.
3523    #[mz_ore::test]
3524    fn update_configuration_merges_replica_override() {
3525        let dyncfg = ConfigSet::default().add(&ENABLE_COLUMN_PAGED_BATCHER);
3526
3527        let replica = ReplicaId::User(1);
3528        let mut override_updates = ConfigUpdates::default();
3529        override_updates.add(&ENABLE_COLUMN_PAGED_BATCHER, true);
3530        let overrides = BTreeMap::from([(replica, override_updates)]);
3531
3532        let cmd = Instance::specialize_command_for_replica(
3533            ComputeCommand::UpdateConfiguration(Box::new(Default::default())),
3534            replica,
3535            &overrides,
3536            &dyncfg,
3537        );
3538        match cmd {
3539            ComputeCommand::UpdateConfiguration(params) => assert_eq!(
3540                params
3541                    .dyncfg_updates
3542                    .updates
3543                    .get(ENABLE_COLUMN_PAGED_BATCHER.name()),
3544                Some(&ConfigVal::Bool(true)),
3545            ),
3546            other => panic!("expected UpdateConfiguration, got {other:?}"),
3547        }
3548    }
3549}