1use std::collections::{BTreeMap, BTreeSet};
36use std::sync::{Arc, Mutex};
37use std::time::Duration;
38
39use mz_build_info::BuildInfo;
40use mz_cluster_client::client::ClusterReplicaLocation;
41use mz_cluster_client::metrics::ControllerMetrics;
42use mz_cluster_client::{ReplicaId, WallclockLagFn};
43use mz_compute_types::ComputeInstanceId;
44use mz_compute_types::config::ComputeReplicaConfig;
45use mz_compute_types::dataflows::DataflowDescription;
46use mz_compute_types::dyncfgs::{
47 COMPUTE_REPLICA_EXPIRATION_OFFSET, ENABLE_ARRANGEMENT_DICTIONARY_COMPRESSION_ALPHA,
48};
49use mz_dyncfg::{ConfigSet, ConfigUpdates};
50use mz_expr::RowSetFinishing;
51use mz_expr::row::RowCollection;
52use mz_ore::cast::CastFrom;
53use mz_ore::metrics::MetricsRegistry;
54use mz_ore::now::NowFn;
55use mz_ore::soft_assert_or_log;
56use mz_ore::soft_panic_or_log;
57use mz_ore::tracing::OpenTelemetryContext;
58use mz_persist_types::PersistLocation;
59use mz_repr::{GlobalId, RelationDesc, Row, Timestamp};
60use mz_storage_client::controller::StorageController;
61use mz_storage_types::dyncfgs::ORE_OVERFLOWING_BEHAVIOR;
62use mz_storage_types::read_holds::ReadHold;
63use mz_storage_types::read_policy::ReadPolicy;
64use mz_storage_types::time_dependence::{TimeDependence, TimeDependenceError};
65use prometheus::proto::LabelPair;
66use serde::{Deserialize, Serialize};
67use timely::PartialOrder;
68use timely::progress::Antichain;
69use tokio::sync::{mpsc, oneshot};
70use tokio::time::{self, MissedTickBehavior};
71use uuid::Uuid;
72
73use crate::controller::error::{
74 CollectionLookupError, CollectionMissing, CollectionUpdateError, DataflowCreationError,
75 HydrationCheckBadTarget, InstanceExists, InstanceMissing, PeekError, ReadPolicyError,
76 ReplicaCreationError, ReplicaDropError,
77};
78use crate::controller::instance::{Instance, SharedCollectionState};
79use crate::controller::introspection::{IntrospectionUpdates, spawn_introspection_sink};
80use crate::controller::replica::ReplicaConfig;
81use crate::logging::{LogVariant, LoggingConfig};
82use crate::metrics::ComputeControllerMetrics;
83use crate::protocol::command::{ComputeParameters, PeekTarget};
84use crate::protocol::response::{PeekResponse, SubscribeBatch};
85
86mod instance;
87mod introspection;
88mod replica;
89mod sequential_hydration;
90
91pub mod error;
92pub mod instance_client;
93pub use instance_client::InstanceClient;
94
95pub(crate) type StorageCollections =
96 Arc<dyn mz_storage_client::storage_collections::StorageCollections + Send + Sync>;
97
98#[derive(Debug)]
100pub enum ComputeControllerResponse {
101 PeekNotification(Uuid, PeekNotification, OpenTelemetryContext),
103 SubscribeResponse(GlobalId, SubscribeBatch),
105 CopyToResponse(GlobalId, Result<u64, anyhow::Error>),
116 FrontierUpper {
121 id: GlobalId,
123 upper: Antichain<Timestamp>,
125 },
126}
127
128#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
130pub enum PeekNotification {
131 Success {
133 rows: u64,
135 result_size: u64,
137 },
138 Error(String),
140 Canceled,
142}
143
144impl PeekNotification {
145 fn new(peek_response: &PeekResponse, offset: usize, limit: Option<usize>) -> Self {
148 match peek_response {
149 PeekResponse::Rows(rows) => {
150 let num_rows = u64::cast_from(RowCollection::offset_limit(
151 rows.iter().map(|r| r.count()).sum(),
152 offset,
153 limit,
154 ));
155 let result_size = u64::cast_from(rows.iter().map(|r| r.byte_len()).sum::<usize>());
156
157 tracing::trace!(?num_rows, ?result_size, "inline peek result");
158
159 Self::Success {
160 rows: num_rows,
161 result_size,
162 }
163 }
164 PeekResponse::Stashed(stashed_response) => {
165 let rows = stashed_response.num_rows(offset, limit);
166 let result_size = stashed_response.size_bytes();
167
168 tracing::trace!(?rows, ?result_size, "stashed peek result");
169
170 Self::Success {
171 rows: u64::cast_from(rows),
172 result_size: u64::cast_from(result_size),
173 }
174 }
175 PeekResponse::Error(err) => Self::Error(err.to_string()),
176 PeekResponse::Canceled => Self::Canceled,
177 }
178 }
179}
180
181pub struct ComputeController {
183 instances: BTreeMap<ComputeInstanceId, InstanceState>,
184 instance_workload_classes: Arc<Mutex<BTreeMap<ComputeInstanceId, Option<String>>>>,
188 build_info: &'static BuildInfo,
189 storage_collections: StorageCollections,
191 initialized: bool,
193 read_only: bool,
199 config: ComputeParameters,
201 peek_stash_persist_location: PersistLocation,
203 stashed_response: Option<ComputeControllerResponse>,
205 metrics: ComputeControllerMetrics,
207 now: NowFn,
209 wallclock_lag: WallclockLagFn<Timestamp>,
211 dyncfg: Arc<ConfigSet>,
216 replica_dyncfg_overrides: BTreeMap<ReplicaId, ConfigUpdates>,
222
223 response_rx: mpsc::UnboundedReceiver<ComputeControllerResponse>,
225 response_tx: mpsc::UnboundedSender<ComputeControllerResponse>,
227 introspection_rx: Option<mpsc::UnboundedReceiver<IntrospectionUpdates>>,
232 introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
234
235 maintenance_ticker: tokio::time::Interval,
237 maintenance_scheduled: bool,
239}
240
241impl ComputeController {
242 pub fn new(
244 build_info: &'static BuildInfo,
245 storage_collections: StorageCollections,
246 read_only: bool,
247 metrics_registry: &MetricsRegistry,
248 peek_stash_persist_location: PersistLocation,
249 controller_metrics: ControllerMetrics,
250 now: NowFn,
251 wallclock_lag: WallclockLagFn<Timestamp>,
252 ) -> Self {
253 let (response_tx, response_rx) = mpsc::unbounded_channel();
254 let (introspection_tx, introspection_rx) = mpsc::unbounded_channel();
255
256 let mut maintenance_ticker = time::interval(Duration::from_secs(1));
257 maintenance_ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);
258
259 let instance_workload_classes = Arc::new(Mutex::new(BTreeMap::<
260 ComputeInstanceId,
261 Option<String>,
262 >::new()));
263
264 metrics_registry.register_postprocessor({
268 let instance_workload_classes = Arc::clone(&instance_workload_classes);
269 move |metrics| {
270 let instance_workload_classes = instance_workload_classes
271 .lock()
272 .expect("lock poisoned")
273 .iter()
274 .map(|(id, workload_class)| (id.to_string(), workload_class.clone()))
275 .collect::<BTreeMap<String, Option<String>>>();
276 for metric in metrics {
277 'metric: for metric in metric.mut_metric() {
278 for label in metric.get_label() {
279 if label.name() == "instance_id" {
280 if let Some(workload_class) = instance_workload_classes
281 .get(label.value())
282 .cloned()
283 .flatten()
284 {
285 let mut label = LabelPair::default();
286 label.set_name("workload_class".into());
287 label.set_value(workload_class.clone());
288
289 let mut labels = metric.take_label();
290 labels.push(label);
291 metric.set_label(labels);
292 }
293 continue 'metric;
294 }
295 }
296 }
297 }
298 }
299 });
300
301 let metrics = ComputeControllerMetrics::new(metrics_registry, controller_metrics);
302
303 Self {
304 instances: BTreeMap::new(),
305 instance_workload_classes,
306 build_info,
307 storage_collections,
308 initialized: false,
309 read_only,
310 config: Default::default(),
311 peek_stash_persist_location,
312 stashed_response: None,
313 metrics,
314 now,
315 wallclock_lag,
316 dyncfg: Arc::new(mz_dyncfgs::all_dyncfgs()),
317 replica_dyncfg_overrides: BTreeMap::new(),
318 response_rx,
319 response_tx,
320 introspection_rx: Some(introspection_rx),
321 introspection_tx,
322 maintenance_ticker,
323 maintenance_scheduled: false,
324 }
325 }
326
327 pub fn start_introspection_sink(&mut self, storage_controller: &dyn StorageController) {
332 if let Some(rx) = self.introspection_rx.take() {
333 spawn_introspection_sink(rx, storage_controller);
334 }
335 }
336
337 pub fn instance_exists(&self, id: ComputeInstanceId) -> bool {
339 self.instances.contains_key(&id)
340 }
341
342 fn instance(&self, id: ComputeInstanceId) -> Result<&InstanceState, InstanceMissing> {
344 self.instances.get(&id).ok_or(InstanceMissing(id))
345 }
346
347 pub fn instance_client(
349 &self,
350 id: ComputeInstanceId,
351 ) -> Result<InstanceClient, InstanceMissing> {
352 self.instance(id).map(|instance| instance.client.clone())
353 }
354
355 fn instance_mut(
357 &mut self,
358 id: ComputeInstanceId,
359 ) -> Result<&mut InstanceState, InstanceMissing> {
360 self.instances.get_mut(&id).ok_or(InstanceMissing(id))
361 }
362
363 pub fn collection_ids(
365 &self,
366 instance_id: ComputeInstanceId,
367 ) -> Result<impl Iterator<Item = GlobalId> + '_, InstanceMissing> {
368 let instance = self.instance(instance_id)?;
369 let ids = instance.collections.keys().copied();
370 Ok(ids)
371 }
372
373 pub fn collection_frontiers(
378 &self,
379 collection_id: GlobalId,
380 instance_id: Option<ComputeInstanceId>,
381 ) -> Result<CollectionFrontiers, CollectionLookupError> {
382 let collection = match instance_id {
383 Some(id) => self.instance(id)?.collection(collection_id)?,
384 None => self
385 .instances
386 .values()
387 .find_map(|i| i.collections.get(&collection_id))
388 .ok_or(CollectionMissing(collection_id))?,
389 };
390
391 Ok(collection.frontiers())
392 }
393
394 pub fn collection_reverse_dependencies(
396 &self,
397 instance_id: ComputeInstanceId,
398 id: GlobalId,
399 ) -> Result<impl Iterator<Item = GlobalId> + '_, InstanceMissing> {
400 let instance = self.instance(instance_id)?;
401 let collections = instance.collections.iter();
402 let ids = collections
403 .filter_map(move |(cid, c)| c.compute_dependencies.contains(&id).then_some(*cid));
404 Ok(ids)
405 }
406
407 pub async fn collection_hydrated(
413 &self,
414 instance_id: ComputeInstanceId,
415 collection_id: GlobalId,
416 ) -> Result<bool, anyhow::Error> {
417 let instance = self.instance(instance_id)?;
418
419 let res = instance
420 .call_sync(move |i| i.collection_hydrated(collection_id))
421 .await?;
422
423 Ok(res)
424 }
425
426 pub fn collections_hydrated_for_replicas(
433 &self,
434 instance_id: ComputeInstanceId,
435 replicas: Vec<ReplicaId>,
436 exclude_collections: BTreeSet<GlobalId>,
437 ) -> Result<oneshot::Receiver<bool>, anyhow::Error> {
438 let instance = self.instance(instance_id)?;
439
440 if !instance.replicas.is_empty()
442 && !replicas.iter().any(|id| instance.replicas.contains(id))
443 {
444 return Err(HydrationCheckBadTarget(replicas).into());
445 }
446
447 let (tx, rx) = oneshot::channel();
448 instance.call(move |i| {
449 let result = i
450 .collections_hydrated_on_replicas(Some(replicas), &exclude_collections)
451 .expect("validated");
452 let _ = tx.send(result);
453 });
454
455 Ok(rx)
456 }
457
458 pub async fn dump(&self) -> Result<serde_json::Value, anyhow::Error> {
462 let Self {
469 instances,
470 instance_workload_classes,
471 build_info: _,
472 storage_collections: _,
473 initialized,
474 read_only,
475 config: _,
476 peek_stash_persist_location: _,
477 stashed_response,
478 metrics: _,
479 now: _,
480 wallclock_lag: _,
481 dyncfg: _,
482 replica_dyncfg_overrides: _,
483 response_rx: _,
484 response_tx: _,
485 introspection_rx: _,
486 introspection_tx: _,
487 maintenance_ticker: _,
488 maintenance_scheduled,
489 } = self;
490
491 let mut instances_dump = BTreeMap::new();
492 for (id, instance) in instances {
493 let dump = instance.dump().await?;
494 instances_dump.insert(id.to_string(), dump);
495 }
496
497 let instance_workload_classes: BTreeMap<_, _> = instance_workload_classes
498 .lock()
499 .expect("lock poisoned")
500 .iter()
501 .map(|(id, wc)| (id.to_string(), format!("{wc:?}")))
502 .collect();
503
504 Ok(serde_json::json!({
505 "instances": instances_dump,
506 "instance_workload_classes": instance_workload_classes,
507 "initialized": initialized,
508 "read_only": read_only,
509 "stashed_response": format!("{stashed_response:?}"),
510 "maintenance_scheduled": maintenance_scheduled,
511 }))
512 }
513}
514
515impl ComputeController {
516 pub fn create_instance(
518 &mut self,
519 id: ComputeInstanceId,
520 arranged_logs: BTreeMap<LogVariant, GlobalId>,
521 workload_class: Option<String>,
522 ) -> Result<(), InstanceExists> {
523 if self.instances.contains_key(&id) {
524 return Err(InstanceExists(id));
525 }
526
527 let mut collections = BTreeMap::new();
528 let mut logs = Vec::with_capacity(arranged_logs.len());
529 for (&log, &id) in &arranged_logs {
530 let collection = Collection::new_log();
531 let shared = collection.shared.clone();
532 collections.insert(id, collection);
533 logs.push((log, id, shared));
534 }
535
536 let client = InstanceClient::spawn(
537 id,
538 self.build_info,
539 Arc::clone(&self.storage_collections),
540 self.peek_stash_persist_location.clone(),
541 logs,
542 self.metrics.for_instance(id),
543 self.now.clone(),
544 self.wallclock_lag.clone(),
545 Arc::clone(&self.dyncfg),
546 self.response_tx.clone(),
547 self.introspection_tx.clone(),
548 self.read_only,
549 );
550
551 let instance = InstanceState::new(client, collections);
552 self.instances.insert(id, instance);
553
554 self.instance_workload_classes
555 .lock()
556 .expect("lock poisoned")
557 .insert(id, workload_class.clone());
558
559 let instance = self.instances.get_mut(&id).expect("instance just added");
560 if self.initialized {
561 instance.call(Instance::initialization_complete);
562 }
563
564 let mut config_params = self.config.clone();
570 config_params.workload_class = Some(workload_class);
571 instance.call(|i| i.update_configuration(config_params));
572
573 Ok(())
574 }
575
576 pub fn update_instance_workload_class(
578 &mut self,
579 id: ComputeInstanceId,
580 workload_class: Option<String>,
581 ) -> Result<(), InstanceMissing> {
582 let _ = self.instance(id)?;
584
585 self.instance_workload_classes
586 .lock()
587 .expect("lock poisoned")
588 .insert(id, workload_class);
589
590 self.update_configuration(Default::default());
592
593 Ok(())
594 }
595
596 pub fn drop_instance(&mut self, id: ComputeInstanceId) {
602 if let Some(instance) = self.instances.remove(&id) {
603 instance.call(|i| i.shutdown());
604 }
605
606 self.instance_workload_classes
607 .lock()
608 .expect("lock poisoned")
609 .remove(&id);
610 }
611
612 pub fn dyncfg(&self) -> &Arc<ConfigSet> {
614 &self.dyncfg
615 }
616
617 pub fn update_configuration(&mut self, config_params: ComputeParameters) {
619 config_params.dyncfg_updates.apply(&self.dyncfg);
621
622 let instance_workload_classes = self
623 .instance_workload_classes
624 .lock()
625 .expect("lock poisoned");
626
627 for (id, instance) in self.instances.iter_mut() {
630 let mut params = config_params.clone();
631 params.workload_class = Some(instance_workload_classes[id].clone());
632 instance.call(|i| i.update_configuration(params));
633 }
634
635 let overflowing_behavior = ORE_OVERFLOWING_BEHAVIOR.get(&self.dyncfg);
636 match overflowing_behavior.parse() {
637 Ok(behavior) => mz_ore::overflowing::set_behavior(behavior),
638 Err(err) => {
639 tracing::error!(
640 err,
641 overflowing_behavior,
642 "Invalid value for ore_overflowing_behavior"
643 );
644 }
645 }
646
647 self.config.update(config_params);
649 }
650
651 pub fn update_replica_dyncfg_overrides(
661 &mut self,
662 mut overrides: BTreeMap<ComputeInstanceId, BTreeMap<ReplicaId, ConfigUpdates>>,
663 ) {
664 self.replica_dyncfg_overrides = overrides
665 .values()
666 .flat_map(|replicas| replicas.iter())
667 .map(|(replica_id, updates)| (*replica_id, updates.clone()))
668 .collect();
669 for (id, instance) in self.instances.iter_mut() {
670 let instance_overrides = overrides.remove(id).unwrap_or_default();
671 instance.call(move |i| i.update_replica_dyncfg_overrides(instance_overrides));
672 }
673 }
674
675 pub fn initialization_complete(&mut self) {
681 self.initialized = true;
682 for instance in self.instances.values_mut() {
683 instance.call(Instance::initialization_complete);
684 }
685 }
686
687 pub async fn ready(&mut self) {
695 if self.stashed_response.is_some() {
696 return;
698 }
699 if self.maintenance_scheduled {
700 return;
702 }
703
704 tokio::select! {
705 resp = self.response_rx.recv() => {
706 let resp = resp.expect("`self.response_tx` not dropped");
707 self.stashed_response = Some(resp);
708 }
709 _ = self.maintenance_ticker.tick() => {
710 self.maintenance_scheduled = true;
711 },
712 }
713 }
714
715 pub fn add_replica_to_instance(
717 &mut self,
718 instance_id: ComputeInstanceId,
719 replica_id: ReplicaId,
720 location: ClusterReplicaLocation,
721 config: ComputeReplicaConfig,
722 ) -> Result<(), ReplicaCreationError> {
723 use ReplicaCreationError::*;
724
725 let instance = self.instance(instance_id)?;
726
727 if instance.replicas.contains(&replica_id) {
729 return Err(ReplicaExists(replica_id));
730 }
731
732 let (enable_logging, interval) = match config.logging.interval {
733 Some(interval) => (true, interval),
734 None => (false, Duration::from_secs(1)),
735 };
736
737 let overrides = self.replica_dyncfg_overrides.get(&replica_id);
745
746 let expiration_offset =
747 COMPUTE_REPLICA_EXPIRATION_OFFSET.get_with_overrides(&self.dyncfg, overrides);
748
749 let arrangement_dictionary_compression = ENABLE_ARRANGEMENT_DICTIONARY_COMPRESSION_ALPHA
756 .get_with_overrides(&self.dyncfg, overrides)
757 && config.arrangement_compression;
758
759 let replica_config = ReplicaConfig {
760 location,
761 logging: LoggingConfig {
762 interval,
763 enable_logging,
764 log_logging: config.logging.log_logging,
765 index_logs: Default::default(),
766 },
767 grpc_client: self.config.grpc_client.clone(),
768 expiration_offset: (!expiration_offset.is_zero()).then_some(expiration_offset),
769 arrangement_dictionary_compression,
770 };
771
772 let instance = self.instance_mut(instance_id).expect("validated");
773 instance.replicas.insert(replica_id);
774
775 instance.call(move |i| {
776 i.add_replica(replica_id, replica_config, None)
777 .expect("validated")
778 });
779
780 Ok(())
781 }
782
783 pub fn drop_replica(
785 &mut self,
786 instance_id: ComputeInstanceId,
787 replica_id: ReplicaId,
788 ) -> Result<(), ReplicaDropError> {
789 use ReplicaDropError::*;
790
791 let instance = self.instance_mut(instance_id)?;
792
793 if !instance.replicas.contains(&replica_id) {
795 return Err(ReplicaMissing(replica_id));
796 }
797
798 instance.replicas.remove(&replica_id);
799
800 self.replica_dyncfg_overrides.remove(&replica_id);
804
805 let instance = self.instance_mut(instance_id).expect("validated");
806 instance.call(move |i| i.remove_replica(replica_id).expect("validated"));
807
808 Ok(())
809 }
810
811 pub fn create_dataflow(
821 &mut self,
822 instance_id: ComputeInstanceId,
823 mut dataflow: DataflowDescription<mz_compute_types::plan::LirRelationExpr, ()>,
824 target_replica: Option<ReplicaId>,
825 ) -> Result<(), DataflowCreationError> {
826 use DataflowCreationError::*;
827
828 let instance = self.instance(instance_id)?;
829
830 if let Some(replica_id) = target_replica {
832 if !instance.replicas.contains(&replica_id) {
833 return Err(ReplicaMissing(replica_id));
834 }
835 assert!(
836 dataflow.exported_index_ids().next().is_none(),
837 "Replica-targeted indexes are not supported"
838 );
839 }
840
841 let as_of = dataflow.as_of.as_ref().ok_or(MissingAsOf)?;
843 if as_of.is_empty() && dataflow.subscribe_ids().next().is_some() {
844 return Err(EmptyAsOfForSubscribe);
845 }
846 if as_of.is_empty() && dataflow.copy_to_ids().next().is_some() {
847 return Err(EmptyAsOfForCopyTo);
848 }
849
850 soft_assert_or_log!(
857 !dataflow.index_exports.is_empty() || !dataflow.sink_exports.is_empty(),
858 "dataflow {} has no exports",
859 dataflow.debug_name,
860 );
861
862 let used_imports = dataflow.used_import_ids();
874
875 soft_assert_or_log!(
882 dataflow.import_ids().all(|id| used_imports.contains(&id)),
883 "dataflow {} imports collections no export reads: imports {:?}, read {:?}",
884 dataflow.debug_name,
885 dataflow.import_ids().collect::<Vec<_>>(),
886 used_imports,
887 );
888
889 let storage_ids = dataflow.imported_source_ids().collect();
891 let mut import_read_holds = self.storage_collections.acquire_read_holds(storage_ids)?;
892 for id in dataflow.imported_index_ids() {
893 let read_hold = instance.acquire_read_hold(id)?;
894 import_read_holds.push(read_hold);
895 }
896 for hold in &import_read_holds {
897 if PartialOrder::less_than(as_of, hold.since()) {
898 return Err(SinceViolation(hold.id()));
899 }
900 }
901
902 for id in dataflow.persist_sink_ids() {
904 if self.storage_collections.check_exists(id).is_err() {
905 return Err(CollectionMissing(id));
906 }
907 }
908 let time_dependence = self
909 .determine_time_dependence(instance_id, &dataflow, &used_imports)
910 .expect("must exist");
911
912 let instance = self.instance_mut(instance_id).expect("validated");
913
914 let mut shared_collection_state = BTreeMap::new();
915 for id in dataflow.export_ids() {
916 let shared = SharedCollectionState::new(as_of.clone());
917 let collection = Collection {
918 write_only: dataflow.sink_exports.contains_key(&id),
919 compute_dependencies: dataflow.imported_index_ids().collect(),
920 shared: shared.clone(),
921 time_dependence: time_dependence.clone(),
922 };
923 instance.collections.insert(id, collection);
924 shared_collection_state.insert(id, shared);
925 }
926
927 dataflow.time_dependence = time_dependence;
928
929 instance.call(move |i| {
930 i.create_dataflow(
931 dataflow,
932 import_read_holds,
933 shared_collection_state,
934 target_replica,
935 )
936 .expect("validated")
937 });
938
939 Ok(())
940 }
941
942 pub fn drop_collections(
945 &mut self,
946 instance_id: ComputeInstanceId,
947 collection_ids: Vec<GlobalId>,
948 ) -> Result<(), CollectionUpdateError> {
949 let instance = self.instance_mut(instance_id)?;
950
951 for id in &collection_ids {
953 instance.collection(*id)?;
954 }
955
956 for id in &collection_ids {
957 instance.collections.remove(id);
958 }
959
960 instance.call(|i| i.drop_collections(collection_ids).expect("validated"));
961
962 Ok(())
963 }
964
965 pub fn peek(
972 &self,
973 instance_id: ComputeInstanceId,
974 peek_target: PeekTarget,
975 literal_constraints: Option<Vec<Row>>,
976 uuid: Uuid,
977 timestamp: Timestamp,
978 result_desc: RelationDesc,
979 finishing: RowSetFinishing,
980 map_filter_project: mz_expr::SafeMfpPlan,
981 read_hold: ReadHold,
982 target_replica: Option<ReplicaId>,
983 peek_response_tx: oneshot::Sender<PeekResponse>,
984 ) -> Result<(), PeekError> {
985 use PeekError::*;
986
987 let instance = self.instance(instance_id)?;
988
989 if let Some(replica_id) = target_replica {
991 if !instance.replicas.contains(&replica_id) {
992 return Err(ReplicaMissing(replica_id));
993 }
994 }
995
996 if read_hold.id() != peek_target.id() {
999 return Err(ReadHoldIdMismatch(read_hold.id()));
1000 }
1001 if !read_hold.since().less_equal(×tamp) {
1002 return Err(SinceViolation(peek_target.id()));
1003 }
1004
1005 instance.call(move |i| {
1006 i.peek(
1007 peek_target,
1008 literal_constraints,
1009 uuid,
1010 timestamp,
1011 result_desc,
1012 finishing,
1013 map_filter_project,
1014 read_hold,
1015 target_replica,
1016 peek_response_tx,
1017 )
1018 .expect("validated")
1019 });
1020
1021 Ok(())
1022 }
1023
1024 pub fn cancel_peek(
1034 &self,
1035 instance_id: ComputeInstanceId,
1036 uuid: Uuid,
1037 reason: PeekResponse,
1038 ) -> Result<(), InstanceMissing> {
1039 self.instance(instance_id)?
1040 .call(move |i| i.cancel_peek(uuid, reason));
1041 Ok(())
1042 }
1043
1044 pub fn set_read_policy(
1056 &self,
1057 instance_id: ComputeInstanceId,
1058 policies: Vec<(GlobalId, ReadPolicy)>,
1059 ) -> Result<(), ReadPolicyError> {
1060 use ReadPolicyError::*;
1061
1062 let instance = self.instance(instance_id)?;
1063
1064 for (id, _) in &policies {
1066 let collection = instance.collection(*id)?;
1067 if collection.write_only {
1068 return Err(WriteOnlyCollection(*id));
1069 }
1070 }
1071
1072 self.instance(instance_id)?
1073 .call(|i| i.set_read_policy(policies).expect("validated"));
1074
1075 Ok(())
1076 }
1077
1078 pub fn acquire_read_hold(
1080 &self,
1081 instance_id: ComputeInstanceId,
1082 collection_id: GlobalId,
1083 ) -> Result<ReadHold, CollectionUpdateError> {
1084 let read_hold = self
1085 .instance(instance_id)?
1086 .acquire_read_hold(collection_id)?;
1087 Ok(read_hold)
1088 }
1089
1090 fn determine_time_dependence(
1104 &self,
1105 instance_id: ComputeInstanceId,
1106 dataflow: &DataflowDescription<mz_compute_types::plan::LirRelationExpr, ()>,
1107 used_imports: &BTreeSet<GlobalId>,
1108 ) -> Result<Option<TimeDependence>, TimeDependenceError> {
1109 let instance = self
1110 .instance(instance_id)
1111 .map_err(|err| TimeDependenceError::InstanceMissing(err.0))?;
1112 let mut time_dependencies = Vec::new();
1113
1114 for id in dataflow
1115 .imported_index_ids()
1116 .filter(|id| used_imports.contains(id))
1117 {
1118 let dependence = instance
1119 .get_time_dependence(id)
1120 .map_err(|err| TimeDependenceError::CollectionMissing(err.0))?;
1121 time_dependencies.push(dependence);
1122 }
1123
1124 'source: for id in dataflow
1125 .imported_source_ids()
1126 .filter(|id| used_imports.contains(id))
1127 {
1128 for instance in self.instances.values() {
1131 if let Ok(dependence) = instance.get_time_dependence(id) {
1132 time_dependencies.push(dependence);
1133 continue 'source;
1134 }
1135 }
1136
1137 time_dependencies.push(self.storage_collections.determine_time_dependence(id)?);
1139 }
1140
1141 Ok(TimeDependence::merge(
1142 time_dependencies,
1143 dataflow.refresh_schedule.as_ref(),
1144 ))
1145 }
1146
1147 #[mz_ore::instrument(level = "debug")]
1149 pub fn process(&mut self) -> Option<ComputeControllerResponse> {
1150 if self.maintenance_scheduled {
1152 self.maintain();
1153 self.maintenance_scheduled = false;
1154 }
1155
1156 self.stashed_response.take()
1158 }
1159
1160 #[mz_ore::instrument(level = "debug")]
1161 fn maintain(&mut self) {
1162 for instance in self.instances.values_mut() {
1164 instance.call(Instance::maintain);
1165 }
1166 }
1167
1168 pub fn allow_writes(
1172 &mut self,
1173 instance_id: ComputeInstanceId,
1174 collection_id: GlobalId,
1175 ) -> Result<(), CollectionUpdateError> {
1176 if self.read_only {
1177 tracing::debug!("Skipping allow_writes in read-only mode");
1178 return Ok(());
1179 }
1180
1181 self.allow_writes_inner(instance_id, collection_id)
1182 }
1183
1184 pub fn allow_writes_in_read_only(
1208 &mut self,
1209 instance_id: ComputeInstanceId,
1210 collection_id: GlobalId,
1211 ) -> Result<(), CollectionUpdateError> {
1212 if self.read_only && !collection_id.is_system() {
1221 soft_panic_or_log!(
1222 "allow_writes_in_read_only called for non-system collection {collection_id}; \
1223 falling back to read-only no-op"
1224 );
1225 return Ok(());
1226 }
1227
1228 self.allow_writes_inner(instance_id, collection_id)
1229 }
1230
1231 fn allow_writes_inner(
1232 &mut self,
1233 instance_id: ComputeInstanceId,
1234 collection_id: GlobalId,
1235 ) -> Result<(), CollectionUpdateError> {
1236 let instance = self.instance_mut(instance_id)?;
1237
1238 instance.collection(collection_id)?;
1240
1241 instance.call(move |i| i.allow_writes(collection_id).expect("validated"));
1242
1243 Ok(())
1244 }
1245}
1246
1247#[derive(Debug)]
1248struct InstanceState {
1249 client: InstanceClient,
1250 replicas: BTreeSet<ReplicaId>,
1251 collections: BTreeMap<GlobalId, Collection>,
1252}
1253
1254impl InstanceState {
1255 fn new(client: InstanceClient, collections: BTreeMap<GlobalId, Collection>) -> Self {
1256 Self {
1257 client,
1258 replicas: Default::default(),
1259 collections,
1260 }
1261 }
1262
1263 fn collection(&self, id: GlobalId) -> Result<&Collection, CollectionMissing> {
1264 self.collections.get(&id).ok_or(CollectionMissing(id))
1265 }
1266
1267 fn call<F>(&self, f: F)
1273 where
1274 F: FnOnce(&mut Instance) + Send + 'static,
1275 {
1276 self.client.call(f).expect("instance not dropped")
1277 }
1278
1279 async fn call_sync<F, R>(&self, f: F) -> R
1285 where
1286 F: FnOnce(&mut Instance) -> R + Send + 'static,
1287 R: Send + 'static,
1288 {
1289 self.client
1290 .call_sync(f)
1291 .await
1292 .expect("instance not dropped")
1293 }
1294
1295 pub fn acquire_read_hold(&self, id: GlobalId) -> Result<ReadHold, CollectionMissing> {
1297 let collection = self.collection(id)?;
1307 let since = collection.shared.lock_read_capabilities(|caps| {
1308 let since = caps.frontier().to_owned();
1309 caps.update_iter(since.iter().map(|t| (t.clone(), 1)));
1310 since
1311 });
1312
1313 let hold = ReadHold::new(id, since, self.client.read_hold_tx());
1314 Ok(hold)
1315 }
1316
1317 fn get_time_dependence(
1319 &self,
1320 id: GlobalId,
1321 ) -> Result<Option<TimeDependence>, CollectionMissing> {
1322 Ok(self.collection(id)?.time_dependence.clone())
1323 }
1324
1325 pub async fn dump(&self) -> Result<serde_json::Value, anyhow::Error> {
1327 let Self {
1329 client: _,
1330 replicas,
1331 collections,
1332 } = self;
1333
1334 let instance = self.call_sync(|i| i.dump()).await?;
1335 let replicas: Vec<_> = replicas.iter().map(|id| id.to_string()).collect();
1336 let collections: BTreeMap<_, _> = collections
1337 .iter()
1338 .map(|(id, c)| (id.to_string(), format!("{c:?}")))
1339 .collect();
1340
1341 Ok(serde_json::json!({
1342 "instance": instance,
1343 "replicas": replicas,
1344 "collections": collections,
1345 }))
1346 }
1347}
1348
1349#[derive(Debug)]
1350struct Collection {
1351 write_only: bool,
1353 compute_dependencies: BTreeSet<GlobalId>,
1354 shared: SharedCollectionState,
1355 time_dependence: Option<TimeDependence>,
1358}
1359
1360impl Collection {
1361 fn new_log() -> Self {
1362 let as_of = Antichain::from_elem(Timestamp::MIN);
1363 Self {
1364 write_only: false,
1365 compute_dependencies: Default::default(),
1366 shared: SharedCollectionState::new(as_of),
1367 time_dependence: Some(TimeDependence::default()),
1368 }
1369 }
1370
1371 fn frontiers(&self) -> CollectionFrontiers {
1372 let read_frontier = self
1373 .shared
1374 .lock_read_capabilities(|c| c.frontier().to_owned());
1375 let write_frontier = self.shared.lock_write_frontier(|f| f.clone());
1376 CollectionFrontiers {
1377 read_frontier,
1378 write_frontier,
1379 }
1380 }
1381}
1382
1383#[derive(Clone, Debug)]
1385pub struct CollectionFrontiers {
1386 pub read_frontier: Antichain<Timestamp>,
1388 pub write_frontier: Antichain<Timestamp>,
1390}
1391
1392impl Default for CollectionFrontiers {
1393 fn default() -> Self {
1394 Self {
1395 read_frontier: Antichain::from_elem(Timestamp::MIN),
1396 write_frontier: Antichain::from_elem(Timestamp::MIN),
1397 }
1398 }
1399}