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