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                ComputeSinkConnection::MetricSink(conn) => ComputeSinkConnection::MetricSink(conn),
1549            };
1550            let desc = ComputeSinkDesc {
1551                from: se.from,
1552                from_desc: se.from_desc,
1553                connection,
1554                with_snapshot: se.with_snapshot,
1555                up_to: se.up_to,
1556                non_null_assertions: se.non_null_assertions,
1557                refresh_schedule: se.refresh_schedule,
1558            };
1559            sink_exports.insert(id, desc);
1560        }
1561
1562        // Flatten the dataflow plans into the representation expected by replicas.
1563        let objects_to_build = dataflow
1564            .objects_to_build
1565            .into_iter()
1566            .map(|object| BuildDesc {
1567                id: object.id,
1568                plan: RenderPlan::try_from(object.plan).expect("valid plan"),
1569            })
1570            .collect();
1571
1572        let augmented_dataflow = DataflowDescription {
1573            source_imports,
1574            sink_exports,
1575            objects_to_build,
1576            // The rest of the fields are identical
1577            index_imports: dataflow.index_imports,
1578            index_exports: dataflow.index_exports,
1579            as_of: dataflow.as_of.clone(),
1580            until: dataflow.until,
1581            initial_storage_as_of: dataflow.initial_storage_as_of,
1582            refresh_schedule: dataflow.refresh_schedule,
1583            debug_name: dataflow.debug_name,
1584            time_dependence: dataflow.time_dependence,
1585        };
1586
1587        if augmented_dataflow.is_transient() {
1588            tracing::debug!(
1589                name = %augmented_dataflow.debug_name,
1590                import_ids = %augmented_dataflow.display_import_ids(),
1591                export_ids = %augmented_dataflow.display_export_ids(),
1592                as_of = ?augmented_dataflow.as_of.as_ref().unwrap().elements(),
1593                until = ?augmented_dataflow.until.elements(),
1594                "creating dataflow",
1595            );
1596        } else {
1597            tracing::info!(
1598                name = %augmented_dataflow.debug_name,
1599                import_ids = %augmented_dataflow.display_import_ids(),
1600                export_ids = %augmented_dataflow.display_export_ids(),
1601                as_of = ?augmented_dataflow.as_of.as_ref().unwrap().elements(),
1602                until = ?augmented_dataflow.until.elements(),
1603                "creating dataflow",
1604            );
1605        }
1606
1607        // Skip the actual dataflow creation for an empty `as_of`. (Happens e.g. for the
1608        // bootstrapping of a REFRESH AT mat view that is past its last refresh.)
1609        if as_of.is_empty() {
1610            tracing::info!(
1611                name = %augmented_dataflow.debug_name,
1612                "not sending `CreateDataflow`, because of empty `as_of`",
1613            );
1614        } else {
1615            let collections: Vec<_> = augmented_dataflow.export_ids().collect();
1616            self.send(ComputeCommand::CreateDataflow(Box::new(augmented_dataflow)));
1617
1618            for id in collections {
1619                self.maybe_schedule_collection(id);
1620            }
1621        }
1622
1623        Ok(())
1624    }
1625
1626    /// Schedule the identified collection if all its inputs are available.
1627    ///
1628    /// # Panics
1629    ///
1630    /// Panics if the identified collection does not exist.
1631    fn maybe_schedule_collection(&mut self, id: GlobalId) {
1632        let collection = self.expect_collection(id);
1633
1634        // Don't schedule collections twice.
1635        if collection.scheduled {
1636            return;
1637        }
1638
1639        let as_of = collection.read_frontier();
1640
1641        // If the collection has an empty `as_of`, it was either never installed on the replica or
1642        // has since been dropped. In either case the replica does not expect any commands for it.
1643        if as_of.is_empty() {
1644            return;
1645        }
1646
1647        let ready = if id.is_transient() {
1648            // Always schedule transient collections immediately. The assumption is that those are
1649            // created by interactive user commands and we want to schedule them as quickly as
1650            // possible. Inputs might not yet be available, but when they become available, we
1651            // don't need to wait for the controller to become aware and for the scheduling check
1652            // to run again.
1653            true
1654        } else {
1655            // Ignore self-dependencies. Any self-dependencies do not need to be
1656            // available at the as_of for the dataflow to make progress, so we
1657            // can ignore them here. At the moment, only continual tasks have
1658            // self-dependencies, but this logic is correct for any dataflow, so
1659            // we don't special case it to CTs.
1660            let not_self_dep = |x: &GlobalId| *x != id;
1661
1662            // Make sure we never schedule a collection before its input compute collections have
1663            // been scheduled. Scheduling in the wrong order can lead to deadlocks.
1664            let mut deps_scheduled = true;
1665
1666            // Check dependency frontiers to determine if all inputs are
1667            // available. An input is available when its frontier is greater
1668            // than the `as_of`, i.e., all input data up to and including the
1669            // `as_of` has been sealed.
1670            let compute_deps = collection.compute_dependency_ids().filter(not_self_dep);
1671            let mut compute_frontiers = Vec::new();
1672            for id in compute_deps {
1673                let dep = &self.expect_collection(id);
1674                deps_scheduled &= dep.scheduled;
1675                compute_frontiers.push(dep.write_frontier());
1676            }
1677
1678            let storage_deps = collection.storage_dependency_ids().filter(not_self_dep);
1679            let storage_frontiers = self
1680                .storage_collections
1681                .collections_frontiers(storage_deps.collect())
1682                .expect("must exist");
1683            let storage_frontiers = storage_frontiers.into_iter().map(|f| f.write_frontier);
1684
1685            let mut frontiers = compute_frontiers.into_iter().chain(storage_frontiers);
1686            let frontiers_ready =
1687                frontiers.all(|frontier| PartialOrder::less_than(&as_of, &frontier));
1688
1689            deps_scheduled && frontiers_ready
1690        };
1691
1692        if ready {
1693            self.send(ComputeCommand::Schedule(id));
1694            let collection = self.expect_collection_mut(id);
1695            collection.scheduled = true;
1696        }
1697    }
1698
1699    /// Schedule any unscheduled collections that are ready.
1700    fn schedule_collections(&mut self) {
1701        let ids: Vec<_> = self.collections.keys().copied().collect();
1702        for id in ids {
1703            self.maybe_schedule_collection(id);
1704        }
1705    }
1706
1707    /// Drops the read capability for the given collections and allows their resources to be
1708    /// reclaimed.
1709    #[mz_ore::instrument(level = "debug")]
1710    pub fn drop_collections(&mut self, ids: Vec<GlobalId>) -> Result<(), CollectionMissing> {
1711        for id in &ids {
1712            let collection = self.collection_mut(*id)?;
1713
1714            // Mark the collection as dropped to allow it to be removed from the controller state.
1715            collection.dropped = true;
1716
1717            // Drop the implied and warmup read holds to announce that clients are not
1718            // interested in the collection anymore.
1719            collection.implied_read_hold.release();
1720            collection.warmup_read_hold.release();
1721
1722            // If the collection is a subscribe, stop tracking it. This ensures that the controller
1723            // ceases to produce `SubscribeResponse`s for this subscribe.
1724            self.subscribes.remove(id);
1725            // If the collection is a copy to, stop tracking it. This ensures that the controller
1726            // ceases to produce `CopyToResponse`s` for this copy to.
1727            self.copy_tos.remove(id);
1728        }
1729
1730        Ok(())
1731    }
1732
1733    /// Initiate a peek request for the contents of `id` at `timestamp`.
1734    ///
1735    /// If this returns an error, then it didn't modify any `Instance` state.
1736    #[mz_ore::instrument(level = "debug")]
1737    pub fn peek(
1738        &mut self,
1739        peek_target: PeekTarget,
1740        literal_constraints: Option<Vec<Row>>,
1741        uuid: Uuid,
1742        timestamp: Timestamp,
1743        result_desc: RelationDesc,
1744        finishing: RowSetFinishing,
1745        map_filter_project: mz_expr::SafeMfpPlan,
1746        mut read_hold: ReadHold,
1747        target_replica: Option<ReplicaId>,
1748        peek_response_tx: oneshot::Sender<PeekResponse>,
1749    ) -> Result<(), PeekError> {
1750        use PeekError::*;
1751
1752        let target_id = peek_target.id();
1753
1754        // Downgrade the provided read hold to the peek time.
1755        if read_hold.id() != target_id {
1756            return Err(ReadHoldIdMismatch(read_hold.id()));
1757        }
1758        read_hold
1759            .try_downgrade(Antichain::from_elem(timestamp.clone()))
1760            .map_err(|_| ReadHoldInsufficient(target_id))?;
1761
1762        if let Some(target) = target_replica {
1763            if !self.replica_exists(target) {
1764                return Err(ReplicaMissing(target));
1765            }
1766        }
1767
1768        let otel_ctx = OpenTelemetryContext::obtain();
1769
1770        self.peeks.insert(
1771            uuid,
1772            PendingPeek {
1773                target_replica,
1774                // TODO(guswynn): can we just hold the `tracing::Span` here instead?
1775                otel_ctx: otel_ctx.clone(),
1776                requested_at: Instant::now(),
1777                read_hold,
1778                peek_response_tx,
1779                limit: finishing.limit.map(usize::cast_from),
1780                offset: finishing.offset,
1781            },
1782        );
1783
1784        let peek = Peek {
1785            literal_constraints,
1786            uuid,
1787            timestamp,
1788            finishing,
1789            map_filter_project,
1790            // Obtain an `OpenTelemetryContext` from the thread-local tracing
1791            // tree to forward it on to the compute worker.
1792            otel_ctx,
1793            target: peek_target,
1794            result_desc,
1795        };
1796        self.send(ComputeCommand::Peek(Box::new(peek)));
1797
1798        Ok(())
1799    }
1800
1801    /// Cancels an existing peek request.
1802    #[mz_ore::instrument(level = "debug")]
1803    pub fn cancel_peek(&mut self, uuid: Uuid, reason: PeekResponse) {
1804        let Some(peek) = self.peeks.get_mut(&uuid) else {
1805            tracing::warn!("did not find pending peek for {uuid}");
1806            return;
1807        };
1808
1809        let duration = peek.requested_at.elapsed();
1810        self.metrics
1811            .observe_peek_response(&PeekResponse::Canceled, duration);
1812
1813        // Enqueue a notification for the cancellation.
1814        let otel_ctx = peek.otel_ctx.clone();
1815        otel_ctx.attach_as_parent();
1816
1817        self.deliver_response(ComputeControllerResponse::PeekNotification(
1818            uuid,
1819            PeekNotification::Canceled,
1820            otel_ctx,
1821        ));
1822
1823        // Finish the peek.
1824        // This will also propagate the cancellation to the replicas.
1825        self.finish_peek(uuid, reason);
1826    }
1827
1828    /// Assigns a read policy to specific identifiers.
1829    ///
1830    /// The policies are assigned in the order presented, and repeated identifiers should
1831    /// conclude with the last policy. Changing a policy will immediately downgrade the read
1832    /// capability if appropriate, but it will not "recover" the read capability if the prior
1833    /// capability is already ahead of it.
1834    ///
1835    /// Identifiers not present in `policies` retain their existing read policies.
1836    ///
1837    /// It is an error to attempt to set a read policy for a collection that is not readable in the
1838    /// context of compute. At this time, only indexes are readable compute collections.
1839    #[mz_ore::instrument(level = "debug")]
1840    pub fn set_read_policy(
1841        &mut self,
1842        policies: Vec<(GlobalId, ReadPolicy)>,
1843    ) -> Result<(), ReadPolicyError> {
1844        // Do error checking upfront, to avoid introducing inconsistencies between a collection's
1845        // `implied_capability` and `read_capabilities`.
1846        for (id, _policy) in &policies {
1847            let collection = self.collection(*id)?;
1848            if collection.read_policy.is_none() {
1849                return Err(ReadPolicyError::WriteOnlyCollection(*id));
1850            }
1851        }
1852
1853        for (id, new_policy) in policies {
1854            let collection = self.expect_collection_mut(id);
1855            let new_since = new_policy.frontier(collection.write_frontier().borrow());
1856            let _ = collection.implied_read_hold.try_downgrade(new_since);
1857            collection.read_policy = Some(new_policy);
1858        }
1859
1860        Ok(())
1861    }
1862
1863    /// Advance the global write frontier of the given collection.
1864    ///
1865    /// Frontier regressions are gracefully ignored.
1866    ///
1867    /// # Panics
1868    ///
1869    /// Panics if the identified collection does not exist.
1870    #[mz_ore::instrument(level = "debug")]
1871    fn maybe_update_global_write_frontier(
1872        &mut self,
1873        id: GlobalId,
1874        new_frontier: Antichain<Timestamp>,
1875    ) {
1876        let collection = self.expect_collection_mut(id);
1877
1878        let advanced = collection.shared.lock_write_frontier(|f| {
1879            let advanced = PartialOrder::less_than(f, &new_frontier);
1880            if advanced {
1881                f.clone_from(&new_frontier);
1882            }
1883            advanced
1884        });
1885
1886        if !advanced {
1887            return;
1888        }
1889
1890        // Relax the implied read hold according to the read policy.
1891        let new_since = match &collection.read_policy {
1892            Some(read_policy) => {
1893                // For readable collections the read frontier is determined by applying the
1894                // client-provided read policy to the write frontier.
1895                read_policy.frontier(new_frontier.borrow())
1896            }
1897            None => {
1898                // Write-only collections cannot be read within the context of the compute
1899                // controller, so their read frontier only controls the read holds taken on their
1900                // inputs. We can safely downgrade the input read holds to any time less than the
1901                // write frontier.
1902                //
1903                // Note that some write-only collections (continual tasks) need to observe changes
1904                // at their current write frontier during hydration. Thus, we cannot downgrade the
1905                // read frontier to the write frontier and instead step it back by one.
1906                Antichain::from_iter(
1907                    new_frontier
1908                        .iter()
1909                        .map(|t| t.step_back().unwrap_or(Timestamp::MIN)),
1910                )
1911            }
1912        };
1913        let _ = collection.implied_read_hold.try_downgrade(new_since);
1914
1915        // Report the frontier advancement.
1916        self.deliver_response(ComputeControllerResponse::FrontierUpper {
1917            id,
1918            upper: new_frontier,
1919        });
1920    }
1921
1922    /// Apply a collection read hold change.
1923    pub(super) fn apply_read_hold_change(
1924        &mut self,
1925        id: GlobalId,
1926        mut update: ChangeBatch<Timestamp>,
1927    ) {
1928        let Some(collection) = self.collections.get_mut(&id) else {
1929            soft_panic_or_log!(
1930                "read hold change for absent collection (id={id}, changes={update:?})"
1931            );
1932            return;
1933        };
1934
1935        let new_since = collection.shared.lock_read_capabilities(|caps| {
1936            // Sanity check to prevent corrupted `read_capabilities`, which can cause hard-to-debug
1937            // issues (usually stuck read frontiers).
1938            let read_frontier = caps.frontier();
1939            for (time, diff) in update.iter() {
1940                let count = caps.count_for(time) + diff;
1941                assert!(
1942                    count >= 0,
1943                    "invalid read capabilities update: negative capability \
1944             (id={id:?}, read_capabilities={caps:?}, update={update:?})",
1945                );
1946                assert!(
1947                    count == 0 || read_frontier.less_equal(time),
1948                    "invalid read capabilities update: frontier regression \
1949             (id={id:?}, read_capabilities={caps:?}, update={update:?})",
1950                );
1951            }
1952
1953            // Apply read capability updates and learn about resulting changes to the read
1954            // frontier.
1955            let changes = caps.update_iter(update.drain());
1956
1957            let changed = changes.count() > 0;
1958            changed.then(|| caps.frontier().to_owned())
1959        });
1960
1961        let Some(new_since) = new_since else {
1962            return; // read frontier did not change
1963        };
1964
1965        // Propagate read frontier update to dependencies.
1966        for read_hold in collection.compute_dependencies.values_mut() {
1967            read_hold
1968                .try_downgrade(new_since.clone())
1969                .expect("frontiers don't regress");
1970        }
1971        for read_hold in collection.storage_dependencies.values_mut() {
1972            read_hold
1973                .try_downgrade(new_since.clone())
1974                .expect("frontiers don't regress");
1975        }
1976
1977        // Produce `AllowCompaction` command.
1978        self.send(ComputeCommand::AllowCompaction {
1979            id,
1980            frontier: new_since,
1981        });
1982    }
1983
1984    /// Fulfills a registered peek and cleans up associated state.
1985    ///
1986    /// As part of this we:
1987    ///  * Send a `PeekResponse` through the peek's response channel.
1988    ///  * Emit a `CancelPeek` command to instruct replicas to stop spending resources on this
1989    ///    peek, and to allow the `ComputeCommandHistory` to reduce away the corresponding `Peek`
1990    ///    command.
1991    ///  * Remove the read hold for this peek, unblocking compaction that might have waited on it.
1992    fn finish_peek(&mut self, uuid: Uuid, response: PeekResponse) {
1993        let Some(peek) = self.peeks.remove(&uuid) else {
1994            return;
1995        };
1996
1997        // The recipient might not be interested in the peek response anymore, which is fine.
1998        let _ = peek.peek_response_tx.send(response);
1999
2000        // NOTE: We need to send the `CancelPeek` command _before_ we release the peek's read hold
2001        // (by dropping it), to avoid the edge case that caused database-issues#4812.
2002        self.send(ComputeCommand::CancelPeek { uuid });
2003
2004        drop(peek.read_hold);
2005    }
2006
2007    /// Handles a response from a replica. Replica IDs are re-used across replica restarts, so we
2008    /// use the replica epoch to drop stale responses.
2009    fn handle_response(&mut self, (replica_id, epoch, response): ReplicaResponse) {
2010        // Filter responses from non-existing or stale replicas.
2011        if self
2012            .replicas
2013            .get(&replica_id)
2014            .filter(|replica| replica.epoch == epoch)
2015            .is_none()
2016        {
2017            return;
2018        }
2019
2020        // Invariant: the replica exists and has the expected epoch.
2021
2022        match response {
2023            ComputeResponse::Frontiers(id, frontiers) => {
2024                self.handle_frontiers_response(id, frontiers, replica_id);
2025            }
2026            ComputeResponse::PeekResponse(uuid, peek_response, otel_ctx) => {
2027                self.handle_peek_response(uuid, peek_response, otel_ctx, replica_id);
2028            }
2029            ComputeResponse::CopyToResponse(id, response) => {
2030                self.handle_copy_to_response(id, response, replica_id);
2031            }
2032            ComputeResponse::SubscribeResponse(id, response) => {
2033                self.handle_subscribe_response(id, response, replica_id);
2034            }
2035            ComputeResponse::Status(response) => {
2036                self.handle_status_response(response, replica_id);
2037            }
2038        }
2039    }
2040
2041    /// Handle new frontiers, returning any compute response that needs to
2042    /// be sent to the client.
2043    fn handle_frontiers_response(
2044        &mut self,
2045        id: GlobalId,
2046        frontiers: FrontiersResponse,
2047        replica_id: ReplicaId,
2048    ) {
2049        if !self.collections.contains_key(&id) {
2050            soft_panic_or_log!(
2051                "frontiers update for an unknown collection \
2052                 (id={id}, replica_id={replica_id}, frontiers={frontiers:?})"
2053            );
2054            return;
2055        }
2056        let Some(replica) = self.replicas.get_mut(&replica_id) else {
2057            soft_panic_or_log!(
2058                "frontiers update for an unknown replica \
2059                 (replica_id={replica_id}, frontiers={frontiers:?})"
2060            );
2061            return;
2062        };
2063        let Some(replica_collection) = replica.collections.get_mut(&id) else {
2064            soft_panic_or_log!(
2065                "frontiers update for an unknown replica collection \
2066                 (id={id}, replica_id={replica_id}, frontiers={frontiers:?})"
2067            );
2068            return;
2069        };
2070
2071        if let Some(new_frontier) = frontiers.input_frontier {
2072            replica_collection.update_input_frontier(new_frontier.clone());
2073        }
2074        if let Some(new_frontier) = frontiers.output_frontier {
2075            replica_collection.update_output_frontier(new_frontier.clone());
2076        }
2077        if let Some(new_frontier) = frontiers.write_frontier {
2078            replica_collection.update_write_frontier(new_frontier.clone());
2079            self.maybe_update_global_write_frontier(id, new_frontier);
2080        }
2081    }
2082
2083    #[mz_ore::instrument(level = "debug")]
2084    fn handle_peek_response(
2085        &mut self,
2086        uuid: Uuid,
2087        response: PeekResponse,
2088        otel_ctx: OpenTelemetryContext,
2089        replica_id: ReplicaId,
2090    ) {
2091        otel_ctx.attach_as_parent();
2092
2093        // We might not be tracking this peek anymore, because we have served a response already or
2094        // because it was canceled. If this is the case, we ignore the response.
2095        let Some(peek) = self.peeks.get(&uuid) else {
2096            return;
2097        };
2098
2099        // If the peek is targeting a replica, ignore responses from other replicas.
2100        let target_replica = peek.target_replica.unwrap_or(replica_id);
2101        if target_replica != replica_id {
2102            return;
2103        }
2104
2105        let duration = peek.requested_at.elapsed();
2106        self.metrics.observe_peek_response(&response, duration);
2107
2108        let notification = PeekNotification::new(&response, peek.offset, peek.limit);
2109        // NOTE: We use the `otel_ctx` from the response, not the pending peek, because we
2110        // currently want the parent to be whatever the compute worker did with this peek.
2111        self.deliver_response(ComputeControllerResponse::PeekNotification(
2112            uuid,
2113            notification,
2114            otel_ctx,
2115        ));
2116
2117        self.finish_peek(uuid, response)
2118    }
2119
2120    fn handle_copy_to_response(
2121        &mut self,
2122        sink_id: GlobalId,
2123        response: CopyToResponse,
2124        replica_id: ReplicaId,
2125    ) {
2126        if !self.collections.contains_key(&sink_id) {
2127            soft_panic_or_log!(
2128                "received response for an unknown copy-to \
2129                 (sink_id={sink_id}, replica_id={replica_id})",
2130            );
2131            return;
2132        }
2133        let Some(replica) = self.replicas.get_mut(&replica_id) else {
2134            soft_panic_or_log!("copy-to response for an unknown replica (replica_id={replica_id})");
2135            return;
2136        };
2137        let Some(replica_collection) = replica.collections.get_mut(&sink_id) else {
2138            soft_panic_or_log!(
2139                "copy-to response for an unknown replica collection \
2140                 (sink_id={sink_id}, replica_id={replica_id})"
2141            );
2142            return;
2143        };
2144
2145        // Downgrade the replica frontiers, to enable dropping of input read holds and clean up of
2146        // collection state.
2147        // TODO(database-issues#4701): report copy-to frontiers through `Frontiers` responses
2148        replica_collection.update_write_frontier(Antichain::new());
2149        replica_collection.update_input_frontier(Antichain::new());
2150        replica_collection.update_output_frontier(Antichain::new());
2151
2152        // We might not be tracking this COPY TO because we have already returned a response
2153        // from one of the replicas. In that case, we ignore the response.
2154        if !self.copy_tos.remove(&sink_id) {
2155            return;
2156        }
2157
2158        let result = match response {
2159            CopyToResponse::RowCount(count) => Ok(count),
2160            CopyToResponse::Error(error) => Err(anyhow::anyhow!(error)),
2161            // We should never get here: Replicas only drop copy to collections in response
2162            // to the controller allowing them to do so, and when the controller drops a
2163            // copy to it also removes it from the list of tracked copy_tos (see
2164            // [`Instance::drop_collections`]).
2165            CopyToResponse::Dropped => {
2166                tracing::error!(
2167                    %sink_id, %replica_id,
2168                    "received `Dropped` response for a tracked copy to",
2169                );
2170                return;
2171            }
2172        };
2173
2174        self.deliver_response(ComputeControllerResponse::CopyToResponse(sink_id, result));
2175    }
2176
2177    fn handle_subscribe_response(
2178        &mut self,
2179        subscribe_id: GlobalId,
2180        response: SubscribeResponse,
2181        replica_id: ReplicaId,
2182    ) {
2183        if !self.collections.contains_key(&subscribe_id) {
2184            soft_panic_or_log!(
2185                "received response for an unknown subscribe \
2186                 (subscribe_id={subscribe_id}, replica_id={replica_id})",
2187            );
2188            return;
2189        }
2190        let Some(replica) = self.replicas.get_mut(&replica_id) else {
2191            soft_panic_or_log!(
2192                "subscribe response for an unknown replica (replica_id={replica_id})"
2193            );
2194            return;
2195        };
2196        let Some(replica_collection) = replica.collections.get_mut(&subscribe_id) else {
2197            soft_panic_or_log!(
2198                "subscribe response for an unknown replica collection \
2199                 (subscribe_id={subscribe_id}, replica_id={replica_id})"
2200            );
2201            return;
2202        };
2203
2204        // Always apply replica write frontier updates. Even if the subscribe is not tracked
2205        // anymore, there might still be replicas reading from its inputs, so we need to track the
2206        // frontiers until all replicas have advanced to the empty one.
2207        let write_frontier = match &response {
2208            SubscribeResponse::Batch(batch) => batch.upper.clone(),
2209            SubscribeResponse::DroppedAt(_) => Antichain::new(),
2210        };
2211
2212        // For subscribes we downgrade all replica frontiers based on write frontiers. This should
2213        // be fine because the input and output frontier of a subscribe track its write frontier.
2214        // TODO(database-issues#4701): report subscribe frontiers through `Frontiers` responses
2215        replica_collection.update_write_frontier(write_frontier.clone());
2216        replica_collection.update_input_frontier(write_frontier.clone());
2217        replica_collection.update_output_frontier(write_frontier.clone());
2218
2219        // If the subscribe is not tracked, or targets a different replica, there is nothing to do.
2220        let Some(mut subscribe) = self.subscribes.get(&subscribe_id).cloned() else {
2221            return;
2222        };
2223
2224        // Apply a global frontier update.
2225        // If this is a replica-targeted subscribe, it is important that we advance the global
2226        // frontier only based on responses from the targeted replica. Otherwise, another replica
2227        // could advance to the empty frontier, making us drop the subscribe on the targeted
2228        // replica prematurely.
2229        self.maybe_update_global_write_frontier(subscribe_id, write_frontier);
2230
2231        match response {
2232            SubscribeResponse::Batch(batch) => {
2233                let upper = batch.upper;
2234                let mut updates = batch.updates;
2235
2236                // If this batch advances the subscribe's frontier, we emit all updates at times
2237                // greater or equal to the last frontier (to avoid emitting duplicate updates).
2238                if PartialOrder::less_than(&subscribe.frontier, &upper) {
2239                    let lower = std::mem::replace(&mut subscribe.frontier, upper.clone());
2240
2241                    if upper.is_empty() {
2242                        // This subscribe cannot produce more data. Stop tracking it.
2243                        self.subscribes.remove(&subscribe_id);
2244                    } else {
2245                        // This subscribe can produce more data. Update our tracking of it.
2246                        self.subscribes.insert(subscribe_id, subscribe);
2247                    }
2248
2249                    if let Ok(updates) = updates.as_mut() {
2250                        updates.retain_mut(|updates| {
2251                            let offset = updates.times().partition_point(|t| {
2252                                // True for times that are strictly less than lower (and should be skipped)
2253                                // and false otherwise.
2254                                !lower.less_equal(t)
2255                            });
2256                            let (_, past_lower) = std::mem::take(updates).split_at(offset);
2257                            *updates = past_lower;
2258                            updates.len() > 0
2259                        });
2260                    }
2261                    self.deliver_response(ComputeControllerResponse::SubscribeResponse(
2262                        subscribe_id,
2263                        SubscribeBatch {
2264                            lower,
2265                            upper,
2266                            updates,
2267                        },
2268                    ));
2269                }
2270            }
2271            SubscribeResponse::DroppedAt(frontier) => {
2272                // We should never get here: Replicas only drop subscribe collections in response
2273                // to the controller allowing them to do so, and when the controller drops a
2274                // subscribe it also removes it from the list of tracked subscribes (see
2275                // [`Instance::drop_collections`]).
2276                tracing::error!(
2277                    %subscribe_id,
2278                    %replica_id,
2279                    frontier = ?frontier.elements(),
2280                    "received `DroppedAt` response for a tracked subscribe",
2281                );
2282                self.subscribes.remove(&subscribe_id);
2283            }
2284        }
2285    }
2286
2287    fn handle_status_response(&self, response: StatusResponse, _replica_id: ReplicaId) {
2288        match response {
2289            StatusResponse::Placeholder => {}
2290        }
2291    }
2292
2293    /// Return the write frontiers of the dependencies of the given collection.
2294    fn dependency_write_frontiers<'b>(
2295        &'b self,
2296        collection: &'b CollectionState,
2297    ) -> impl Iterator<Item = Antichain<Timestamp>> + 'b {
2298        let compute_frontiers = collection.compute_dependency_ids().filter_map(|dep_id| {
2299            let collection = self.collections.get(&dep_id);
2300            collection.map(|c| c.write_frontier())
2301        });
2302        let storage_frontiers = collection.storage_dependency_ids().filter_map(|dep_id| {
2303            let frontiers = self.storage_collections.collection_frontiers(dep_id).ok();
2304            frontiers.map(|f| f.write_frontier)
2305        });
2306
2307        compute_frontiers.chain(storage_frontiers)
2308    }
2309
2310    /// Return the write frontiers of transitive storage dependencies of the given collection.
2311    fn transitive_storage_dependency_write_frontiers<'b>(
2312        &'b self,
2313        collection: &'b CollectionState,
2314    ) -> impl Iterator<Item = Antichain<Timestamp>> + 'b {
2315        let mut storage_ids: BTreeSet<_> = collection.storage_dependency_ids().collect();
2316        let mut todo: Vec<_> = collection.compute_dependency_ids().collect();
2317        let mut done = BTreeSet::new();
2318
2319        while let Some(id) = todo.pop() {
2320            if done.contains(&id) {
2321                continue;
2322            }
2323            if let Some(dep) = self.collections.get(&id) {
2324                storage_ids.extend(dep.storage_dependency_ids());
2325                todo.extend(dep.compute_dependency_ids())
2326            }
2327            done.insert(id);
2328        }
2329
2330        let storage_frontiers = storage_ids.into_iter().filter_map(|id| {
2331            let frontiers = self.storage_collections.collection_frontiers(id).ok();
2332            frontiers.map(|f| f.write_frontier)
2333        });
2334
2335        storage_frontiers
2336    }
2337
2338    /// Downgrade the warmup capabilities of collections as much as possible.
2339    ///
2340    /// The only requirement we have for a collection's warmup capability is that it is for a time
2341    /// that is available in all of the collection's inputs. For each input the latest time that is
2342    /// the case for is `write_frontier - 1`. So the farthest we can downgrade a collection's
2343    /// warmup capability is the minimum of `write_frontier - 1` of all its inputs.
2344    ///
2345    /// This method expects to be periodically called as part of instance maintenance work.
2346    /// We would like to instead update the warmup capabilities synchronously in response to
2347    /// frontier updates of dependency collections, but that is not generally possible because we
2348    /// don't learn about frontier updates of storage collections synchronously. We could do
2349    /// synchronous updates for compute dependencies, but we refrain from doing for simplicity.
2350    fn downgrade_warmup_capabilities(&mut self) {
2351        let mut new_capabilities = BTreeMap::new();
2352        for (id, collection) in &self.collections {
2353            // For write-only collections that have advanced to the empty frontier, we can drop the
2354            // warmup capability entirely. There is no reason why we would need to hydrate those
2355            // collections again, so being able to warm them up is not useful.
2356            if collection.read_policy.is_none()
2357                && collection.shared.lock_write_frontier(|f| f.is_empty())
2358            {
2359                new_capabilities.insert(*id, Antichain::new());
2360                continue;
2361            }
2362
2363            let mut new_capability = Antichain::new();
2364            for frontier in self.dependency_write_frontiers(collection) {
2365                for time in frontier {
2366                    new_capability.insert(time.step_back().unwrap_or(time));
2367                }
2368            }
2369
2370            new_capabilities.insert(*id, new_capability);
2371        }
2372
2373        for (id, new_capability) in new_capabilities {
2374            let collection = self.expect_collection_mut(id);
2375            let _ = collection.warmup_read_hold.try_downgrade(new_capability);
2376        }
2377    }
2378
2379    /// Forward the implied capabilities of collections, if possible.
2380    ///
2381    /// The implied capability of a collection controls (a) which times are still readable (for
2382    /// indexes) and (b) with which as-of the collection gets installed on a new replica. We are
2383    /// usually not allowed to advance an implied capability beyond the frontier that follows from
2384    /// the collection's read policy applied to its write frontier:
2385    ///
2386    ///  * For sink collections, some external consumer might rely on seeing all distinct times in
2387    ///    the input reflected in the output. If we'd forward the implied capability of a sink,
2388    ///    we'd risk skipping times in the output across replica restarts.
2389    ///  * For index collections, we might make the index unreadable by advancing its read frontier
2390    ///    beyond its write frontier.
2391    ///
2392    /// There is one case where forwarding an implied capability is fine though: an index installed
2393    /// on a cluster that has no replicas. Such indexes are not readable anyway until a new replica
2394    /// is added, so advancing its read frontier can't make it unreadable. We can thus advance the
2395    /// implied capability as long as we make sure that when a new replica is added, the expected
2396    /// relationship between write frontier, read policy, and implied capability can be restored
2397    /// immediately (modulo computation time).
2398    ///
2399    /// Forwarding implied capabilities is not necessary for the correct functioning of the
2400    /// controller but an optimization that is beneficial in two ways:
2401    ///
2402    ///  * It relaxes read holds on inputs to forwarded collections, allowing their compaction.
2403    ///  * It reduces the amount of historical detail new replicas need to process when computing
2404    ///    forwarded collections, as forwarding the implied capability also forwards the corresponding
2405    ///    dataflow as-of.
2406    fn forward_implied_capabilities(&mut self) {
2407        if !ENABLE_PAUSED_CLUSTER_READHOLD_DOWNGRADE.get(&self.dyncfg) {
2408            return;
2409        }
2410        if !self.replicas.is_empty() {
2411            return;
2412        }
2413
2414        let mut new_capabilities = BTreeMap::new();
2415        for (id, collection) in &self.collections {
2416            let Some(read_policy) = &collection.read_policy else {
2417                // Collection is write-only, i.e. a sink.
2418                continue;
2419            };
2420
2421            // When a new replica is started, it will immediately be able to compute all collection
2422            // output up to the write frontier of its transitive storage inputs. So the new implied
2423            // read capability should be the read policy applied to that frontier.
2424            let mut dep_frontier = Antichain::new();
2425            for frontier in self.transitive_storage_dependency_write_frontiers(collection) {
2426                dep_frontier.extend(frontier);
2427            }
2428
2429            let new_capability = read_policy.frontier(dep_frontier.borrow());
2430            if PartialOrder::less_than(collection.implied_read_hold.since(), &new_capability) {
2431                new_capabilities.insert(*id, new_capability);
2432            }
2433        }
2434
2435        for (id, new_capability) in new_capabilities {
2436            let collection = self.expect_collection_mut(id);
2437            let _ = collection.implied_read_hold.try_downgrade(new_capability);
2438        }
2439    }
2440
2441    /// Acquires a `ReadHold` for the identified compute collection.
2442    ///
2443    /// This mirrors the logic used by the controller-side `InstanceState::acquire_read_hold`,
2444    /// but executes on the instance task itself.
2445    pub(super) fn acquire_read_hold(&self, id: GlobalId) -> Result<ReadHold, CollectionMissing> {
2446        // Similarly to InstanceState::acquire_read_hold and StorageCollections::acquire_read_holds,
2447        // we acquire read holds at the earliest possible time rather than returning a copy
2448        // of the implied read hold. This is so that dependents can acquire read holds on
2449        // compute dependencies at frontiers that are held back by other read holds the caller
2450        // has previously taken.
2451        let collection = self.collection(id)?;
2452        let since = collection.shared.lock_read_capabilities(|caps| {
2453            let since = caps.frontier().to_owned();
2454            caps.update_iter(since.iter().map(|t| (t.clone(), 1)));
2455            since
2456        });
2457        let hold = ReadHold::new(id, since, Arc::clone(&self.read_hold_tx));
2458        Ok(hold)
2459    }
2460
2461    /// Process pending maintenance work.
2462    ///
2463    /// This method is invoked periodically by the global controller.
2464    /// It is a good place to perform maintenance work that arises from various controller state
2465    /// changes and that cannot conveniently be handled synchronously with those state changes.
2466    #[mz_ore::instrument(level = "debug")]
2467    pub fn maintain(&mut self) {
2468        self.rehydrate_failed_replicas();
2469        self.downgrade_warmup_capabilities();
2470        self.forward_implied_capabilities();
2471        self.schedule_collections();
2472        self.cleanup_collections();
2473        self.update_frontier_introspection();
2474        self.refresh_state_metrics();
2475        self.refresh_wallclock_lag();
2476    }
2477}
2478
2479/// State maintained about individual compute collections.
2480///
2481/// A compute collection is either an index, or a storage sink, or a subscribe, exported by a
2482/// compute dataflow.
2483#[derive(Debug)]
2484struct CollectionState {
2485    /// If set, this collection is only maintained by the specified replica.
2486    target_replica: Option<ReplicaId>,
2487    /// Whether this collection is a log collection.
2488    ///
2489    /// Log collections are special in that they are only maintained by a subset of all replicas.
2490    log_collection: bool,
2491    /// Whether this collection has been dropped by a controller client.
2492    ///
2493    /// The controller is allowed to remove the `CollectionState` for a collection only when
2494    /// `dropped == true`. Otherwise, clients might still expect to be able to query information
2495    /// about this collection.
2496    dropped: bool,
2497    /// Whether this collection has been scheduled, i.e., the controller has sent a `Schedule`
2498    /// command for it.
2499    scheduled: bool,
2500
2501    /// Whether this collection is in read-only mode.
2502    ///
2503    /// When in read-only mode, the dataflow is not allowed to affect external state (largely persist).
2504    read_only: bool,
2505
2506    /// State shared with the `ComputeController`.
2507    shared: SharedCollectionState,
2508
2509    /// A read hold maintaining the implicit capability of the collection.
2510    ///
2511    /// This capability is kept to ensure that the collection remains readable according to its
2512    /// `read_policy`. It also ensures that read holds on the collection's dependencies are kept at
2513    /// some time not greater than the collection's `write_frontier`, guaranteeing that the
2514    /// collection's next outputs can always be computed without skipping times.
2515    implied_read_hold: ReadHold,
2516    /// A read hold held to enable dataflow warmup.
2517    ///
2518    /// Dataflow warmup is an optimization that allows dataflows to immediately start hydrating
2519    /// even when their next output time (as implied by the `write_frontier`) is in the future.
2520    /// By installing a read capability derived from the write frontiers of the collection's
2521    /// inputs, we ensure that the as-of of new dataflows installed for the collection is at a time
2522    /// that is immediately available, so hydration can begin immediately too.
2523    warmup_read_hold: ReadHold,
2524    /// The policy to use to downgrade `self.implied_read_hold`.
2525    ///
2526    /// If `None`, the collection is a write-only collection (i.e. a sink). For write-only
2527    /// collections, the `implied_read_hold` is only required for maintaining read holds on the
2528    /// inputs, so we can immediately downgrade it to the `write_frontier`.
2529    read_policy: Option<ReadPolicy>,
2530
2531    /// Storage identifiers on which this collection depends, and read holds this collection
2532    /// requires on them.
2533    storage_dependencies: BTreeMap<GlobalId, ReadHold>,
2534    /// Compute identifiers on which this collection depends, and read holds this collection
2535    /// requires on them.
2536    compute_dependencies: BTreeMap<GlobalId, ReadHold>,
2537
2538    /// Introspection state associated with this collection.
2539    introspection: CollectionIntrospection,
2540
2541    /// Frontier wallclock lag measurements stashed until the next `WallclockLagHistogram`
2542    /// introspection update.
2543    ///
2544    /// Keys are `(period, lag, labels)` triples, values are counts.
2545    ///
2546    /// If this is `None`, wallclock lag is not tracked for this collection.
2547    wallclock_lag_histogram_stash: Option<
2548        BTreeMap<
2549            (
2550                WallclockLagHistogramPeriod,
2551                WallclockLag,
2552                BTreeMap<&'static str, String>,
2553            ),
2554            Diff,
2555        >,
2556    >,
2557}
2558
2559impl CollectionState {
2560    /// Creates a new collection state, with an initial read policy valid from `since`.
2561    fn new(
2562        collection_id: GlobalId,
2563        as_of: Antichain<Timestamp>,
2564        shared: SharedCollectionState,
2565        storage_dependencies: BTreeMap<GlobalId, ReadHold>,
2566        compute_dependencies: BTreeMap<GlobalId, ReadHold>,
2567        read_hold_tx: read_holds::ChangeTx,
2568        introspection: CollectionIntrospection,
2569    ) -> Self {
2570        // A collection is not readable before the `as_of`.
2571        let since = as_of.clone();
2572        // A collection won't produce updates for times before the `as_of`.
2573        let upper = as_of;
2574
2575        // Ensure that the provided `shared` is valid for the given `as_of`.
2576        assert!(shared.lock_read_capabilities(|c| c.frontier() == since.borrow()));
2577        assert!(shared.lock_write_frontier(|f| f == &upper));
2578
2579        // Initialize collection read holds.
2580        // Note that the implied read hold was already added to the `read_capabilities` when
2581        // `shared` was created, so we only need to add the warmup read hold here.
2582        let implied_read_hold =
2583            ReadHold::new(collection_id, since.clone(), Arc::clone(&read_hold_tx));
2584        let warmup_read_hold = ReadHold::new(collection_id, since.clone(), read_hold_tx);
2585
2586        let updates = warmup_read_hold.since().iter().map(|t| (t.clone(), 1));
2587        shared.lock_read_capabilities(|c| {
2588            c.update_iter(updates);
2589        });
2590
2591        // In an effort to keep the produced wallclock lag introspection data small and
2592        // predictable, we disable wallclock lag tracking for transient collections, i.e. slow-path
2593        // select indexes and subscribes.
2594        let wallclock_lag_histogram_stash = match collection_id.is_transient() {
2595            true => None,
2596            false => Some(Default::default()),
2597        };
2598
2599        Self {
2600            target_replica: None,
2601            log_collection: false,
2602            dropped: false,
2603            scheduled: false,
2604            read_only: true,
2605            shared,
2606            implied_read_hold,
2607            warmup_read_hold,
2608            read_policy: Some(ReadPolicy::ValidFrom(since)),
2609            storage_dependencies,
2610            compute_dependencies,
2611            introspection,
2612            wallclock_lag_histogram_stash,
2613        }
2614    }
2615
2616    /// Creates a new collection state for a log collection.
2617    fn new_log_collection(
2618        id: GlobalId,
2619        shared: SharedCollectionState,
2620        read_hold_tx: read_holds::ChangeTx,
2621        introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
2622    ) -> Self {
2623        let since = Antichain::from_elem(Timestamp::MIN);
2624        let introspection = CollectionIntrospection::new(
2625            id,
2626            introspection_tx,
2627            since.clone(),
2628            false,
2629            None,
2630            None,
2631            Vec::new(),
2632        );
2633        let mut state = Self::new(
2634            id,
2635            since,
2636            shared,
2637            Default::default(),
2638            Default::default(),
2639            read_hold_tx,
2640            introspection,
2641        );
2642        state.log_collection = true;
2643        // Log collections are created and scheduled implicitly as part of replica initialization.
2644        state.scheduled = true;
2645        state
2646    }
2647
2648    /// Reports the current read frontier.
2649    fn read_frontier(&self) -> Antichain<Timestamp> {
2650        self.shared
2651            .lock_read_capabilities(|c| c.frontier().to_owned())
2652    }
2653
2654    /// Reports the current write frontier.
2655    fn write_frontier(&self) -> Antichain<Timestamp> {
2656        self.shared.lock_write_frontier(|f| f.clone())
2657    }
2658
2659    fn storage_dependency_ids(&self) -> impl Iterator<Item = GlobalId> + '_ {
2660        self.storage_dependencies.keys().copied()
2661    }
2662
2663    fn compute_dependency_ids(&self) -> impl Iterator<Item = GlobalId> + '_ {
2664        self.compute_dependencies.keys().copied()
2665    }
2666}
2667
2668/// Collection state shared with the `ComputeController`.
2669///
2670/// Having this allows certain controller APIs, such as `ComputeController::collection_frontiers`
2671/// and `ComputeController::acquire_read_hold` to be non-`async`. This comes at the cost of
2672/// complexity (by introducing shared mutable state) and performance (by introducing locking). We
2673/// should aim to reduce the amount of shared state over time, rather than expand it.
2674///
2675/// Note that [`SharedCollectionState`]s are initialized by the `ComputeController` prior to the
2676/// collection's creation in the [`Instance`]. This is to allow compute clients to query frontiers
2677/// and take new read holds immediately, without having to wait for the [`Instance`] to update.
2678#[derive(Clone, Debug)]
2679pub(super) struct SharedCollectionState {
2680    /// Accumulation of read capabilities for the collection.
2681    ///
2682    /// This accumulation contains the capabilities held by all [`ReadHold`]s given out for the
2683    /// collection, including `implied_read_hold` and `warmup_read_hold`.
2684    ///
2685    /// NOTE: This field may only be modified by [`Instance::apply_read_hold_change`],
2686    /// [`Instance::acquire_read_hold`], and `ComputeController::acquire_read_hold`.
2687    /// Nobody else should modify read capabilities directly. Instead, collection users should
2688    /// manage read holds through [`ReadHold`] objects acquired through
2689    /// `ComputeController::acquire_read_hold`.
2690    ///
2691    /// TODO(teskje): Restructure the code to enforce the above in the type system.
2692    read_capabilities: Arc<Mutex<MutableAntichain<Timestamp>>>,
2693    /// The write frontier of this collection.
2694    write_frontier: Arc<Mutex<Antichain<Timestamp>>>,
2695}
2696
2697impl SharedCollectionState {
2698    pub fn new(as_of: Antichain<Timestamp>) -> Self {
2699        // A collection is not readable before the `as_of`.
2700        let since = as_of.clone();
2701        // A collection won't produce updates for times before the `as_of`.
2702        let upper = as_of;
2703
2704        // Initialize read capabilities to the `since`.
2705        // The is the implied read capability. The corresponding [`ReadHold`] is created in
2706        // [`CollectionState::new`].
2707        let mut read_capabilities = MutableAntichain::new();
2708        read_capabilities.update_iter(since.iter().map(|time| (time.clone(), 1)));
2709
2710        Self {
2711            read_capabilities: Arc::new(Mutex::new(read_capabilities)),
2712            write_frontier: Arc::new(Mutex::new(upper)),
2713        }
2714    }
2715
2716    pub fn lock_read_capabilities<F, R>(&self, f: F) -> R
2717    where
2718        F: FnOnce(&mut MutableAntichain<Timestamp>) -> R,
2719    {
2720        let mut caps = self.read_capabilities.lock().expect("poisoned");
2721        f(&mut *caps)
2722    }
2723
2724    pub fn lock_write_frontier<F, R>(&self, f: F) -> R
2725    where
2726        F: FnOnce(&mut Antichain<Timestamp>) -> R,
2727    {
2728        let mut frontier = self.write_frontier.lock().expect("poisoned");
2729        f(&mut *frontier)
2730    }
2731}
2732
2733/// Manages certain introspection relations associated with a collection. Upon creation, it adds
2734/// rows to introspection relations. When dropped, it retracts its managed rows.
2735#[derive(Debug)]
2736struct CollectionIntrospection {
2737    /// The ID of the compute collection.
2738    collection_id: GlobalId,
2739    /// A channel through which introspection updates are delivered.
2740    introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
2741    /// Introspection state for `IntrospectionType::Frontiers`.
2742    ///
2743    /// `Some` if the collection does _not_ sink into a storage collection (i.e. is not an MV). If
2744    /// the collection sinks into storage, the storage controller reports its frontiers instead.
2745    frontiers: Option<FrontiersIntrospectionState>,
2746    /// Introspection state for `IntrospectionType::ComputeMaterializedViewRefreshes`.
2747    ///
2748    /// `Some` if the collection is a REFRESH MV.
2749    refresh: Option<RefreshIntrospectionState>,
2750    /// The IDs of the collection's dependencies, for `IntrospectionType::ComputeDependencies`.
2751    dependency_ids: Vec<GlobalId>,
2752}
2753
2754impl CollectionIntrospection {
2755    fn new(
2756        collection_id: GlobalId,
2757        introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
2758        as_of: Antichain<Timestamp>,
2759        storage_sink: bool,
2760        initial_as_of: Option<Antichain<Timestamp>>,
2761        refresh_schedule: Option<RefreshSchedule>,
2762        dependency_ids: Vec<GlobalId>,
2763    ) -> Self {
2764        let refresh =
2765            match (refresh_schedule, initial_as_of) {
2766                (Some(refresh_schedule), Some(initial_as_of)) => Some(
2767                    RefreshIntrospectionState::new(refresh_schedule, initial_as_of, &as_of),
2768                ),
2769                (refresh_schedule, _) => {
2770                    // If we have a `refresh_schedule`, then the collection is a MV, so we should also have
2771                    // an `initial_as_of`.
2772                    soft_assert_or_log!(
2773                        refresh_schedule.is_none(),
2774                        "`refresh_schedule` without an `initial_as_of`: {collection_id}"
2775                    );
2776                    None
2777                }
2778            };
2779        let frontiers = (!storage_sink).then(|| FrontiersIntrospectionState::new(as_of));
2780
2781        let self_ = Self {
2782            collection_id,
2783            introspection_tx,
2784            frontiers,
2785            refresh,
2786            dependency_ids,
2787        };
2788
2789        self_.report_initial_state();
2790        self_
2791    }
2792
2793    /// Reports the initial introspection state.
2794    fn report_initial_state(&self) {
2795        if let Some(frontiers) = &self.frontiers {
2796            let row = frontiers.row_for_collection(self.collection_id);
2797            let updates = vec![(row, Diff::ONE)];
2798            self.send(IntrospectionType::Frontiers, updates);
2799        }
2800
2801        if let Some(refresh) = &self.refresh {
2802            let row = refresh.row_for_collection(self.collection_id);
2803            let updates = vec![(row, Diff::ONE)];
2804            self.send(IntrospectionType::ComputeMaterializedViewRefreshes, updates);
2805        }
2806
2807        if !self.dependency_ids.is_empty() {
2808            let updates = self.dependency_rows(Diff::ONE);
2809            self.send(IntrospectionType::ComputeDependencies, updates);
2810        }
2811    }
2812
2813    /// Produces rows for the `ComputeDependencies` introspection relation.
2814    fn dependency_rows(&self, diff: Diff) -> Vec<(Row, Diff)> {
2815        self.dependency_ids
2816            .iter()
2817            .map(|dependency_id| {
2818                let row = Row::pack_slice(&[
2819                    Datum::String(&self.collection_id.to_string()),
2820                    Datum::String(&dependency_id.to_string()),
2821                ]);
2822                (row, diff)
2823            })
2824            .collect()
2825    }
2826
2827    /// Observe the given current collection frontiers and update the introspection state as
2828    /// necessary.
2829    fn observe_frontiers(
2830        &mut self,
2831        read_frontier: &Antichain<Timestamp>,
2832        write_frontier: &Antichain<Timestamp>,
2833    ) {
2834        self.update_frontier_introspection(read_frontier, write_frontier);
2835        self.update_refresh_introspection(write_frontier);
2836    }
2837
2838    fn update_frontier_introspection(
2839        &mut self,
2840        read_frontier: &Antichain<Timestamp>,
2841        write_frontier: &Antichain<Timestamp>,
2842    ) {
2843        let Some(frontiers) = &mut self.frontiers else {
2844            return;
2845        };
2846
2847        if &frontiers.read_frontier == read_frontier && &frontiers.write_frontier == write_frontier
2848        {
2849            return; // no change
2850        };
2851
2852        let retraction = frontiers.row_for_collection(self.collection_id);
2853        frontiers.update(read_frontier, write_frontier);
2854        let insertion = frontiers.row_for_collection(self.collection_id);
2855        let updates = vec![(retraction, Diff::MINUS_ONE), (insertion, Diff::ONE)];
2856        self.send(IntrospectionType::Frontiers, updates);
2857    }
2858
2859    fn update_refresh_introspection(&mut self, write_frontier: &Antichain<Timestamp>) {
2860        let Some(refresh) = &mut self.refresh else {
2861            return;
2862        };
2863
2864        let retraction = refresh.row_for_collection(self.collection_id);
2865        refresh.frontier_update(write_frontier);
2866        let insertion = refresh.row_for_collection(self.collection_id);
2867
2868        if retraction == insertion {
2869            return; // no change
2870        }
2871
2872        let updates = vec![(retraction, Diff::MINUS_ONE), (insertion, Diff::ONE)];
2873        self.send(IntrospectionType::ComputeMaterializedViewRefreshes, updates);
2874    }
2875
2876    fn send(&self, introspection_type: IntrospectionType, updates: Vec<(Row, Diff)>) {
2877        // Failure to send means the `ComputeController` has been dropped and doesn't care about
2878        // introspection updates anymore.
2879        let _ = self.introspection_tx.send((introspection_type, updates));
2880    }
2881}
2882
2883impl Drop for CollectionIntrospection {
2884    fn drop(&mut self) {
2885        // Retract collection frontiers.
2886        if let Some(frontiers) = &self.frontiers {
2887            let row = frontiers.row_for_collection(self.collection_id);
2888            let updates = vec![(row, Diff::MINUS_ONE)];
2889            self.send(IntrospectionType::Frontiers, updates);
2890        }
2891
2892        // Retract MV refresh state.
2893        if let Some(refresh) = &self.refresh {
2894            let retraction = refresh.row_for_collection(self.collection_id);
2895            let updates = vec![(retraction, Diff::MINUS_ONE)];
2896            self.send(IntrospectionType::ComputeMaterializedViewRefreshes, updates);
2897        }
2898
2899        // Retract collection dependencies.
2900        if !self.dependency_ids.is_empty() {
2901            let updates = self.dependency_rows(Diff::MINUS_ONE);
2902            self.send(IntrospectionType::ComputeDependencies, updates);
2903        }
2904    }
2905}
2906
2907#[derive(Debug)]
2908struct FrontiersIntrospectionState {
2909    read_frontier: Antichain<Timestamp>,
2910    write_frontier: Antichain<Timestamp>,
2911}
2912
2913impl FrontiersIntrospectionState {
2914    fn new(as_of: Antichain<Timestamp>) -> Self {
2915        Self {
2916            read_frontier: as_of.clone(),
2917            write_frontier: as_of,
2918        }
2919    }
2920
2921    /// Return a `Row` reflecting the current collection frontiers.
2922    fn row_for_collection(&self, collection_id: GlobalId) -> Row {
2923        let read_frontier = self
2924            .read_frontier
2925            .as_option()
2926            .map_or(Datum::Null, |ts| ts.clone().into());
2927        let write_frontier = self
2928            .write_frontier
2929            .as_option()
2930            .map_or(Datum::Null, |ts| ts.clone().into());
2931        Row::pack_slice(&[
2932            Datum::String(&collection_id.to_string()),
2933            read_frontier,
2934            write_frontier,
2935        ])
2936    }
2937
2938    /// Update the introspection state with the given new frontiers.
2939    fn update(
2940        &mut self,
2941        read_frontier: &Antichain<Timestamp>,
2942        write_frontier: &Antichain<Timestamp>,
2943    ) {
2944        if read_frontier != &self.read_frontier {
2945            self.read_frontier.clone_from(read_frontier);
2946        }
2947        if write_frontier != &self.write_frontier {
2948            self.write_frontier.clone_from(write_frontier);
2949        }
2950    }
2951}
2952
2953/// Information needed to compute introspection updates for a REFRESH materialized view when the
2954/// write frontier advances.
2955#[derive(Debug)]
2956struct RefreshIntrospectionState {
2957    // Immutable properties of the MV
2958    refresh_schedule: RefreshSchedule,
2959    initial_as_of: Antichain<Timestamp>,
2960    // Refresh state
2961    next_refresh: Datum<'static>,           // Null or an MzTimestamp
2962    last_completed_refresh: Datum<'static>, // Null or an MzTimestamp
2963}
2964
2965impl RefreshIntrospectionState {
2966    /// Return a `Row` reflecting the current refresh introspection state.
2967    fn row_for_collection(&self, collection_id: GlobalId) -> Row {
2968        Row::pack_slice(&[
2969            Datum::String(&collection_id.to_string()),
2970            self.last_completed_refresh,
2971            self.next_refresh,
2972        ])
2973    }
2974}
2975
2976impl RefreshIntrospectionState {
2977    /// Construct a new [`RefreshIntrospectionState`], and apply an initial `frontier_update()` at
2978    /// the `upper`.
2979    fn new(
2980        refresh_schedule: RefreshSchedule,
2981        initial_as_of: Antichain<Timestamp>,
2982        upper: &Antichain<Timestamp>,
2983    ) -> Self {
2984        let mut self_ = Self {
2985            refresh_schedule: refresh_schedule.clone(),
2986            initial_as_of: initial_as_of.clone(),
2987            next_refresh: Datum::Null,
2988            last_completed_refresh: Datum::Null,
2989        };
2990        self_.frontier_update(upper);
2991        self_
2992    }
2993
2994    /// Should be called whenever the write frontier of the collection advances. It updates the
2995    /// state that should be recorded in introspection relations, but doesn't send the updates yet.
2996    fn frontier_update(&mut self, write_frontier: &Antichain<Timestamp>) {
2997        if write_frontier.is_empty() {
2998            self.last_completed_refresh =
2999                if let Some(last_refresh) = self.refresh_schedule.last_refresh() {
3000                    last_refresh.into()
3001                } else {
3002                    // If there is no last refresh, then we have a `REFRESH EVERY`, in which case
3003                    // the saturating roundup puts a refresh at the maximum possible timestamp.
3004                    Timestamp::MAX.into()
3005                };
3006            self.next_refresh = Datum::Null;
3007        } else {
3008            if PartialOrder::less_equal(write_frontier, &self.initial_as_of) {
3009                // We are before the first refresh.
3010                self.last_completed_refresh = Datum::Null;
3011                let initial_as_of = self.initial_as_of.as_option().expect(
3012                    "initial_as_of can't be [], because then there would be no refreshes at all",
3013                );
3014                let first_refresh = self
3015                    .refresh_schedule
3016                    .round_up_timestamp(*initial_as_of)
3017                    .expect("sequencing makes sure that REFRESH MVs always have a first refresh");
3018                soft_assert_or_log!(
3019                    first_refresh == *initial_as_of,
3020                    "initial_as_of should be set to the first refresh"
3021                );
3022                self.next_refresh = first_refresh.into();
3023            } else {
3024                // The first refresh has already happened.
3025                let write_frontier = write_frontier.as_option().expect("checked above");
3026                self.last_completed_refresh = self
3027                    .refresh_schedule
3028                    .round_down_timestamp_m1(*write_frontier)
3029                    .map_or_else(
3030                        || {
3031                            soft_panic_or_log!(
3032                                "rounding down should have returned the first refresh or later"
3033                            );
3034                            Datum::Null
3035                        },
3036                        |last_completed_refresh| last_completed_refresh.into(),
3037                    );
3038                self.next_refresh = write_frontier.clone().into();
3039            }
3040        }
3041    }
3042}
3043
3044/// A note of an outstanding peek response.
3045#[derive(Debug)]
3046struct PendingPeek {
3047    /// For replica-targeted peeks, this specifies the replica whose response we should pass on.
3048    ///
3049    /// If this value is `None`, we pass on the first response.
3050    target_replica: Option<ReplicaId>,
3051    /// The OpenTelemetry context for this peek.
3052    otel_ctx: OpenTelemetryContext,
3053    /// The time at which the peek was requested.
3054    ///
3055    /// Used to track peek durations.
3056    requested_at: Instant,
3057    /// The read hold installed to serve this peek.
3058    read_hold: ReadHold,
3059    /// The channel to send peek results.
3060    peek_response_tx: oneshot::Sender<PeekResponse>,
3061    /// An optional limit of the peek's result size.
3062    limit: Option<usize>,
3063    /// The offset into the peek's result.
3064    offset: usize,
3065}
3066
3067#[derive(Debug, Clone)]
3068struct ActiveSubscribe {
3069    /// Current upper frontier of this subscribe.
3070    frontier: Antichain<Timestamp>,
3071}
3072
3073impl Default for ActiveSubscribe {
3074    fn default() -> Self {
3075        Self {
3076            frontier: Antichain::from_elem(Timestamp::MIN),
3077        }
3078    }
3079}
3080
3081/// State maintained about individual replicas.
3082#[derive(Debug)]
3083struct ReplicaState {
3084    /// The ID of the replica.
3085    id: ReplicaId,
3086    /// Client for the running replica task.
3087    client: ReplicaClient,
3088    /// The replica configuration.
3089    config: ReplicaConfig,
3090    /// Replica metrics.
3091    metrics: ReplicaMetrics,
3092    /// A channel through which introspection updates are delivered.
3093    introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
3094    /// Per-replica collection state.
3095    collections: BTreeMap<GlobalId, ReplicaCollectionState>,
3096    /// The epoch of the replica.
3097    epoch: u64,
3098}
3099
3100impl ReplicaState {
3101    fn new(
3102        id: ReplicaId,
3103        client: ReplicaClient,
3104        config: ReplicaConfig,
3105        metrics: ReplicaMetrics,
3106        introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
3107        epoch: u64,
3108    ) -> Self {
3109        Self {
3110            id,
3111            client,
3112            config,
3113            metrics,
3114            introspection_tx,
3115            epoch,
3116            collections: Default::default(),
3117        }
3118    }
3119
3120    /// Add a collection to the replica state.
3121    ///
3122    /// # Panics
3123    ///
3124    /// Panics if a collection with the same ID exists already.
3125    fn add_collection(
3126        &mut self,
3127        id: GlobalId,
3128        as_of: Antichain<Timestamp>,
3129        input_read_holds: Vec<ReadHold>,
3130    ) {
3131        let metrics = self.metrics.for_collection(id);
3132        let introspection = ReplicaCollectionIntrospection::new(
3133            self.id,
3134            id,
3135            self.introspection_tx.clone(),
3136            as_of.clone(),
3137        );
3138        let mut state =
3139            ReplicaCollectionState::new(metrics, as_of, introspection, input_read_holds);
3140
3141        // In an effort to keep the produced wallclock lag introspection data small and
3142        // predictable, we disable wallclock lag tracking for transient collections, i.e. slow-path
3143        // select indexes and subscribes.
3144        if id.is_transient() {
3145            state.wallclock_lag_max = None;
3146        }
3147
3148        if let Some(previous) = self.collections.insert(id, state) {
3149            panic!("attempt to add a collection with existing ID {id} (previous={previous:?}");
3150        }
3151    }
3152
3153    /// Remove state for a collection.
3154    fn remove_collection(&mut self, id: GlobalId) -> Option<ReplicaCollectionState> {
3155        self.collections.remove(&id)
3156    }
3157
3158    /// Returns whether all replica frontiers of the given collection are empty.
3159    fn collection_frontiers_empty(&self, id: GlobalId) -> bool {
3160        self.collections.get(&id).map_or(true, |c| {
3161            c.write_frontier.is_empty()
3162                && c.input_frontier.is_empty()
3163                && c.output_frontier.is_empty()
3164        })
3165    }
3166
3167    /// Returns the state of the [`ReplicaState`] formatted as JSON.
3168    ///
3169    /// The returned value is not guaranteed to be stable and may change at any point in time.
3170    #[mz_ore::instrument(level = "debug")]
3171    pub fn dump(&self) -> Result<serde_json::Value, anyhow::Error> {
3172        // Note: We purposefully use the `Debug` formatting for the value of all fields in the
3173        // returned object as a tradeoff between usability and stability. `serde_json` will fail
3174        // to serialize an object if the keys aren't strings, so `Debug` formatting the values
3175        // prevents a future unrelated change from silently breaking this method.
3176
3177        // Destructure `self` here so we don't forget to consider dumping newly added fields.
3178        let Self {
3179            id,
3180            client: _,
3181            config: _,
3182            metrics: _,
3183            introspection_tx: _,
3184            epoch,
3185            collections,
3186        } = self;
3187
3188        let collections: BTreeMap<_, _> = collections
3189            .iter()
3190            .map(|(id, collection)| (id.to_string(), format!("{collection:?}")))
3191            .collect();
3192
3193        Ok(serde_json::json!({
3194            "id": id.to_string(),
3195            "collections": collections,
3196            "epoch": epoch,
3197        }))
3198    }
3199}
3200
3201#[derive(Debug)]
3202struct ReplicaCollectionState {
3203    /// The replica write frontier of this collection.
3204    ///
3205    /// See [`FrontiersResponse::write_frontier`].
3206    write_frontier: Antichain<Timestamp>,
3207    /// The replica input frontier of this collection.
3208    ///
3209    /// See [`FrontiersResponse::input_frontier`].
3210    input_frontier: Antichain<Timestamp>,
3211    /// The replica output frontier of this collection.
3212    ///
3213    /// See [`FrontiersResponse::output_frontier`].
3214    output_frontier: Antichain<Timestamp>,
3215
3216    /// Metrics tracked for this collection.
3217    ///
3218    /// If this is `None`, no metrics are collected.
3219    metrics: Option<ReplicaCollectionMetrics>,
3220    /// As-of frontier with which this collection was installed on the replica.
3221    as_of: Antichain<Timestamp>,
3222    /// Tracks introspection state for this collection.
3223    introspection: ReplicaCollectionIntrospection,
3224    /// Read holds on storage inputs to this collection.
3225    ///
3226    /// These read holds are kept to ensure that the replica is able to read from storage inputs at
3227    /// all times it hasn't read yet. We only need to install read holds for storage inputs since
3228    /// compaction of compute inputs is implicitly held back by Timely/DD.
3229    input_read_holds: Vec<ReadHold>,
3230
3231    /// Maximum frontier wallclock lag since the last `WallclockLagHistory` introspection update.
3232    ///
3233    /// If this is `None`, wallclock lag is not tracked for this collection.
3234    wallclock_lag_max: Option<WallclockLag>,
3235}
3236
3237impl ReplicaCollectionState {
3238    fn new(
3239        metrics: Option<ReplicaCollectionMetrics>,
3240        as_of: Antichain<Timestamp>,
3241        introspection: ReplicaCollectionIntrospection,
3242        input_read_holds: Vec<ReadHold>,
3243    ) -> Self {
3244        Self {
3245            write_frontier: as_of.clone(),
3246            input_frontier: as_of.clone(),
3247            output_frontier: as_of.clone(),
3248            metrics,
3249            as_of,
3250            introspection,
3251            input_read_holds,
3252            wallclock_lag_max: Some(WallclockLag::MIN),
3253        }
3254    }
3255
3256    /// Returns whether this collection is hydrated.
3257    fn hydrated(&self) -> bool {
3258        // If the observed frontier is greater than the collection's as-of, the collection has
3259        // produced some output and is therefore hydrated.
3260        //
3261        // We need to consider the edge case where the as-of is the empty frontier. Such an as-of
3262        // is not useful for indexes, because they wouldn't be readable. For write-only
3263        // collections, an empty as-of means that the collection has been fully written and no new
3264        // dataflow needs to be created for it. Consequently, no hydration will happen either.
3265        //
3266        // Based on this, we could respond in two ways:
3267        //  * `false`, as in "the dataflow was never created"
3268        //  * `true`, as in "the dataflow completed immediately"
3269        //
3270        // Since hydration is often used as a measure of dataflow progress and we don't want to
3271        // give the impression that certain dataflows are somehow stuck when they are not, we go
3272        // with the second interpretation here.
3273        self.as_of.is_empty() || PartialOrder::less_than(&self.as_of, &self.output_frontier)
3274    }
3275
3276    /// Updates the replica write frontier of this collection.
3277    fn update_write_frontier(&mut self, new_frontier: Antichain<Timestamp>) {
3278        if PartialOrder::less_than(&new_frontier, &self.write_frontier) {
3279            soft_panic_or_log!(
3280                "replica collection write frontier regression (old={:?}, new={new_frontier:?})",
3281                self.write_frontier,
3282            );
3283            return;
3284        } else if new_frontier == self.write_frontier {
3285            return;
3286        }
3287
3288        self.write_frontier = new_frontier;
3289    }
3290
3291    /// Updates the replica input frontier of this collection.
3292    fn update_input_frontier(&mut self, new_frontier: Antichain<Timestamp>) {
3293        if PartialOrder::less_than(&new_frontier, &self.input_frontier) {
3294            soft_panic_or_log!(
3295                "replica collection input frontier regression (old={:?}, new={new_frontier:?})",
3296                self.input_frontier,
3297            );
3298            return;
3299        } else if new_frontier == self.input_frontier {
3300            return;
3301        }
3302
3303        self.input_frontier = new_frontier;
3304
3305        // Relax our read holds on collection inputs.
3306        for read_hold in &mut self.input_read_holds {
3307            let result = read_hold.try_downgrade(self.input_frontier.clone());
3308            soft_assert_or_log!(
3309                result.is_ok(),
3310                "read hold downgrade failed (read_hold={read_hold:?}, new_since={:?})",
3311                self.input_frontier,
3312            );
3313        }
3314    }
3315
3316    /// Updates the replica output frontier of this collection.
3317    fn update_output_frontier(&mut self, new_frontier: Antichain<Timestamp>) {
3318        if PartialOrder::less_than(&new_frontier, &self.output_frontier) {
3319            soft_panic_or_log!(
3320                "replica collection output frontier regression (old={:?}, new={new_frontier:?})",
3321                self.output_frontier,
3322            );
3323            return;
3324        } else if new_frontier == self.output_frontier {
3325            return;
3326        }
3327
3328        self.output_frontier = new_frontier;
3329    }
3330}
3331
3332/// Maintains the introspection state for a given replica and collection, and ensures that reported
3333/// introspection data is retracted when the collection is dropped.
3334#[derive(Debug)]
3335struct ReplicaCollectionIntrospection {
3336    /// The ID of the replica.
3337    replica_id: ReplicaId,
3338    /// The ID of the compute collection.
3339    collection_id: GlobalId,
3340    /// The collection's reported replica write frontier.
3341    write_frontier: Antichain<Timestamp>,
3342    /// A channel through which introspection updates are delivered.
3343    introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
3344}
3345
3346impl ReplicaCollectionIntrospection {
3347    /// Create a new `HydrationState` and initialize introspection.
3348    fn new(
3349        replica_id: ReplicaId,
3350        collection_id: GlobalId,
3351        introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
3352        as_of: Antichain<Timestamp>,
3353    ) -> Self {
3354        let self_ = Self {
3355            replica_id,
3356            collection_id,
3357            write_frontier: as_of,
3358            introspection_tx,
3359        };
3360
3361        self_.report_initial_state();
3362        self_
3363    }
3364
3365    /// Reports the initial introspection state.
3366    fn report_initial_state(&self) {
3367        let row = self.write_frontier_row();
3368        let updates = vec![(row, Diff::ONE)];
3369        self.send(IntrospectionType::ReplicaFrontiers, updates);
3370    }
3371
3372    /// Observe the given current write frontier and update the introspection state as necessary.
3373    fn observe_frontier(&mut self, write_frontier: &Antichain<Timestamp>) {
3374        if self.write_frontier == *write_frontier {
3375            return; // no change
3376        }
3377
3378        let retraction = self.write_frontier_row();
3379        self.write_frontier.clone_from(write_frontier);
3380        let insertion = self.write_frontier_row();
3381
3382        let updates = vec![(retraction, Diff::MINUS_ONE), (insertion, Diff::ONE)];
3383        self.send(IntrospectionType::ReplicaFrontiers, updates);
3384    }
3385
3386    /// Return a `Row` reflecting the current replica write frontier.
3387    fn write_frontier_row(&self) -> Row {
3388        let write_frontier = self
3389            .write_frontier
3390            .as_option()
3391            .map_or(Datum::Null, |ts| ts.clone().into());
3392        Row::pack_slice(&[
3393            Datum::String(&self.collection_id.to_string()),
3394            Datum::String(&self.replica_id.to_string()),
3395            write_frontier,
3396        ])
3397    }
3398
3399    fn send(&self, introspection_type: IntrospectionType, updates: Vec<(Row, Diff)>) {
3400        // Failure to send means the `ComputeController` has been dropped and doesn't care about
3401        // introspection updates anymore.
3402        let _ = self.introspection_tx.send((introspection_type, updates));
3403    }
3404}
3405
3406impl Drop for ReplicaCollectionIntrospection {
3407    fn drop(&mut self) {
3408        // Retract the write frontier.
3409        let row = self.write_frontier_row();
3410        let updates = vec![(row, Diff::MINUS_ONE)];
3411        self.send(IntrospectionType::ReplicaFrontiers, updates);
3412    }
3413}
3414
3415#[cfg(test)]
3416mod tests {
3417    use std::collections::BTreeMap;
3418
3419    use mz_compute_types::dyncfgs::{ENABLE_COLUMN_PAGED_BATCHER, ENABLE_MZ_JOIN_CORE};
3420    use mz_dyncfg::{ConfigSet, ConfigUpdates, ConfigVal};
3421    use mz_persist_types::PersistLocation;
3422
3423    use crate::protocol::command::{ComputeCommand, InstanceConfig};
3424
3425    use super::{Instance, ReplicaId};
3426
3427    fn create_instance_command() -> ComputeCommand {
3428        ComputeCommand::CreateInstance(Box::new(InstanceConfig {
3429            logging: Default::default(),
3430            expiration_offset: None,
3431            peek_stash_persist_location: PersistLocation::new_in_mem(),
3432            arrangement_dictionary_compression: false,
3433            initial_config: Default::default(),
3434        }))
3435    }
3436
3437    fn initial_config(cmd: &ComputeCommand) -> &ConfigUpdates {
3438        match cmd {
3439            ComputeCommand::CreateInstance(config) => &config.initial_config,
3440            other => panic!("expected CreateInstance, got {other:?}"),
3441        }
3442    }
3443
3444    /// `CreateInstance` is specialized with a full snapshot of the instance-wide dyncfg, so the
3445    /// replica seeds its worker config at create time rather than waiting for the first
3446    /// `UpdateConfiguration`. This is the regression guard for create-time setup observing dyncfg
3447    /// defaults.
3448    #[mz_ore::test]
3449    fn create_instance_snapshots_instance_wide_dyncfg() {
3450        let dyncfg = ConfigSet::default()
3451            .add(&ENABLE_COLUMN_PAGED_BATCHER)
3452            .add(&ENABLE_MZ_JOIN_CORE);
3453        let mut updates = ConfigUpdates::default();
3454        updates.add(&ENABLE_COLUMN_PAGED_BATCHER, true);
3455        updates.add(&ENABLE_MZ_JOIN_CORE, false);
3456        updates.apply(&dyncfg);
3457
3458        // A replica without an override sees exactly the instance-wide values.
3459        let overrides = BTreeMap::new();
3460        let cmd = Instance::specialize_command_for_replica(
3461            create_instance_command(),
3462            ReplicaId::User(1),
3463            &overrides,
3464            &dyncfg,
3465        );
3466        let snapshot = initial_config(&cmd);
3467        assert_eq!(
3468            snapshot.updates.get(ENABLE_COLUMN_PAGED_BATCHER.name()),
3469            Some(&ConfigVal::Bool(true)),
3470        );
3471        assert_eq!(
3472            snapshot.updates.get(ENABLE_MZ_JOIN_CORE.name()),
3473            Some(&ConfigVal::Bool(false)),
3474        );
3475    }
3476
3477    /// A replica-scoped override beats the instance-wide value in the create-time snapshot, so a
3478    /// create-time-frozen scoped flag reaches the replica with its override applied.
3479    #[mz_ore::test]
3480    fn create_instance_snapshot_applies_replica_override() {
3481        let dyncfg = ConfigSet::default().add(&ENABLE_COLUMN_PAGED_BATCHER);
3482        let mut updates = ConfigUpdates::default();
3483        updates.add(&ENABLE_COLUMN_PAGED_BATCHER, true);
3484        updates.apply(&dyncfg);
3485
3486        let replica = ReplicaId::User(1);
3487        let mut override_updates = ConfigUpdates::default();
3488        override_updates.add(&ENABLE_COLUMN_PAGED_BATCHER, false);
3489        let overrides = BTreeMap::from([(replica, override_updates)]);
3490
3491        let cmd = Instance::specialize_command_for_replica(
3492            create_instance_command(),
3493            replica,
3494            &overrides,
3495            &dyncfg,
3496        );
3497        assert_eq!(
3498            initial_config(&cmd)
3499                .updates
3500                .get(ENABLE_COLUMN_PAGED_BATCHER.name()),
3501            Some(&ConfigVal::Bool(false)),
3502            "replica override should win over the instance-wide value",
3503        );
3504    }
3505
3506    /// `UpdateConfiguration` continues to merge the replica's override into the update.
3507    #[mz_ore::test]
3508    fn update_configuration_merges_replica_override() {
3509        let dyncfg = ConfigSet::default().add(&ENABLE_COLUMN_PAGED_BATCHER);
3510
3511        let replica = ReplicaId::User(1);
3512        let mut override_updates = ConfigUpdates::default();
3513        override_updates.add(&ENABLE_COLUMN_PAGED_BATCHER, true);
3514        let overrides = BTreeMap::from([(replica, override_updates)]);
3515
3516        let cmd = Instance::specialize_command_for_replica(
3517            ComputeCommand::UpdateConfiguration(Box::new(Default::default())),
3518            replica,
3519            &overrides,
3520            &dyncfg,
3521        );
3522        match cmd {
3523            ComputeCommand::UpdateConfiguration(params) => assert_eq!(
3524                params
3525                    .dyncfg_updates
3526                    .updates
3527                    .get(ENABLE_COLUMN_PAGED_BATCHER.name()),
3528                Some(&ConfigVal::Bool(true)),
3529            ),
3530            other => panic!("expected UpdateConfiguration, got {other:?}"),
3531        }
3532    }
3533}