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 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 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 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 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 fn maybe_schedule_collection(&mut self, id: GlobalId) {
1632 let collection = self.expect_collection(id);
1633
1634 if collection.scheduled {
1636 return;
1637 }
1638
1639 let as_of = collection.read_frontier();
1640
1641 if as_of.is_empty() {
1644 return;
1645 }
1646
1647 let ready = if id.is_transient() {
1648 true
1654 } else {
1655 let not_self_dep = |x: &GlobalId| *x != id;
1661
1662 let mut deps_scheduled = true;
1665
1666 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 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 #[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 collection.dropped = true;
1716
1717 collection.implied_read_hold.release();
1720 collection.warmup_read_hold.release();
1721
1722 self.subscribes.remove(id);
1725 self.copy_tos.remove(id);
1728 }
1729
1730 Ok(())
1731 }
1732
1733 #[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 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 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 otel_ctx,
1793 target: peek_target,
1794 result_desc,
1795 };
1796 self.send(ComputeCommand::Peek(Box::new(peek)));
1797
1798 Ok(())
1799 }
1800
1801 #[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 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 self.finish_peek(uuid, reason);
1826 }
1827
1828 #[mz_ore::instrument(level = "debug")]
1840 pub fn set_read_policy(
1841 &mut self,
1842 policies: Vec<(GlobalId, ReadPolicy)>,
1843 ) -> Result<(), ReadPolicyError> {
1844 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 #[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 let new_since = match &collection.read_policy {
1892 Some(read_policy) => {
1893 read_policy.frontier(new_frontier.borrow())
1896 }
1897 None => {
1898 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 self.deliver_response(ComputeControllerResponse::FrontierUpper {
1917 id,
1918 upper: new_frontier,
1919 });
1920 }
1921
1922 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 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 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; };
1964
1965 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 self.send(ComputeCommand::AllowCompaction {
1979 id,
1980 frontier: new_since,
1981 });
1982 }
1983
1984 fn finish_peek(&mut self, uuid: Uuid, response: PeekResponse) {
1993 let Some(peek) = self.peeks.remove(&uuid) else {
1994 return;
1995 };
1996
1997 let _ = peek.peek_response_tx.send(response);
1999
2000 self.send(ComputeCommand::CancelPeek { uuid });
2003
2004 drop(peek.read_hold);
2005 }
2006
2007 fn handle_response(&mut self, (replica_id, epoch, response): ReplicaResponse) {
2010 if self
2012 .replicas
2013 .get(&replica_id)
2014 .filter(|replica| replica.epoch == epoch)
2015 .is_none()
2016 {
2017 return;
2018 }
2019
2020 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 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 let Some(peek) = self.peeks.get(&uuid) else {
2096 return;
2097 };
2098
2099 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 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 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 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 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 let write_frontier = match &response {
2208 SubscribeResponse::Batch(batch) => batch.upper.clone(),
2209 SubscribeResponse::DroppedAt(_) => Antichain::new(),
2210 };
2211
2212 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 let Some(mut subscribe) = self.subscribes.get(&subscribe_id).cloned() else {
2221 return;
2222 };
2223
2224 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 PartialOrder::less_than(&subscribe.frontier, &upper) {
2239 let lower = std::mem::replace(&mut subscribe.frontier, upper.clone());
2240
2241 if upper.is_empty() {
2242 self.subscribes.remove(&subscribe_id);
2244 } else {
2245 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 !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 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 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 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 fn downgrade_warmup_capabilities(&mut self) {
2351 let mut new_capabilities = BTreeMap::new();
2352 for (id, collection) in &self.collections {
2353 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 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 continue;
2419 };
2420
2421 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 pub(super) fn acquire_read_hold(&self, id: GlobalId) -> Result<ReadHold, CollectionMissing> {
2446 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 #[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#[derive(Debug)]
2484struct CollectionState {
2485 target_replica: Option<ReplicaId>,
2487 log_collection: bool,
2491 dropped: bool,
2497 scheduled: bool,
2500
2501 read_only: bool,
2505
2506 shared: SharedCollectionState,
2508
2509 implied_read_hold: ReadHold,
2516 warmup_read_hold: ReadHold,
2524 read_policy: Option<ReadPolicy>,
2530
2531 storage_dependencies: BTreeMap<GlobalId, ReadHold>,
2534 compute_dependencies: BTreeMap<GlobalId, ReadHold>,
2537
2538 introspection: CollectionIntrospection,
2540
2541 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 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 let since = as_of.clone();
2572 let upper = as_of;
2574
2575 assert!(shared.lock_read_capabilities(|c| c.frontier() == since.borrow()));
2577 assert!(shared.lock_write_frontier(|f| f == &upper));
2578
2579 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 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 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 state.scheduled = true;
2645 state
2646 }
2647
2648 fn read_frontier(&self) -> Antichain<Timestamp> {
2650 self.shared
2651 .lock_read_capabilities(|c| c.frontier().to_owned())
2652 }
2653
2654 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#[derive(Clone, Debug)]
2679pub(super) struct SharedCollectionState {
2680 read_capabilities: Arc<Mutex<MutableAntichain<Timestamp>>>,
2693 write_frontier: Arc<Mutex<Antichain<Timestamp>>>,
2695}
2696
2697impl SharedCollectionState {
2698 pub fn new(as_of: Antichain<Timestamp>) -> Self {
2699 let since = as_of.clone();
2701 let upper = as_of;
2703
2704 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#[derive(Debug)]
2736struct CollectionIntrospection {
2737 collection_id: GlobalId,
2739 introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
2741 frontiers: Option<FrontiersIntrospectionState>,
2746 refresh: Option<RefreshIntrospectionState>,
2750 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 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 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 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 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; };
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; }
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 let _ = self.introspection_tx.send((introspection_type, updates));
2880 }
2881}
2882
2883impl Drop for CollectionIntrospection {
2884 fn drop(&mut self) {
2885 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 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 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 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 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#[derive(Debug)]
2956struct RefreshIntrospectionState {
2957 refresh_schedule: RefreshSchedule,
2959 initial_as_of: Antichain<Timestamp>,
2960 next_refresh: Datum<'static>, last_completed_refresh: Datum<'static>, }
2964
2965impl RefreshIntrospectionState {
2966 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 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 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 Timestamp::MAX.into()
3005 };
3006 self.next_refresh = Datum::Null;
3007 } else {
3008 if PartialOrder::less_equal(write_frontier, &self.initial_as_of) {
3009 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 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#[derive(Debug)]
3046struct PendingPeek {
3047 target_replica: Option<ReplicaId>,
3051 otel_ctx: OpenTelemetryContext,
3053 requested_at: Instant,
3057 read_hold: ReadHold,
3059 peek_response_tx: oneshot::Sender<PeekResponse>,
3061 limit: Option<usize>,
3063 offset: usize,
3065}
3066
3067#[derive(Debug, Clone)]
3068struct ActiveSubscribe {
3069 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#[derive(Debug)]
3083struct ReplicaState {
3084 id: ReplicaId,
3086 client: ReplicaClient,
3088 config: ReplicaConfig,
3090 metrics: ReplicaMetrics,
3092 introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
3094 collections: BTreeMap<GlobalId, ReplicaCollectionState>,
3096 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 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 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 fn remove_collection(&mut self, id: GlobalId) -> Option<ReplicaCollectionState> {
3155 self.collections.remove(&id)
3156 }
3157
3158 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 #[mz_ore::instrument(level = "debug")]
3171 pub fn dump(&self) -> Result<serde_json::Value, anyhow::Error> {
3172 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 write_frontier: Antichain<Timestamp>,
3207 input_frontier: Antichain<Timestamp>,
3211 output_frontier: Antichain<Timestamp>,
3215
3216 metrics: Option<ReplicaCollectionMetrics>,
3220 as_of: Antichain<Timestamp>,
3222 introspection: ReplicaCollectionIntrospection,
3224 input_read_holds: Vec<ReadHold>,
3230
3231 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 fn hydrated(&self) -> bool {
3258 self.as_of.is_empty() || PartialOrder::less_than(&self.as_of, &self.output_frontier)
3274 }
3275
3276 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 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 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 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#[derive(Debug)]
3335struct ReplicaCollectionIntrospection {
3336 replica_id: ReplicaId,
3338 collection_id: GlobalId,
3340 write_frontier: Antichain<Timestamp>,
3342 introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
3344}
3345
3346impl ReplicaCollectionIntrospection {
3347 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 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 fn observe_frontier(&mut self, write_frontier: &Antichain<Timestamp>) {
3374 if self.write_frontier == *write_frontier {
3375 return; }
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 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 let _ = self.introspection_tx.send((introspection_type, updates));
3403 }
3404}
3405
3406impl Drop for ReplicaCollectionIntrospection {
3407 fn drop(&mut self) {
3408 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 #[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 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 #[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 #[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}