1use 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
117pub(super) type Command = Box<dyn FnOnce(&mut Instance) + Send>;
119
120pub(super) type ReplicaResponse = (ReplicaId, u64, ComputeResponse);
123
124pub(super) struct Instance {
126 build_info: &'static BuildInfo,
128 storage_collections: StorageCollections,
130 initialized: bool,
132 read_only: bool,
137 workload_class: Option<String>,
141 replicas: BTreeMap<ReplicaId, ReplicaState>,
143 replica_dyncfg_overrides: BTreeMap<ReplicaId, ConfigUpdates>,
150 collections: BTreeMap<GlobalId, CollectionState>,
158 log_sources: BTreeMap<LogVariant, GlobalId>,
160 peeks: BTreeMap<Uuid, PendingPeek>,
169 subscribes: BTreeMap<GlobalId, ActiveSubscribe>,
183 copy_tos: BTreeSet<GlobalId>,
191 history: ComputeCommandHistory<UIntGauge>,
193 command_rx: mpsc::UnboundedReceiver<Command>,
195 response_tx: mpsc::UnboundedSender<ComputeControllerResponse>,
197 introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
199 metrics: InstanceMetrics,
201 dyncfg: Arc<ConfigSet>,
203
204 peek_stash_persist_location: PersistLocation,
206
207 now: NowFn,
209 wallclock_lag: WallclockLagFn<Timestamp>,
211 wallclock_lag_last_recorded: DateTime<Utc>,
213
214 read_hold_tx: read_holds::ChangeTx,
219 replica_tx: mz_ore::channel::InstrumentedUnboundedSender<ReplicaResponse, IntCounter>,
221 replica_rx: mz_ore::channel::InstrumentedUnboundedReceiver<ReplicaResponse, IntCounter>,
223}
224
225impl Instance {
226 fn collection(&self, id: GlobalId) -> Result<&CollectionState, CollectionMissing> {
228 self.collections.get(&id).ok_or(CollectionMissing(id))
229 }
230
231 fn collection_mut(&mut self, id: GlobalId) -> Result<&mut CollectionState, CollectionMissing> {
233 self.collections.get_mut(&id).ok_or(CollectionMissing(id))
234 }
235
236 fn expect_collection(&self, id: GlobalId) -> &CollectionState {
242 self.collections.get(&id).expect("collection must exist")
243 }
244
245 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 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 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 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 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 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 for replica in self.replicas.values_mut() {
342 replica.remove_collection(id);
343 }
344
345 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 let mut shutdown_input = None;
370 for (collection_id, collection) in &self.collections {
371 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 Antichain::from_elem(Timestamp::MIN)
385 } else {
386 collection.read_frontier().to_owned()
387 };
388
389 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 fn deliver_response(&self, response: ComputeControllerResponse) {
429 let _ = self.response_tx.send(response);
432 }
433
434 fn deliver_introspection_updates(&self, type_: IntrospectionType, updates: Vec<(Row, Diff)>) {
436 let _ = self.introspection_tx.send((type_, updates));
439 }
440
441 fn replica_exists(&self, id: ReplicaId) -> bool {
443 self.replicas.contains_key(&id)
444 }
445
446 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 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 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 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 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 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 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 for replica in self.replicas.values_mut() {
597 for (id, collection) in &mut replica.collections {
598 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 let secs = lag.unwrap_seconds_or(u64::MAX);
619 metrics.wallclock_lag.observe(secs);
620 };
621 }
622 }
623
624 self.maybe_record_wallclock_lag();
626 }
627
628 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 #[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 #[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 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 }
797 }
798
799 Ok(all_hydrated)
800 }
801
802 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 #[mz_ore::instrument(level = "debug")]
840 pub fn dump(&self) -> Result<serde_json::Value, anyhow::Error> {
841 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 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 nonce: Uuid::default(),
987 });
988
989 let instance_config = InstanceConfig {
990 peek_stash_persist_location: self.peek_stash_persist_location.clone(),
991 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 #[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 #[mz_ore::instrument(level = "debug")]
1032 pub fn initialization_complete(&mut self) {
1033 if !self.initialized {
1035 self.send(ComputeCommand::InitializationComplete);
1036 self.initialized = true;
1037 }
1038 }
1039
1040 #[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 if !collection.read_only {
1049 return Ok(());
1050 }
1051
1052 let as_of = collection.read_frontier();
1054
1055 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 #[mz_ore::instrument(level = "debug")]
1077 pub fn shutdown(&mut self) {
1078 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 fn initiate_shutdown(&mut self) {
1095 let (_tx, rx) = mpsc::unbounded_channel();
1098 self.command_rx = rx;
1099 }
1100
1101 #[mz_ore::instrument(level = "debug")]
1103 fn send(&mut self, cmd: ComputeCommand) {
1104 self.history.push(cmd.clone());
1109
1110 let target_replica = self.target_replica(&cmd);
1111
1112 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 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 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 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 ComputeCommand::Peek(_)
1203 | ComputeCommand::Hello { .. }
1204 | ComputeCommand::CreateInstance(_)
1205 | ComputeCommand::InitializationComplete
1206 | ComputeCommand::UpdateConfiguration(_)
1207 | ComputeCommand::CancelPeek { .. } => None,
1208 }
1209 }
1210
1211 #[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 self.history.reduce();
1239
1240 self.history.update_source_uppers(&self.storage_collections);
1242
1243 for command in self.history.iter() {
1245 if let Some(target_replica) = self.target_replica(command)
1247 && target_replica != id
1248 {
1249 continue;
1250 }
1251
1252 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 tracing::warn!("Replica {:?} connection terminated during hydration", id);
1264 break;
1265 }
1266 }
1267
1268 if self.add_replica_state(id, client, config, epoch).is_err() {
1270 self.initiate_shutdown();
1276 }
1277
1278 Ok(())
1279 }
1280
1281 #[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 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 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 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 self.forward_implied_capabilities();
1356
1357 Ok(())
1358 }
1359
1360 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 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 #[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 if let Some(replica_id) = target_replica {
1408 if !self.replica_exists(replica_id) {
1409 return Err(ReplicaMissing(replica_id));
1410 }
1411 }
1412
1413 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 let mut storage_dependencies = BTreeMap::new();
1424 let mut compute_dependencies = BTreeMap::new();
1425
1426 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 as_of.is_empty() {
1456 replica_input_read_holds = Default::default();
1457 }
1458
1459 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 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 for subscribe_id in dataflow.subscribe_ids() {
1490 self.subscribes
1491 .insert(subscribe_id, ActiveSubscribe::default());
1492 }
1493
1494 for copy_to_id in dataflow.copy_to_ids() {
1496 self.copy_tos.insert(copy_to_id);
1497 }
1498
1499 let mut source_imports = BTreeMap::new();
1502 for (id, import) in dataflow.source_imports {
1503 let frontiers = self
1504 .storage_collections
1505 .collection_frontiers(id)
1506 .expect("collection exists");
1507
1508 let collection_metadata = self
1509 .storage_collections
1510 .collection_metadata(id)
1511 .expect("we have a read hold on this collection");
1512
1513 let desc = SourceInstanceDesc {
1514 storage_metadata: collection_metadata.clone(),
1515 arguments: import.desc.arguments,
1516 typ: import.desc.typ.clone(),
1517 };
1518 source_imports.insert(
1519 id,
1520 mz_compute_types::dataflows::SourceImport {
1521 desc,
1522 monotonic: import.monotonic,
1523 with_snapshot: import.with_snapshot,
1524 upper: frontiers.write_frontier,
1525 },
1526 );
1527 }
1528
1529 let mut sink_exports = BTreeMap::new();
1530 for (id, se) in dataflow.sink_exports {
1531 let connection = match se.connection {
1532 ComputeSinkConnection::MaterializedView(conn) => {
1533 let metadata = self
1534 .storage_collections
1535 .collection_metadata(id)
1536 .map_err(|_| CollectionMissing(id))?
1537 .clone();
1538 let conn = MaterializedViewSinkConnection {
1539 value_desc: conn.value_desc,
1540 storage_metadata: metadata,
1541 };
1542 ComputeSinkConnection::MaterializedView(conn)
1543 }
1544 ComputeSinkConnection::Subscribe(conn) => ComputeSinkConnection::Subscribe(conn),
1545 ComputeSinkConnection::CopyToS3Oneshot(conn) => {
1546 ComputeSinkConnection::CopyToS3Oneshot(conn)
1547 }
1548 };
1549 let desc = ComputeSinkDesc {
1550 from: se.from,
1551 from_desc: se.from_desc,
1552 connection,
1553 with_snapshot: se.with_snapshot,
1554 up_to: se.up_to,
1555 non_null_assertions: se.non_null_assertions,
1556 refresh_schedule: se.refresh_schedule,
1557 };
1558 sink_exports.insert(id, desc);
1559 }
1560
1561 let objects_to_build = dataflow
1563 .objects_to_build
1564 .into_iter()
1565 .map(|object| BuildDesc {
1566 id: object.id,
1567 plan: RenderPlan::try_from(object.plan).expect("valid plan"),
1568 })
1569 .collect();
1570
1571 let augmented_dataflow = DataflowDescription {
1572 source_imports,
1573 sink_exports,
1574 objects_to_build,
1575 index_imports: dataflow.index_imports,
1577 index_exports: dataflow.index_exports,
1578 as_of: dataflow.as_of.clone(),
1579 until: dataflow.until,
1580 initial_storage_as_of: dataflow.initial_storage_as_of,
1581 refresh_schedule: dataflow.refresh_schedule,
1582 debug_name: dataflow.debug_name,
1583 time_dependence: dataflow.time_dependence,
1584 };
1585
1586 if augmented_dataflow.is_transient() {
1587 tracing::debug!(
1588 name = %augmented_dataflow.debug_name,
1589 import_ids = %augmented_dataflow.display_import_ids(),
1590 export_ids = %augmented_dataflow.display_export_ids(),
1591 as_of = ?augmented_dataflow.as_of.as_ref().unwrap().elements(),
1592 until = ?augmented_dataflow.until.elements(),
1593 "creating dataflow",
1594 );
1595 } else {
1596 tracing::info!(
1597 name = %augmented_dataflow.debug_name,
1598 import_ids = %augmented_dataflow.display_import_ids(),
1599 export_ids = %augmented_dataflow.display_export_ids(),
1600 as_of = ?augmented_dataflow.as_of.as_ref().unwrap().elements(),
1601 until = ?augmented_dataflow.until.elements(),
1602 "creating dataflow",
1603 );
1604 }
1605
1606 if as_of.is_empty() {
1609 tracing::info!(
1610 name = %augmented_dataflow.debug_name,
1611 "not sending `CreateDataflow`, because of empty `as_of`",
1612 );
1613 } else {
1614 let collections: Vec<_> = augmented_dataflow.export_ids().collect();
1615 self.send(ComputeCommand::CreateDataflow(Box::new(augmented_dataflow)));
1616
1617 for id in collections {
1618 self.maybe_schedule_collection(id);
1619 }
1620 }
1621
1622 Ok(())
1623 }
1624
1625 fn maybe_schedule_collection(&mut self, id: GlobalId) {
1631 let collection = self.expect_collection(id);
1632
1633 if collection.scheduled {
1635 return;
1636 }
1637
1638 let as_of = collection.read_frontier();
1639
1640 if as_of.is_empty() {
1643 return;
1644 }
1645
1646 let ready = if id.is_transient() {
1647 true
1653 } else {
1654 let not_self_dep = |x: &GlobalId| *x != id;
1660
1661 let mut deps_scheduled = true;
1664
1665 let compute_deps = collection.compute_dependency_ids().filter(not_self_dep);
1670 let mut compute_frontiers = Vec::new();
1671 for id in compute_deps {
1672 let dep = &self.expect_collection(id);
1673 deps_scheduled &= dep.scheduled;
1674 compute_frontiers.push(dep.write_frontier());
1675 }
1676
1677 let storage_deps = collection.storage_dependency_ids().filter(not_self_dep);
1678 let storage_frontiers = self
1679 .storage_collections
1680 .collections_frontiers(storage_deps.collect())
1681 .expect("must exist");
1682 let storage_frontiers = storage_frontiers.into_iter().map(|f| f.write_frontier);
1683
1684 let mut frontiers = compute_frontiers.into_iter().chain(storage_frontiers);
1685 let frontiers_ready =
1686 frontiers.all(|frontier| PartialOrder::less_than(&as_of, &frontier));
1687
1688 deps_scheduled && frontiers_ready
1689 };
1690
1691 if ready {
1692 self.send(ComputeCommand::Schedule(id));
1693 let collection = self.expect_collection_mut(id);
1694 collection.scheduled = true;
1695 }
1696 }
1697
1698 fn schedule_collections(&mut self) {
1700 let ids: Vec<_> = self.collections.keys().copied().collect();
1701 for id in ids {
1702 self.maybe_schedule_collection(id);
1703 }
1704 }
1705
1706 #[mz_ore::instrument(level = "debug")]
1709 pub fn drop_collections(&mut self, ids: Vec<GlobalId>) -> Result<(), CollectionMissing> {
1710 for id in &ids {
1711 let collection = self.collection_mut(*id)?;
1712
1713 collection.dropped = true;
1715
1716 collection.implied_read_hold.release();
1719 collection.warmup_read_hold.release();
1720
1721 self.subscribes.remove(id);
1724 self.copy_tos.remove(id);
1727 }
1728
1729 Ok(())
1730 }
1731
1732 #[mz_ore::instrument(level = "debug")]
1736 pub fn peek(
1737 &mut self,
1738 peek_target: PeekTarget,
1739 literal_constraints: Option<Vec<Row>>,
1740 uuid: Uuid,
1741 timestamp: Timestamp,
1742 result_desc: RelationDesc,
1743 finishing: RowSetFinishing,
1744 map_filter_project: mz_expr::SafeMfpPlan,
1745 mut read_hold: ReadHold,
1746 target_replica: Option<ReplicaId>,
1747 peek_response_tx: oneshot::Sender<PeekResponse>,
1748 ) -> Result<(), PeekError> {
1749 use PeekError::*;
1750
1751 let target_id = peek_target.id();
1752
1753 if read_hold.id() != target_id {
1755 return Err(ReadHoldIdMismatch(read_hold.id()));
1756 }
1757 read_hold
1758 .try_downgrade(Antichain::from_elem(timestamp.clone()))
1759 .map_err(|_| ReadHoldInsufficient(target_id))?;
1760
1761 if let Some(target) = target_replica {
1762 if !self.replica_exists(target) {
1763 return Err(ReplicaMissing(target));
1764 }
1765 }
1766
1767 let otel_ctx = OpenTelemetryContext::obtain();
1768
1769 self.peeks.insert(
1770 uuid,
1771 PendingPeek {
1772 target_replica,
1773 otel_ctx: otel_ctx.clone(),
1775 requested_at: Instant::now(),
1776 read_hold,
1777 peek_response_tx,
1778 limit: finishing.limit.map(usize::cast_from),
1779 offset: finishing.offset,
1780 },
1781 );
1782
1783 let peek = Peek {
1784 literal_constraints,
1785 uuid,
1786 timestamp,
1787 finishing,
1788 map_filter_project,
1789 otel_ctx,
1792 target: peek_target,
1793 result_desc,
1794 };
1795 self.send(ComputeCommand::Peek(Box::new(peek)));
1796
1797 Ok(())
1798 }
1799
1800 #[mz_ore::instrument(level = "debug")]
1802 pub fn cancel_peek(&mut self, uuid: Uuid, reason: PeekResponse) {
1803 let Some(peek) = self.peeks.get_mut(&uuid) else {
1804 tracing::warn!("did not find pending peek for {uuid}");
1805 return;
1806 };
1807
1808 let duration = peek.requested_at.elapsed();
1809 self.metrics
1810 .observe_peek_response(&PeekResponse::Canceled, duration);
1811
1812 let otel_ctx = peek.otel_ctx.clone();
1814 otel_ctx.attach_as_parent();
1815
1816 self.deliver_response(ComputeControllerResponse::PeekNotification(
1817 uuid,
1818 PeekNotification::Canceled,
1819 otel_ctx,
1820 ));
1821
1822 self.finish_peek(uuid, reason);
1825 }
1826
1827 #[mz_ore::instrument(level = "debug")]
1839 pub fn set_read_policy(
1840 &mut self,
1841 policies: Vec<(GlobalId, ReadPolicy)>,
1842 ) -> Result<(), ReadPolicyError> {
1843 for (id, _policy) in &policies {
1846 let collection = self.collection(*id)?;
1847 if collection.read_policy.is_none() {
1848 return Err(ReadPolicyError::WriteOnlyCollection(*id));
1849 }
1850 }
1851
1852 for (id, new_policy) in policies {
1853 let collection = self.expect_collection_mut(id);
1854 let new_since = new_policy.frontier(collection.write_frontier().borrow());
1855 let _ = collection.implied_read_hold.try_downgrade(new_since);
1856 collection.read_policy = Some(new_policy);
1857 }
1858
1859 Ok(())
1860 }
1861
1862 #[mz_ore::instrument(level = "debug")]
1870 fn maybe_update_global_write_frontier(
1871 &mut self,
1872 id: GlobalId,
1873 new_frontier: Antichain<Timestamp>,
1874 ) {
1875 let collection = self.expect_collection_mut(id);
1876
1877 let advanced = collection.shared.lock_write_frontier(|f| {
1878 let advanced = PartialOrder::less_than(f, &new_frontier);
1879 if advanced {
1880 f.clone_from(&new_frontier);
1881 }
1882 advanced
1883 });
1884
1885 if !advanced {
1886 return;
1887 }
1888
1889 let new_since = match &collection.read_policy {
1891 Some(read_policy) => {
1892 read_policy.frontier(new_frontier.borrow())
1895 }
1896 None => {
1897 Antichain::from_iter(
1906 new_frontier
1907 .iter()
1908 .map(|t| t.step_back().unwrap_or(Timestamp::MIN)),
1909 )
1910 }
1911 };
1912 let _ = collection.implied_read_hold.try_downgrade(new_since);
1913
1914 self.deliver_response(ComputeControllerResponse::FrontierUpper {
1916 id,
1917 upper: new_frontier,
1918 });
1919 }
1920
1921 pub(super) fn apply_read_hold_change(
1923 &mut self,
1924 id: GlobalId,
1925 mut update: ChangeBatch<Timestamp>,
1926 ) {
1927 let Some(collection) = self.collections.get_mut(&id) else {
1928 soft_panic_or_log!(
1929 "read hold change for absent collection (id={id}, changes={update:?})"
1930 );
1931 return;
1932 };
1933
1934 let new_since = collection.shared.lock_read_capabilities(|caps| {
1935 let read_frontier = caps.frontier();
1938 for (time, diff) in update.iter() {
1939 let count = caps.count_for(time) + diff;
1940 assert!(
1941 count >= 0,
1942 "invalid read capabilities update: negative capability \
1943 (id={id:?}, read_capabilities={caps:?}, update={update:?})",
1944 );
1945 assert!(
1946 count == 0 || read_frontier.less_equal(time),
1947 "invalid read capabilities update: frontier regression \
1948 (id={id:?}, read_capabilities={caps:?}, update={update:?})",
1949 );
1950 }
1951
1952 let changes = caps.update_iter(update.drain());
1955
1956 let changed = changes.count() > 0;
1957 changed.then(|| caps.frontier().to_owned())
1958 });
1959
1960 let Some(new_since) = new_since else {
1961 return; };
1963
1964 for read_hold in collection.compute_dependencies.values_mut() {
1966 read_hold
1967 .try_downgrade(new_since.clone())
1968 .expect("frontiers don't regress");
1969 }
1970 for read_hold in collection.storage_dependencies.values_mut() {
1971 read_hold
1972 .try_downgrade(new_since.clone())
1973 .expect("frontiers don't regress");
1974 }
1975
1976 self.send(ComputeCommand::AllowCompaction {
1978 id,
1979 frontier: new_since,
1980 });
1981 }
1982
1983 fn finish_peek(&mut self, uuid: Uuid, response: PeekResponse) {
1992 let Some(peek) = self.peeks.remove(&uuid) else {
1993 return;
1994 };
1995
1996 let _ = peek.peek_response_tx.send(response);
1998
1999 self.send(ComputeCommand::CancelPeek { uuid });
2002
2003 drop(peek.read_hold);
2004 }
2005
2006 fn handle_response(&mut self, (replica_id, epoch, response): ReplicaResponse) {
2009 if self
2011 .replicas
2012 .get(&replica_id)
2013 .filter(|replica| replica.epoch == epoch)
2014 .is_none()
2015 {
2016 return;
2017 }
2018
2019 match response {
2022 ComputeResponse::Frontiers(id, frontiers) => {
2023 self.handle_frontiers_response(id, frontiers, replica_id);
2024 }
2025 ComputeResponse::PeekResponse(uuid, peek_response, otel_ctx) => {
2026 self.handle_peek_response(uuid, peek_response, otel_ctx, replica_id);
2027 }
2028 ComputeResponse::CopyToResponse(id, response) => {
2029 self.handle_copy_to_response(id, response, replica_id);
2030 }
2031 ComputeResponse::SubscribeResponse(id, response) => {
2032 self.handle_subscribe_response(id, response, replica_id);
2033 }
2034 ComputeResponse::Status(response) => {
2035 self.handle_status_response(response, replica_id);
2036 }
2037 }
2038 }
2039
2040 fn handle_frontiers_response(
2043 &mut self,
2044 id: GlobalId,
2045 frontiers: FrontiersResponse,
2046 replica_id: ReplicaId,
2047 ) {
2048 if !self.collections.contains_key(&id) {
2049 soft_panic_or_log!(
2050 "frontiers update for an unknown collection \
2051 (id={id}, replica_id={replica_id}, frontiers={frontiers:?})"
2052 );
2053 return;
2054 }
2055 let Some(replica) = self.replicas.get_mut(&replica_id) else {
2056 soft_panic_or_log!(
2057 "frontiers update for an unknown replica \
2058 (replica_id={replica_id}, frontiers={frontiers:?})"
2059 );
2060 return;
2061 };
2062 let Some(replica_collection) = replica.collections.get_mut(&id) else {
2063 soft_panic_or_log!(
2064 "frontiers update for an unknown replica collection \
2065 (id={id}, replica_id={replica_id}, frontiers={frontiers:?})"
2066 );
2067 return;
2068 };
2069
2070 if let Some(new_frontier) = frontiers.input_frontier {
2071 replica_collection.update_input_frontier(new_frontier.clone());
2072 }
2073 if let Some(new_frontier) = frontiers.output_frontier {
2074 replica_collection.update_output_frontier(new_frontier.clone());
2075 }
2076 if let Some(new_frontier) = frontiers.write_frontier {
2077 replica_collection.update_write_frontier(new_frontier.clone());
2078 self.maybe_update_global_write_frontier(id, new_frontier);
2079 }
2080 }
2081
2082 #[mz_ore::instrument(level = "debug")]
2083 fn handle_peek_response(
2084 &mut self,
2085 uuid: Uuid,
2086 response: PeekResponse,
2087 otel_ctx: OpenTelemetryContext,
2088 replica_id: ReplicaId,
2089 ) {
2090 otel_ctx.attach_as_parent();
2091
2092 let Some(peek) = self.peeks.get(&uuid) else {
2095 return;
2096 };
2097
2098 let target_replica = peek.target_replica.unwrap_or(replica_id);
2100 if target_replica != replica_id {
2101 return;
2102 }
2103
2104 let duration = peek.requested_at.elapsed();
2105 self.metrics.observe_peek_response(&response, duration);
2106
2107 let notification = PeekNotification::new(&response, peek.offset, peek.limit);
2108 self.deliver_response(ComputeControllerResponse::PeekNotification(
2111 uuid,
2112 notification,
2113 otel_ctx,
2114 ));
2115
2116 self.finish_peek(uuid, response)
2117 }
2118
2119 fn handle_copy_to_response(
2120 &mut self,
2121 sink_id: GlobalId,
2122 response: CopyToResponse,
2123 replica_id: ReplicaId,
2124 ) {
2125 if !self.collections.contains_key(&sink_id) {
2126 soft_panic_or_log!(
2127 "received response for an unknown copy-to \
2128 (sink_id={sink_id}, replica_id={replica_id})",
2129 );
2130 return;
2131 }
2132 let Some(replica) = self.replicas.get_mut(&replica_id) else {
2133 soft_panic_or_log!("copy-to response for an unknown replica (replica_id={replica_id})");
2134 return;
2135 };
2136 let Some(replica_collection) = replica.collections.get_mut(&sink_id) else {
2137 soft_panic_or_log!(
2138 "copy-to response for an unknown replica collection \
2139 (sink_id={sink_id}, replica_id={replica_id})"
2140 );
2141 return;
2142 };
2143
2144 replica_collection.update_write_frontier(Antichain::new());
2148 replica_collection.update_input_frontier(Antichain::new());
2149 replica_collection.update_output_frontier(Antichain::new());
2150
2151 if !self.copy_tos.remove(&sink_id) {
2154 return;
2155 }
2156
2157 let result = match response {
2158 CopyToResponse::RowCount(count) => Ok(count),
2159 CopyToResponse::Error(error) => Err(anyhow::anyhow!(error)),
2160 CopyToResponse::Dropped => {
2165 tracing::error!(
2166 %sink_id, %replica_id,
2167 "received `Dropped` response for a tracked copy to",
2168 );
2169 return;
2170 }
2171 };
2172
2173 self.deliver_response(ComputeControllerResponse::CopyToResponse(sink_id, result));
2174 }
2175
2176 fn handle_subscribe_response(
2177 &mut self,
2178 subscribe_id: GlobalId,
2179 response: SubscribeResponse,
2180 replica_id: ReplicaId,
2181 ) {
2182 if !self.collections.contains_key(&subscribe_id) {
2183 soft_panic_or_log!(
2184 "received response for an unknown subscribe \
2185 (subscribe_id={subscribe_id}, replica_id={replica_id})",
2186 );
2187 return;
2188 }
2189 let Some(replica) = self.replicas.get_mut(&replica_id) else {
2190 soft_panic_or_log!(
2191 "subscribe response for an unknown replica (replica_id={replica_id})"
2192 );
2193 return;
2194 };
2195 let Some(replica_collection) = replica.collections.get_mut(&subscribe_id) else {
2196 soft_panic_or_log!(
2197 "subscribe response for an unknown replica collection \
2198 (subscribe_id={subscribe_id}, replica_id={replica_id})"
2199 );
2200 return;
2201 };
2202
2203 let write_frontier = match &response {
2207 SubscribeResponse::Batch(batch) => batch.upper.clone(),
2208 SubscribeResponse::DroppedAt(_) => Antichain::new(),
2209 };
2210
2211 replica_collection.update_write_frontier(write_frontier.clone());
2215 replica_collection.update_input_frontier(write_frontier.clone());
2216 replica_collection.update_output_frontier(write_frontier.clone());
2217
2218 let Some(mut subscribe) = self.subscribes.get(&subscribe_id).cloned() else {
2220 return;
2221 };
2222
2223 self.maybe_update_global_write_frontier(subscribe_id, write_frontier);
2229
2230 match response {
2231 SubscribeResponse::Batch(batch) => {
2232 let upper = batch.upper;
2233 let mut updates = batch.updates;
2234
2235 if PartialOrder::less_than(&subscribe.frontier, &upper) {
2238 let lower = std::mem::replace(&mut subscribe.frontier, upper.clone());
2239
2240 if upper.is_empty() {
2241 self.subscribes.remove(&subscribe_id);
2243 } else {
2244 self.subscribes.insert(subscribe_id, subscribe);
2246 }
2247
2248 if let Ok(updates) = updates.as_mut() {
2249 updates.retain_mut(|updates| {
2250 let offset = updates.times().partition_point(|t| {
2251 !lower.less_equal(t)
2254 });
2255 let (_, past_lower) = std::mem::take(updates).split_at(offset);
2256 *updates = past_lower;
2257 updates.len() > 0
2258 });
2259 }
2260 self.deliver_response(ComputeControllerResponse::SubscribeResponse(
2261 subscribe_id,
2262 SubscribeBatch {
2263 lower,
2264 upper,
2265 updates,
2266 },
2267 ));
2268 }
2269 }
2270 SubscribeResponse::DroppedAt(frontier) => {
2271 tracing::error!(
2276 %subscribe_id,
2277 %replica_id,
2278 frontier = ?frontier.elements(),
2279 "received `DroppedAt` response for a tracked subscribe",
2280 );
2281 self.subscribes.remove(&subscribe_id);
2282 }
2283 }
2284 }
2285
2286 fn handle_status_response(&self, response: StatusResponse, _replica_id: ReplicaId) {
2287 match response {
2288 StatusResponse::Placeholder => {}
2289 }
2290 }
2291
2292 fn dependency_write_frontiers<'b>(
2294 &'b self,
2295 collection: &'b CollectionState,
2296 ) -> impl Iterator<Item = Antichain<Timestamp>> + 'b {
2297 let compute_frontiers = collection.compute_dependency_ids().filter_map(|dep_id| {
2298 let collection = self.collections.get(&dep_id);
2299 collection.map(|c| c.write_frontier())
2300 });
2301 let storage_frontiers = collection.storage_dependency_ids().filter_map(|dep_id| {
2302 let frontiers = self.storage_collections.collection_frontiers(dep_id).ok();
2303 frontiers.map(|f| f.write_frontier)
2304 });
2305
2306 compute_frontiers.chain(storage_frontiers)
2307 }
2308
2309 fn transitive_storage_dependency_write_frontiers<'b>(
2311 &'b self,
2312 collection: &'b CollectionState,
2313 ) -> impl Iterator<Item = Antichain<Timestamp>> + 'b {
2314 let mut storage_ids: BTreeSet<_> = collection.storage_dependency_ids().collect();
2315 let mut todo: Vec<_> = collection.compute_dependency_ids().collect();
2316 let mut done = BTreeSet::new();
2317
2318 while let Some(id) = todo.pop() {
2319 if done.contains(&id) {
2320 continue;
2321 }
2322 if let Some(dep) = self.collections.get(&id) {
2323 storage_ids.extend(dep.storage_dependency_ids());
2324 todo.extend(dep.compute_dependency_ids())
2325 }
2326 done.insert(id);
2327 }
2328
2329 let storage_frontiers = storage_ids.into_iter().filter_map(|id| {
2330 let frontiers = self.storage_collections.collection_frontiers(id).ok();
2331 frontiers.map(|f| f.write_frontier)
2332 });
2333
2334 storage_frontiers
2335 }
2336
2337 fn downgrade_warmup_capabilities(&mut self) {
2350 let mut new_capabilities = BTreeMap::new();
2351 for (id, collection) in &self.collections {
2352 if collection.read_policy.is_none()
2356 && collection.shared.lock_write_frontier(|f| f.is_empty())
2357 {
2358 new_capabilities.insert(*id, Antichain::new());
2359 continue;
2360 }
2361
2362 let mut new_capability = Antichain::new();
2363 for frontier in self.dependency_write_frontiers(collection) {
2364 for time in frontier {
2365 new_capability.insert(time.step_back().unwrap_or(time));
2366 }
2367 }
2368
2369 new_capabilities.insert(*id, new_capability);
2370 }
2371
2372 for (id, new_capability) in new_capabilities {
2373 let collection = self.expect_collection_mut(id);
2374 let _ = collection.warmup_read_hold.try_downgrade(new_capability);
2375 }
2376 }
2377
2378 fn forward_implied_capabilities(&mut self) {
2406 if !ENABLE_PAUSED_CLUSTER_READHOLD_DOWNGRADE.get(&self.dyncfg) {
2407 return;
2408 }
2409 if !self.replicas.is_empty() {
2410 return;
2411 }
2412
2413 let mut new_capabilities = BTreeMap::new();
2414 for (id, collection) in &self.collections {
2415 let Some(read_policy) = &collection.read_policy else {
2416 continue;
2418 };
2419
2420 let mut dep_frontier = Antichain::new();
2424 for frontier in self.transitive_storage_dependency_write_frontiers(collection) {
2425 dep_frontier.extend(frontier);
2426 }
2427
2428 let new_capability = read_policy.frontier(dep_frontier.borrow());
2429 if PartialOrder::less_than(collection.implied_read_hold.since(), &new_capability) {
2430 new_capabilities.insert(*id, new_capability);
2431 }
2432 }
2433
2434 for (id, new_capability) in new_capabilities {
2435 let collection = self.expect_collection_mut(id);
2436 let _ = collection.implied_read_hold.try_downgrade(new_capability);
2437 }
2438 }
2439
2440 pub(super) fn acquire_read_hold(&self, id: GlobalId) -> Result<ReadHold, CollectionMissing> {
2445 let collection = self.collection(id)?;
2451 let since = collection.shared.lock_read_capabilities(|caps| {
2452 let since = caps.frontier().to_owned();
2453 caps.update_iter(since.iter().map(|t| (t.clone(), 1)));
2454 since
2455 });
2456 let hold = ReadHold::new(id, since, Arc::clone(&self.read_hold_tx));
2457 Ok(hold)
2458 }
2459
2460 #[mz_ore::instrument(level = "debug")]
2466 pub fn maintain(&mut self) {
2467 self.rehydrate_failed_replicas();
2468 self.downgrade_warmup_capabilities();
2469 self.forward_implied_capabilities();
2470 self.schedule_collections();
2471 self.cleanup_collections();
2472 self.update_frontier_introspection();
2473 self.refresh_state_metrics();
2474 self.refresh_wallclock_lag();
2475 }
2476}
2477
2478#[derive(Debug)]
2483struct CollectionState {
2484 target_replica: Option<ReplicaId>,
2486 log_collection: bool,
2490 dropped: bool,
2496 scheduled: bool,
2499
2500 read_only: bool,
2504
2505 shared: SharedCollectionState,
2507
2508 implied_read_hold: ReadHold,
2515 warmup_read_hold: ReadHold,
2523 read_policy: Option<ReadPolicy>,
2529
2530 storage_dependencies: BTreeMap<GlobalId, ReadHold>,
2533 compute_dependencies: BTreeMap<GlobalId, ReadHold>,
2536
2537 introspection: CollectionIntrospection,
2539
2540 wallclock_lag_histogram_stash: Option<
2547 BTreeMap<
2548 (
2549 WallclockLagHistogramPeriod,
2550 WallclockLag,
2551 BTreeMap<&'static str, String>,
2552 ),
2553 Diff,
2554 >,
2555 >,
2556}
2557
2558impl CollectionState {
2559 fn new(
2561 collection_id: GlobalId,
2562 as_of: Antichain<Timestamp>,
2563 shared: SharedCollectionState,
2564 storage_dependencies: BTreeMap<GlobalId, ReadHold>,
2565 compute_dependencies: BTreeMap<GlobalId, ReadHold>,
2566 read_hold_tx: read_holds::ChangeTx,
2567 introspection: CollectionIntrospection,
2568 ) -> Self {
2569 let since = as_of.clone();
2571 let upper = as_of;
2573
2574 assert!(shared.lock_read_capabilities(|c| c.frontier() == since.borrow()));
2576 assert!(shared.lock_write_frontier(|f| f == &upper));
2577
2578 let implied_read_hold =
2582 ReadHold::new(collection_id, since.clone(), Arc::clone(&read_hold_tx));
2583 let warmup_read_hold = ReadHold::new(collection_id, since.clone(), read_hold_tx);
2584
2585 let updates = warmup_read_hold.since().iter().map(|t| (t.clone(), 1));
2586 shared.lock_read_capabilities(|c| {
2587 c.update_iter(updates);
2588 });
2589
2590 let wallclock_lag_histogram_stash = match collection_id.is_transient() {
2594 true => None,
2595 false => Some(Default::default()),
2596 };
2597
2598 Self {
2599 target_replica: None,
2600 log_collection: false,
2601 dropped: false,
2602 scheduled: false,
2603 read_only: true,
2604 shared,
2605 implied_read_hold,
2606 warmup_read_hold,
2607 read_policy: Some(ReadPolicy::ValidFrom(since)),
2608 storage_dependencies,
2609 compute_dependencies,
2610 introspection,
2611 wallclock_lag_histogram_stash,
2612 }
2613 }
2614
2615 fn new_log_collection(
2617 id: GlobalId,
2618 shared: SharedCollectionState,
2619 read_hold_tx: read_holds::ChangeTx,
2620 introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
2621 ) -> Self {
2622 let since = Antichain::from_elem(Timestamp::MIN);
2623 let introspection = CollectionIntrospection::new(
2624 id,
2625 introspection_tx,
2626 since.clone(),
2627 false,
2628 None,
2629 None,
2630 Vec::new(),
2631 );
2632 let mut state = Self::new(
2633 id,
2634 since,
2635 shared,
2636 Default::default(),
2637 Default::default(),
2638 read_hold_tx,
2639 introspection,
2640 );
2641 state.log_collection = true;
2642 state.scheduled = true;
2644 state
2645 }
2646
2647 fn read_frontier(&self) -> Antichain<Timestamp> {
2649 self.shared
2650 .lock_read_capabilities(|c| c.frontier().to_owned())
2651 }
2652
2653 fn write_frontier(&self) -> Antichain<Timestamp> {
2655 self.shared.lock_write_frontier(|f| f.clone())
2656 }
2657
2658 fn storage_dependency_ids(&self) -> impl Iterator<Item = GlobalId> + '_ {
2659 self.storage_dependencies.keys().copied()
2660 }
2661
2662 fn compute_dependency_ids(&self) -> impl Iterator<Item = GlobalId> + '_ {
2663 self.compute_dependencies.keys().copied()
2664 }
2665}
2666
2667#[derive(Clone, Debug)]
2678pub(super) struct SharedCollectionState {
2679 read_capabilities: Arc<Mutex<MutableAntichain<Timestamp>>>,
2692 write_frontier: Arc<Mutex<Antichain<Timestamp>>>,
2694}
2695
2696impl SharedCollectionState {
2697 pub fn new(as_of: Antichain<Timestamp>) -> Self {
2698 let since = as_of.clone();
2700 let upper = as_of;
2702
2703 let mut read_capabilities = MutableAntichain::new();
2707 read_capabilities.update_iter(since.iter().map(|time| (time.clone(), 1)));
2708
2709 Self {
2710 read_capabilities: Arc::new(Mutex::new(read_capabilities)),
2711 write_frontier: Arc::new(Mutex::new(upper)),
2712 }
2713 }
2714
2715 pub fn lock_read_capabilities<F, R>(&self, f: F) -> R
2716 where
2717 F: FnOnce(&mut MutableAntichain<Timestamp>) -> R,
2718 {
2719 let mut caps = self.read_capabilities.lock().expect("poisoned");
2720 f(&mut *caps)
2721 }
2722
2723 pub fn lock_write_frontier<F, R>(&self, f: F) -> R
2724 where
2725 F: FnOnce(&mut Antichain<Timestamp>) -> R,
2726 {
2727 let mut frontier = self.write_frontier.lock().expect("poisoned");
2728 f(&mut *frontier)
2729 }
2730}
2731
2732#[derive(Debug)]
2735struct CollectionIntrospection {
2736 collection_id: GlobalId,
2738 introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
2740 frontiers: Option<FrontiersIntrospectionState>,
2745 refresh: Option<RefreshIntrospectionState>,
2749 dependency_ids: Vec<GlobalId>,
2751}
2752
2753impl CollectionIntrospection {
2754 fn new(
2755 collection_id: GlobalId,
2756 introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
2757 as_of: Antichain<Timestamp>,
2758 storage_sink: bool,
2759 initial_as_of: Option<Antichain<Timestamp>>,
2760 refresh_schedule: Option<RefreshSchedule>,
2761 dependency_ids: Vec<GlobalId>,
2762 ) -> Self {
2763 let refresh =
2764 match (refresh_schedule, initial_as_of) {
2765 (Some(refresh_schedule), Some(initial_as_of)) => Some(
2766 RefreshIntrospectionState::new(refresh_schedule, initial_as_of, &as_of),
2767 ),
2768 (refresh_schedule, _) => {
2769 soft_assert_or_log!(
2772 refresh_schedule.is_none(),
2773 "`refresh_schedule` without an `initial_as_of`: {collection_id}"
2774 );
2775 None
2776 }
2777 };
2778 let frontiers = (!storage_sink).then(|| FrontiersIntrospectionState::new(as_of));
2779
2780 let self_ = Self {
2781 collection_id,
2782 introspection_tx,
2783 frontiers,
2784 refresh,
2785 dependency_ids,
2786 };
2787
2788 self_.report_initial_state();
2789 self_
2790 }
2791
2792 fn report_initial_state(&self) {
2794 if let Some(frontiers) = &self.frontiers {
2795 let row = frontiers.row_for_collection(self.collection_id);
2796 let updates = vec![(row, Diff::ONE)];
2797 self.send(IntrospectionType::Frontiers, updates);
2798 }
2799
2800 if let Some(refresh) = &self.refresh {
2801 let row = refresh.row_for_collection(self.collection_id);
2802 let updates = vec![(row, Diff::ONE)];
2803 self.send(IntrospectionType::ComputeMaterializedViewRefreshes, updates);
2804 }
2805
2806 if !self.dependency_ids.is_empty() {
2807 let updates = self.dependency_rows(Diff::ONE);
2808 self.send(IntrospectionType::ComputeDependencies, updates);
2809 }
2810 }
2811
2812 fn dependency_rows(&self, diff: Diff) -> Vec<(Row, Diff)> {
2814 self.dependency_ids
2815 .iter()
2816 .map(|dependency_id| {
2817 let row = Row::pack_slice(&[
2818 Datum::String(&self.collection_id.to_string()),
2819 Datum::String(&dependency_id.to_string()),
2820 ]);
2821 (row, diff)
2822 })
2823 .collect()
2824 }
2825
2826 fn observe_frontiers(
2829 &mut self,
2830 read_frontier: &Antichain<Timestamp>,
2831 write_frontier: &Antichain<Timestamp>,
2832 ) {
2833 self.update_frontier_introspection(read_frontier, write_frontier);
2834 self.update_refresh_introspection(write_frontier);
2835 }
2836
2837 fn update_frontier_introspection(
2838 &mut self,
2839 read_frontier: &Antichain<Timestamp>,
2840 write_frontier: &Antichain<Timestamp>,
2841 ) {
2842 let Some(frontiers) = &mut self.frontiers else {
2843 return;
2844 };
2845
2846 if &frontiers.read_frontier == read_frontier && &frontiers.write_frontier == write_frontier
2847 {
2848 return; };
2850
2851 let retraction = frontiers.row_for_collection(self.collection_id);
2852 frontiers.update(read_frontier, write_frontier);
2853 let insertion = frontiers.row_for_collection(self.collection_id);
2854 let updates = vec![(retraction, Diff::MINUS_ONE), (insertion, Diff::ONE)];
2855 self.send(IntrospectionType::Frontiers, updates);
2856 }
2857
2858 fn update_refresh_introspection(&mut self, write_frontier: &Antichain<Timestamp>) {
2859 let Some(refresh) = &mut self.refresh else {
2860 return;
2861 };
2862
2863 let retraction = refresh.row_for_collection(self.collection_id);
2864 refresh.frontier_update(write_frontier);
2865 let insertion = refresh.row_for_collection(self.collection_id);
2866
2867 if retraction == insertion {
2868 return; }
2870
2871 let updates = vec![(retraction, Diff::MINUS_ONE), (insertion, Diff::ONE)];
2872 self.send(IntrospectionType::ComputeMaterializedViewRefreshes, updates);
2873 }
2874
2875 fn send(&self, introspection_type: IntrospectionType, updates: Vec<(Row, Diff)>) {
2876 let _ = self.introspection_tx.send((introspection_type, updates));
2879 }
2880}
2881
2882impl Drop for CollectionIntrospection {
2883 fn drop(&mut self) {
2884 if let Some(frontiers) = &self.frontiers {
2886 let row = frontiers.row_for_collection(self.collection_id);
2887 let updates = vec![(row, Diff::MINUS_ONE)];
2888 self.send(IntrospectionType::Frontiers, updates);
2889 }
2890
2891 if let Some(refresh) = &self.refresh {
2893 let retraction = refresh.row_for_collection(self.collection_id);
2894 let updates = vec![(retraction, Diff::MINUS_ONE)];
2895 self.send(IntrospectionType::ComputeMaterializedViewRefreshes, updates);
2896 }
2897
2898 if !self.dependency_ids.is_empty() {
2900 let updates = self.dependency_rows(Diff::MINUS_ONE);
2901 self.send(IntrospectionType::ComputeDependencies, updates);
2902 }
2903 }
2904}
2905
2906#[derive(Debug)]
2907struct FrontiersIntrospectionState {
2908 read_frontier: Antichain<Timestamp>,
2909 write_frontier: Antichain<Timestamp>,
2910}
2911
2912impl FrontiersIntrospectionState {
2913 fn new(as_of: Antichain<Timestamp>) -> Self {
2914 Self {
2915 read_frontier: as_of.clone(),
2916 write_frontier: as_of,
2917 }
2918 }
2919
2920 fn row_for_collection(&self, collection_id: GlobalId) -> Row {
2922 let read_frontier = self
2923 .read_frontier
2924 .as_option()
2925 .map_or(Datum::Null, |ts| ts.clone().into());
2926 let write_frontier = self
2927 .write_frontier
2928 .as_option()
2929 .map_or(Datum::Null, |ts| ts.clone().into());
2930 Row::pack_slice(&[
2931 Datum::String(&collection_id.to_string()),
2932 read_frontier,
2933 write_frontier,
2934 ])
2935 }
2936
2937 fn update(
2939 &mut self,
2940 read_frontier: &Antichain<Timestamp>,
2941 write_frontier: &Antichain<Timestamp>,
2942 ) {
2943 if read_frontier != &self.read_frontier {
2944 self.read_frontier.clone_from(read_frontier);
2945 }
2946 if write_frontier != &self.write_frontier {
2947 self.write_frontier.clone_from(write_frontier);
2948 }
2949 }
2950}
2951
2952#[derive(Debug)]
2955struct RefreshIntrospectionState {
2956 refresh_schedule: RefreshSchedule,
2958 initial_as_of: Antichain<Timestamp>,
2959 next_refresh: Datum<'static>, last_completed_refresh: Datum<'static>, }
2963
2964impl RefreshIntrospectionState {
2965 fn row_for_collection(&self, collection_id: GlobalId) -> Row {
2967 Row::pack_slice(&[
2968 Datum::String(&collection_id.to_string()),
2969 self.last_completed_refresh,
2970 self.next_refresh,
2971 ])
2972 }
2973}
2974
2975impl RefreshIntrospectionState {
2976 fn new(
2979 refresh_schedule: RefreshSchedule,
2980 initial_as_of: Antichain<Timestamp>,
2981 upper: &Antichain<Timestamp>,
2982 ) -> Self {
2983 let mut self_ = Self {
2984 refresh_schedule: refresh_schedule.clone(),
2985 initial_as_of: initial_as_of.clone(),
2986 next_refresh: Datum::Null,
2987 last_completed_refresh: Datum::Null,
2988 };
2989 self_.frontier_update(upper);
2990 self_
2991 }
2992
2993 fn frontier_update(&mut self, write_frontier: &Antichain<Timestamp>) {
2996 if write_frontier.is_empty() {
2997 self.last_completed_refresh =
2998 if let Some(last_refresh) = self.refresh_schedule.last_refresh() {
2999 last_refresh.into()
3000 } else {
3001 Timestamp::MAX.into()
3004 };
3005 self.next_refresh = Datum::Null;
3006 } else {
3007 if PartialOrder::less_equal(write_frontier, &self.initial_as_of) {
3008 self.last_completed_refresh = Datum::Null;
3010 let initial_as_of = self.initial_as_of.as_option().expect(
3011 "initial_as_of can't be [], because then there would be no refreshes at all",
3012 );
3013 let first_refresh = self
3014 .refresh_schedule
3015 .round_up_timestamp(*initial_as_of)
3016 .expect("sequencing makes sure that REFRESH MVs always have a first refresh");
3017 soft_assert_or_log!(
3018 first_refresh == *initial_as_of,
3019 "initial_as_of should be set to the first refresh"
3020 );
3021 self.next_refresh = first_refresh.into();
3022 } else {
3023 let write_frontier = write_frontier.as_option().expect("checked above");
3025 self.last_completed_refresh = self
3026 .refresh_schedule
3027 .round_down_timestamp_m1(*write_frontier)
3028 .map_or_else(
3029 || {
3030 soft_panic_or_log!(
3031 "rounding down should have returned the first refresh or later"
3032 );
3033 Datum::Null
3034 },
3035 |last_completed_refresh| last_completed_refresh.into(),
3036 );
3037 self.next_refresh = write_frontier.clone().into();
3038 }
3039 }
3040 }
3041}
3042
3043#[derive(Debug)]
3045struct PendingPeek {
3046 target_replica: Option<ReplicaId>,
3050 otel_ctx: OpenTelemetryContext,
3052 requested_at: Instant,
3056 read_hold: ReadHold,
3058 peek_response_tx: oneshot::Sender<PeekResponse>,
3060 limit: Option<usize>,
3062 offset: usize,
3064}
3065
3066#[derive(Debug, Clone)]
3067struct ActiveSubscribe {
3068 frontier: Antichain<Timestamp>,
3070}
3071
3072impl Default for ActiveSubscribe {
3073 fn default() -> Self {
3074 Self {
3075 frontier: Antichain::from_elem(Timestamp::MIN),
3076 }
3077 }
3078}
3079
3080#[derive(Debug)]
3082struct ReplicaState {
3083 id: ReplicaId,
3085 client: ReplicaClient,
3087 config: ReplicaConfig,
3089 metrics: ReplicaMetrics,
3091 introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
3093 collections: BTreeMap<GlobalId, ReplicaCollectionState>,
3095 epoch: u64,
3097}
3098
3099impl ReplicaState {
3100 fn new(
3101 id: ReplicaId,
3102 client: ReplicaClient,
3103 config: ReplicaConfig,
3104 metrics: ReplicaMetrics,
3105 introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
3106 epoch: u64,
3107 ) -> Self {
3108 Self {
3109 id,
3110 client,
3111 config,
3112 metrics,
3113 introspection_tx,
3114 epoch,
3115 collections: Default::default(),
3116 }
3117 }
3118
3119 fn add_collection(
3125 &mut self,
3126 id: GlobalId,
3127 as_of: Antichain<Timestamp>,
3128 input_read_holds: Vec<ReadHold>,
3129 ) {
3130 let metrics = self.metrics.for_collection(id);
3131 let introspection = ReplicaCollectionIntrospection::new(
3132 self.id,
3133 id,
3134 self.introspection_tx.clone(),
3135 as_of.clone(),
3136 );
3137 let mut state =
3138 ReplicaCollectionState::new(metrics, as_of, introspection, input_read_holds);
3139
3140 if id.is_transient() {
3144 state.wallclock_lag_max = None;
3145 }
3146
3147 if let Some(previous) = self.collections.insert(id, state) {
3148 panic!("attempt to add a collection with existing ID {id} (previous={previous:?}");
3149 }
3150 }
3151
3152 fn remove_collection(&mut self, id: GlobalId) -> Option<ReplicaCollectionState> {
3154 self.collections.remove(&id)
3155 }
3156
3157 fn collection_frontiers_empty(&self, id: GlobalId) -> bool {
3159 self.collections.get(&id).map_or(true, |c| {
3160 c.write_frontier.is_empty()
3161 && c.input_frontier.is_empty()
3162 && c.output_frontier.is_empty()
3163 })
3164 }
3165
3166 #[mz_ore::instrument(level = "debug")]
3170 pub fn dump(&self) -> Result<serde_json::Value, anyhow::Error> {
3171 let Self {
3178 id,
3179 client: _,
3180 config: _,
3181 metrics: _,
3182 introspection_tx: _,
3183 epoch,
3184 collections,
3185 } = self;
3186
3187 let collections: BTreeMap<_, _> = collections
3188 .iter()
3189 .map(|(id, collection)| (id.to_string(), format!("{collection:?}")))
3190 .collect();
3191
3192 Ok(serde_json::json!({
3193 "id": id.to_string(),
3194 "collections": collections,
3195 "epoch": epoch,
3196 }))
3197 }
3198}
3199
3200#[derive(Debug)]
3201struct ReplicaCollectionState {
3202 write_frontier: Antichain<Timestamp>,
3206 input_frontier: Antichain<Timestamp>,
3210 output_frontier: Antichain<Timestamp>,
3214
3215 metrics: Option<ReplicaCollectionMetrics>,
3219 as_of: Antichain<Timestamp>,
3221 introspection: ReplicaCollectionIntrospection,
3223 input_read_holds: Vec<ReadHold>,
3229
3230 wallclock_lag_max: Option<WallclockLag>,
3234}
3235
3236impl ReplicaCollectionState {
3237 fn new(
3238 metrics: Option<ReplicaCollectionMetrics>,
3239 as_of: Antichain<Timestamp>,
3240 introspection: ReplicaCollectionIntrospection,
3241 input_read_holds: Vec<ReadHold>,
3242 ) -> Self {
3243 Self {
3244 write_frontier: as_of.clone(),
3245 input_frontier: as_of.clone(),
3246 output_frontier: as_of.clone(),
3247 metrics,
3248 as_of,
3249 introspection,
3250 input_read_holds,
3251 wallclock_lag_max: Some(WallclockLag::MIN),
3252 }
3253 }
3254
3255 fn hydrated(&self) -> bool {
3257 self.as_of.is_empty() || PartialOrder::less_than(&self.as_of, &self.output_frontier)
3273 }
3274
3275 fn update_write_frontier(&mut self, new_frontier: Antichain<Timestamp>) {
3277 if PartialOrder::less_than(&new_frontier, &self.write_frontier) {
3278 soft_panic_or_log!(
3279 "replica collection write frontier regression (old={:?}, new={new_frontier:?})",
3280 self.write_frontier,
3281 );
3282 return;
3283 } else if new_frontier == self.write_frontier {
3284 return;
3285 }
3286
3287 self.write_frontier = new_frontier;
3288 }
3289
3290 fn update_input_frontier(&mut self, new_frontier: Antichain<Timestamp>) {
3292 if PartialOrder::less_than(&new_frontier, &self.input_frontier) {
3293 soft_panic_or_log!(
3294 "replica collection input frontier regression (old={:?}, new={new_frontier:?})",
3295 self.input_frontier,
3296 );
3297 return;
3298 } else if new_frontier == self.input_frontier {
3299 return;
3300 }
3301
3302 self.input_frontier = new_frontier;
3303
3304 for read_hold in &mut self.input_read_holds {
3306 let result = read_hold.try_downgrade(self.input_frontier.clone());
3307 soft_assert_or_log!(
3308 result.is_ok(),
3309 "read hold downgrade failed (read_hold={read_hold:?}, new_since={:?})",
3310 self.input_frontier,
3311 );
3312 }
3313 }
3314
3315 fn update_output_frontier(&mut self, new_frontier: Antichain<Timestamp>) {
3317 if PartialOrder::less_than(&new_frontier, &self.output_frontier) {
3318 soft_panic_or_log!(
3319 "replica collection output frontier regression (old={:?}, new={new_frontier:?})",
3320 self.output_frontier,
3321 );
3322 return;
3323 } else if new_frontier == self.output_frontier {
3324 return;
3325 }
3326
3327 self.output_frontier = new_frontier;
3328 }
3329}
3330
3331#[derive(Debug)]
3334struct ReplicaCollectionIntrospection {
3335 replica_id: ReplicaId,
3337 collection_id: GlobalId,
3339 write_frontier: Antichain<Timestamp>,
3341 introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
3343}
3344
3345impl ReplicaCollectionIntrospection {
3346 fn new(
3348 replica_id: ReplicaId,
3349 collection_id: GlobalId,
3350 introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
3351 as_of: Antichain<Timestamp>,
3352 ) -> Self {
3353 let self_ = Self {
3354 replica_id,
3355 collection_id,
3356 write_frontier: as_of,
3357 introspection_tx,
3358 };
3359
3360 self_.report_initial_state();
3361 self_
3362 }
3363
3364 fn report_initial_state(&self) {
3366 let row = self.write_frontier_row();
3367 let updates = vec![(row, Diff::ONE)];
3368 self.send(IntrospectionType::ReplicaFrontiers, updates);
3369 }
3370
3371 fn observe_frontier(&mut self, write_frontier: &Antichain<Timestamp>) {
3373 if self.write_frontier == *write_frontier {
3374 return; }
3376
3377 let retraction = self.write_frontier_row();
3378 self.write_frontier.clone_from(write_frontier);
3379 let insertion = self.write_frontier_row();
3380
3381 let updates = vec![(retraction, Diff::MINUS_ONE), (insertion, Diff::ONE)];
3382 self.send(IntrospectionType::ReplicaFrontiers, updates);
3383 }
3384
3385 fn write_frontier_row(&self) -> Row {
3387 let write_frontier = self
3388 .write_frontier
3389 .as_option()
3390 .map_or(Datum::Null, |ts| ts.clone().into());
3391 Row::pack_slice(&[
3392 Datum::String(&self.collection_id.to_string()),
3393 Datum::String(&self.replica_id.to_string()),
3394 write_frontier,
3395 ])
3396 }
3397
3398 fn send(&self, introspection_type: IntrospectionType, updates: Vec<(Row, Diff)>) {
3399 let _ = self.introspection_tx.send((introspection_type, updates));
3402 }
3403}
3404
3405impl Drop for ReplicaCollectionIntrospection {
3406 fn drop(&mut self) {
3407 let row = self.write_frontier_row();
3409 let updates = vec![(row, Diff::MINUS_ONE)];
3410 self.send(IntrospectionType::ReplicaFrontiers, updates);
3411 }
3412}
3413
3414#[cfg(test)]
3415mod tests {
3416 use std::collections::BTreeMap;
3417
3418 use mz_compute_types::dyncfgs::{ENABLE_COLUMN_PAGED_BATCHER, ENABLE_MZ_JOIN_CORE};
3419 use mz_dyncfg::{ConfigSet, ConfigUpdates, ConfigVal};
3420 use mz_persist_types::PersistLocation;
3421
3422 use crate::protocol::command::{ComputeCommand, InstanceConfig};
3423
3424 use super::{Instance, ReplicaId};
3425
3426 fn create_instance_command() -> ComputeCommand {
3427 ComputeCommand::CreateInstance(Box::new(InstanceConfig {
3428 logging: Default::default(),
3429 expiration_offset: None,
3430 peek_stash_persist_location: PersistLocation::new_in_mem(),
3431 arrangement_dictionary_compression: false,
3432 initial_config: Default::default(),
3433 }))
3434 }
3435
3436 fn initial_config(cmd: &ComputeCommand) -> &ConfigUpdates {
3437 match cmd {
3438 ComputeCommand::CreateInstance(config) => &config.initial_config,
3439 other => panic!("expected CreateInstance, got {other:?}"),
3440 }
3441 }
3442
3443 #[mz_ore::test]
3448 fn create_instance_snapshots_instance_wide_dyncfg() {
3449 let dyncfg = ConfigSet::default()
3450 .add(&ENABLE_COLUMN_PAGED_BATCHER)
3451 .add(&ENABLE_MZ_JOIN_CORE);
3452 let mut updates = ConfigUpdates::default();
3453 updates.add(&ENABLE_COLUMN_PAGED_BATCHER, true);
3454 updates.add(&ENABLE_MZ_JOIN_CORE, false);
3455 updates.apply(&dyncfg);
3456
3457 let overrides = BTreeMap::new();
3459 let cmd = Instance::specialize_command_for_replica(
3460 create_instance_command(),
3461 ReplicaId::User(1),
3462 &overrides,
3463 &dyncfg,
3464 );
3465 let snapshot = initial_config(&cmd);
3466 assert_eq!(
3467 snapshot.updates.get(ENABLE_COLUMN_PAGED_BATCHER.name()),
3468 Some(&ConfigVal::Bool(true)),
3469 );
3470 assert_eq!(
3471 snapshot.updates.get(ENABLE_MZ_JOIN_CORE.name()),
3472 Some(&ConfigVal::Bool(false)),
3473 );
3474 }
3475
3476 #[mz_ore::test]
3479 fn create_instance_snapshot_applies_replica_override() {
3480 let dyncfg = ConfigSet::default().add(&ENABLE_COLUMN_PAGED_BATCHER);
3481 let mut updates = ConfigUpdates::default();
3482 updates.add(&ENABLE_COLUMN_PAGED_BATCHER, true);
3483 updates.apply(&dyncfg);
3484
3485 let replica = ReplicaId::User(1);
3486 let mut override_updates = ConfigUpdates::default();
3487 override_updates.add(&ENABLE_COLUMN_PAGED_BATCHER, false);
3488 let overrides = BTreeMap::from([(replica, override_updates)]);
3489
3490 let cmd = Instance::specialize_command_for_replica(
3491 create_instance_command(),
3492 replica,
3493 &overrides,
3494 &dyncfg,
3495 );
3496 assert_eq!(
3497 initial_config(&cmd)
3498 .updates
3499 .get(ENABLE_COLUMN_PAGED_BATCHER.name()),
3500 Some(&ConfigVal::Bool(false)),
3501 "replica override should win over the instance-wide value",
3502 );
3503 }
3504
3505 #[mz_ore::test]
3507 fn update_configuration_merges_replica_override() {
3508 let dyncfg = ConfigSet::default().add(&ENABLE_COLUMN_PAGED_BATCHER);
3509
3510 let replica = ReplicaId::User(1);
3511 let mut override_updates = ConfigUpdates::default();
3512 override_updates.add(&ENABLE_COLUMN_PAGED_BATCHER, true);
3513 let overrides = BTreeMap::from([(replica, override_updates)]);
3514
3515 let cmd = Instance::specialize_command_for_replica(
3516 ComputeCommand::UpdateConfiguration(Box::new(Default::default())),
3517 replica,
3518 &overrides,
3519 &dyncfg,
3520 );
3521 match cmd {
3522 ComputeCommand::UpdateConfiguration(params) => assert_eq!(
3523 params
3524 .dyncfg_updates
3525 .updates
3526 .get(ENABLE_COLUMN_PAGED_BATCHER.name()),
3527 Some(&ConfigVal::Bool(true)),
3528 ),
3529 other => panic!("expected UpdateConfiguration, got {other:?}"),
3530 }
3531 }
3532}