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::tracing::OpenTelemetryContext;
57use mz_persist_types::PersistLocation;
58use mz_repr::{GlobalId, RelationDesc, Row, Timestamp};
59use mz_storage_client::controller::StorageController;
60use mz_storage_types::dyncfgs::ORE_OVERFLOWING_BEHAVIOR;
61use mz_storage_types::read_holds::ReadHold;
62use mz_storage_types::read_policy::ReadPolicy;
63use mz_storage_types::time_dependence::{TimeDependence, TimeDependenceError};
64use prometheus::proto::LabelPair;
65use serde::{Deserialize, Serialize};
66use timely::PartialOrder;
67use timely::progress::Antichain;
68use tokio::sync::{mpsc, oneshot};
69use tokio::time::{self, MissedTickBehavior};
70use uuid::Uuid;
71
72use crate::controller::error::{
73 CollectionLookupError, CollectionMissing, CollectionUpdateError, DataflowCreationError,
74 HydrationCheckBadTarget, InstanceExists, InstanceMissing, PeekError, ReadPolicyError,
75 ReplicaCreationError, ReplicaDropError,
76};
77use crate::controller::instance::{Instance, SharedCollectionState};
78use crate::controller::introspection::{IntrospectionUpdates, spawn_introspection_sink};
79use crate::controller::replica::ReplicaConfig;
80use crate::logging::{LogVariant, LoggingConfig};
81use crate::metrics::ComputeControllerMetrics;
82use crate::protocol::command::{ComputeParameters, PeekTarget};
83use crate::protocol::response::{PeekResponse, SubscribeBatch};
84
85mod instance;
86mod introspection;
87mod replica;
88mod sequential_hydration;
89
90pub mod error;
91pub mod instance_client;
92pub use instance_client::InstanceClient;
93
94pub(crate) type StorageCollections =
95 Arc<dyn mz_storage_client::storage_collections::StorageCollections + Send + Sync>;
96
97#[derive(Debug)]
99pub enum ComputeControllerResponse {
100 PeekNotification(Uuid, PeekNotification, OpenTelemetryContext),
102 SubscribeResponse(GlobalId, SubscribeBatch),
104 CopyToResponse(GlobalId, Result<u64, anyhow::Error>),
115 FrontierUpper {
120 id: GlobalId,
122 upper: Antichain<Timestamp>,
124 },
125}
126
127#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
129pub enum PeekNotification {
130 Success {
132 rows: u64,
134 result_size: u64,
136 },
137 Error(String),
139 Canceled,
141}
142
143impl PeekNotification {
144 fn new(peek_response: &PeekResponse, offset: usize, limit: Option<usize>) -> Self {
147 match peek_response {
148 PeekResponse::Rows(rows) => {
149 let num_rows = u64::cast_from(RowCollection::offset_limit(
150 rows.iter().map(|r| r.count()).sum(),
151 offset,
152 limit,
153 ));
154 let result_size = u64::cast_from(rows.iter().map(|r| r.byte_len()).sum::<usize>());
155
156 tracing::trace!(?num_rows, ?result_size, "inline peek result");
157
158 Self::Success {
159 rows: num_rows,
160 result_size,
161 }
162 }
163 PeekResponse::Stashed(stashed_response) => {
164 let rows = stashed_response.num_rows(offset, limit);
165 let result_size = stashed_response.size_bytes();
166
167 tracing::trace!(?rows, ?result_size, "stashed peek result");
168
169 Self::Success {
170 rows: u64::cast_from(rows),
171 result_size: u64::cast_from(result_size),
172 }
173 }
174 PeekResponse::Error(err) => Self::Error(err.clone()),
175 PeekResponse::Canceled => Self::Canceled,
176 }
177 }
178}
179
180pub struct ComputeController {
182 instances: BTreeMap<ComputeInstanceId, InstanceState>,
183 instance_workload_classes: Arc<Mutex<BTreeMap<ComputeInstanceId, Option<String>>>>,
187 build_info: &'static BuildInfo,
188 storage_collections: StorageCollections,
190 initialized: bool,
192 read_only: bool,
198 config: ComputeParameters,
200 peek_stash_persist_location: PersistLocation,
202 stashed_response: Option<ComputeControllerResponse>,
204 metrics: ComputeControllerMetrics,
206 now: NowFn,
208 wallclock_lag: WallclockLagFn<Timestamp>,
210 dyncfg: Arc<ConfigSet>,
215 replica_dyncfg_overrides: BTreeMap<ReplicaId, ConfigUpdates>,
221
222 response_rx: mpsc::UnboundedReceiver<ComputeControllerResponse>,
224 response_tx: mpsc::UnboundedSender<ComputeControllerResponse>,
226 introspection_rx: Option<mpsc::UnboundedReceiver<IntrospectionUpdates>>,
231 introspection_tx: mpsc::UnboundedSender<IntrospectionUpdates>,
233
234 maintenance_ticker: tokio::time::Interval,
236 maintenance_scheduled: bool,
238}
239
240impl ComputeController {
241 pub fn new(
243 build_info: &'static BuildInfo,
244 storage_collections: StorageCollections,
245 read_only: bool,
246 metrics_registry: &MetricsRegistry,
247 peek_stash_persist_location: PersistLocation,
248 controller_metrics: ControllerMetrics,
249 now: NowFn,
250 wallclock_lag: WallclockLagFn<Timestamp>,
251 ) -> Self {
252 let (response_tx, response_rx) = mpsc::unbounded_channel();
253 let (introspection_tx, introspection_rx) = mpsc::unbounded_channel();
254
255 let mut maintenance_ticker = time::interval(Duration::from_secs(1));
256 maintenance_ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);
257
258 let instance_workload_classes = Arc::new(Mutex::new(BTreeMap::<
259 ComputeInstanceId,
260 Option<String>,
261 >::new()));
262
263 metrics_registry.register_postprocessor({
267 let instance_workload_classes = Arc::clone(&instance_workload_classes);
268 move |metrics| {
269 let instance_workload_classes = instance_workload_classes
270 .lock()
271 .expect("lock poisoned")
272 .iter()
273 .map(|(id, workload_class)| (id.to_string(), workload_class.clone()))
274 .collect::<BTreeMap<String, Option<String>>>();
275 for metric in metrics {
276 'metric: for metric in metric.mut_metric() {
277 for label in metric.get_label() {
278 if label.name() == "instance_id" {
279 if let Some(workload_class) = instance_workload_classes
280 .get(label.value())
281 .cloned()
282 .flatten()
283 {
284 let mut label = LabelPair::default();
285 label.set_name("workload_class".into());
286 label.set_value(workload_class.clone());
287
288 let mut labels = metric.take_label();
289 labels.push(label);
290 metric.set_label(labels);
291 }
292 continue 'metric;
293 }
294 }
295 }
296 }
297 }
298 });
299
300 let metrics = ComputeControllerMetrics::new(metrics_registry, controller_metrics);
301
302 Self {
303 instances: BTreeMap::new(),
304 instance_workload_classes,
305 build_info,
306 storage_collections,
307 initialized: false,
308 read_only,
309 config: Default::default(),
310 peek_stash_persist_location,
311 stashed_response: None,
312 metrics,
313 now,
314 wallclock_lag,
315 dyncfg: Arc::new(mz_dyncfgs::all_dyncfgs()),
316 replica_dyncfg_overrides: BTreeMap::new(),
317 response_rx,
318 response_tx,
319 introspection_rx: Some(introspection_rx),
320 introspection_tx,
321 maintenance_ticker,
322 maintenance_scheduled: false,
323 }
324 }
325
326 pub fn start_introspection_sink(&mut self, storage_controller: &dyn StorageController) {
331 if let Some(rx) = self.introspection_rx.take() {
332 spawn_introspection_sink(rx, storage_controller);
333 }
334 }
335
336 pub fn instance_exists(&self, id: ComputeInstanceId) -> bool {
338 self.instances.contains_key(&id)
339 }
340
341 fn instance(&self, id: ComputeInstanceId) -> Result<&InstanceState, InstanceMissing> {
343 self.instances.get(&id).ok_or(InstanceMissing(id))
344 }
345
346 pub fn instance_client(
348 &self,
349 id: ComputeInstanceId,
350 ) -> Result<InstanceClient, InstanceMissing> {
351 self.instance(id).map(|instance| instance.client.clone())
352 }
353
354 fn instance_mut(
356 &mut self,
357 id: ComputeInstanceId,
358 ) -> Result<&mut InstanceState, InstanceMissing> {
359 self.instances.get_mut(&id).ok_or(InstanceMissing(id))
360 }
361
362 pub fn collection_ids(
364 &self,
365 instance_id: ComputeInstanceId,
366 ) -> Result<impl Iterator<Item = GlobalId> + '_, InstanceMissing> {
367 let instance = self.instance(instance_id)?;
368 let ids = instance.collections.keys().copied();
369 Ok(ids)
370 }
371
372 pub fn collection_frontiers(
377 &self,
378 collection_id: GlobalId,
379 instance_id: Option<ComputeInstanceId>,
380 ) -> Result<CollectionFrontiers, CollectionLookupError> {
381 let collection = match instance_id {
382 Some(id) => self.instance(id)?.collection(collection_id)?,
383 None => self
384 .instances
385 .values()
386 .find_map(|i| i.collections.get(&collection_id))
387 .ok_or(CollectionMissing(collection_id))?,
388 };
389
390 Ok(collection.frontiers())
391 }
392
393 pub fn collection_reverse_dependencies(
395 &self,
396 instance_id: ComputeInstanceId,
397 id: GlobalId,
398 ) -> Result<impl Iterator<Item = GlobalId> + '_, InstanceMissing> {
399 let instance = self.instance(instance_id)?;
400 let collections = instance.collections.iter();
401 let ids = collections
402 .filter_map(move |(cid, c)| c.compute_dependencies.contains(&id).then_some(*cid));
403 Ok(ids)
404 }
405
406 pub async fn collection_hydrated(
412 &self,
413 instance_id: ComputeInstanceId,
414 collection_id: GlobalId,
415 ) -> Result<bool, anyhow::Error> {
416 let instance = self.instance(instance_id)?;
417
418 let res = instance
419 .call_sync(move |i| i.collection_hydrated(collection_id))
420 .await?;
421
422 Ok(res)
423 }
424
425 pub fn collections_hydrated_for_replicas(
432 &self,
433 instance_id: ComputeInstanceId,
434 replicas: Vec<ReplicaId>,
435 exclude_collections: BTreeSet<GlobalId>,
436 ) -> Result<oneshot::Receiver<bool>, anyhow::Error> {
437 let instance = self.instance(instance_id)?;
438
439 if !instance.replicas.is_empty()
441 && !replicas.iter().any(|id| instance.replicas.contains(id))
442 {
443 return Err(HydrationCheckBadTarget(replicas).into());
444 }
445
446 let (tx, rx) = oneshot::channel();
447 instance.call(move |i| {
448 let result = i
449 .collections_hydrated_on_replicas(Some(replicas), &exclude_collections)
450 .expect("validated");
451 let _ = tx.send(result);
452 });
453
454 Ok(rx)
455 }
456
457 pub async fn dump(&self) -> Result<serde_json::Value, anyhow::Error> {
461 let Self {
468 instances,
469 instance_workload_classes,
470 build_info: _,
471 storage_collections: _,
472 initialized,
473 read_only,
474 config: _,
475 peek_stash_persist_location: _,
476 stashed_response,
477 metrics: _,
478 now: _,
479 wallclock_lag: _,
480 dyncfg: _,
481 replica_dyncfg_overrides: _,
482 response_rx: _,
483 response_tx: _,
484 introspection_rx: _,
485 introspection_tx: _,
486 maintenance_ticker: _,
487 maintenance_scheduled,
488 } = self;
489
490 let mut instances_dump = BTreeMap::new();
491 for (id, instance) in instances {
492 let dump = instance.dump().await?;
493 instances_dump.insert(id.to_string(), dump);
494 }
495
496 let instance_workload_classes: BTreeMap<_, _> = instance_workload_classes
497 .lock()
498 .expect("lock poisoned")
499 .iter()
500 .map(|(id, wc)| (id.to_string(), format!("{wc:?}")))
501 .collect();
502
503 Ok(serde_json::json!({
504 "instances": instances_dump,
505 "instance_workload_classes": instance_workload_classes,
506 "initialized": initialized,
507 "read_only": read_only,
508 "stashed_response": format!("{stashed_response:?}"),
509 "maintenance_scheduled": maintenance_scheduled,
510 }))
511 }
512}
513
514impl ComputeController {
515 pub fn create_instance(
517 &mut self,
518 id: ComputeInstanceId,
519 arranged_logs: BTreeMap<LogVariant, GlobalId>,
520 workload_class: Option<String>,
521 ) -> Result<(), InstanceExists> {
522 if self.instances.contains_key(&id) {
523 return Err(InstanceExists(id));
524 }
525
526 let mut collections = BTreeMap::new();
527 let mut logs = Vec::with_capacity(arranged_logs.len());
528 for (&log, &id) in &arranged_logs {
529 let collection = Collection::new_log();
530 let shared = collection.shared.clone();
531 collections.insert(id, collection);
532 logs.push((log, id, shared));
533 }
534
535 let client = InstanceClient::spawn(
536 id,
537 self.build_info,
538 Arc::clone(&self.storage_collections),
539 self.peek_stash_persist_location.clone(),
540 logs,
541 self.metrics.for_instance(id),
542 self.now.clone(),
543 self.wallclock_lag.clone(),
544 Arc::clone(&self.dyncfg),
545 self.response_tx.clone(),
546 self.introspection_tx.clone(),
547 self.read_only,
548 );
549
550 let instance = InstanceState::new(client, collections);
551 self.instances.insert(id, instance);
552
553 self.instance_workload_classes
554 .lock()
555 .expect("lock poisoned")
556 .insert(id, workload_class.clone());
557
558 let instance = self.instances.get_mut(&id).expect("instance just added");
559 if self.initialized {
560 instance.call(Instance::initialization_complete);
561 }
562
563 let mut config_params = self.config.clone();
569 config_params.workload_class = Some(workload_class);
570 instance.call(|i| i.update_configuration(config_params));
571
572 Ok(())
573 }
574
575 pub fn update_instance_workload_class(
577 &mut self,
578 id: ComputeInstanceId,
579 workload_class: Option<String>,
580 ) -> Result<(), InstanceMissing> {
581 let _ = self.instance(id)?;
583
584 self.instance_workload_classes
585 .lock()
586 .expect("lock poisoned")
587 .insert(id, workload_class);
588
589 self.update_configuration(Default::default());
591
592 Ok(())
593 }
594
595 pub fn drop_instance(&mut self, id: ComputeInstanceId) {
601 if let Some(instance) = self.instances.remove(&id) {
602 instance.call(|i| i.shutdown());
603 }
604
605 self.instance_workload_classes
606 .lock()
607 .expect("lock poisoned")
608 .remove(&id);
609 }
610
611 pub fn dyncfg(&self) -> &Arc<ConfigSet> {
613 &self.dyncfg
614 }
615
616 pub fn update_configuration(&mut self, config_params: ComputeParameters) {
618 config_params.dyncfg_updates.apply(&self.dyncfg);
620
621 let instance_workload_classes = self
622 .instance_workload_classes
623 .lock()
624 .expect("lock poisoned");
625
626 for (id, instance) in self.instances.iter_mut() {
629 let mut params = config_params.clone();
630 params.workload_class = Some(instance_workload_classes[id].clone());
631 instance.call(|i| i.update_configuration(params));
632 }
633
634 let overflowing_behavior = ORE_OVERFLOWING_BEHAVIOR.get(&self.dyncfg);
635 match overflowing_behavior.parse() {
636 Ok(behavior) => mz_ore::overflowing::set_behavior(behavior),
637 Err(err) => {
638 tracing::error!(
639 err,
640 overflowing_behavior,
641 "Invalid value for ore_overflowing_behavior"
642 );
643 }
644 }
645
646 self.config.update(config_params);
648 }
649
650 pub fn update_replica_dyncfg_overrides(
660 &mut self,
661 mut overrides: BTreeMap<ComputeInstanceId, BTreeMap<ReplicaId, ConfigUpdates>>,
662 ) {
663 self.replica_dyncfg_overrides = overrides
664 .values()
665 .flat_map(|replicas| replicas.iter())
666 .map(|(replica_id, updates)| (*replica_id, updates.clone()))
667 .collect();
668 for (id, instance) in self.instances.iter_mut() {
669 let instance_overrides = overrides.remove(id).unwrap_or_default();
670 instance.call(move |i| i.update_replica_dyncfg_overrides(instance_overrides));
671 }
672 }
673
674 pub fn initialization_complete(&mut self) {
680 self.initialized = true;
681 for instance in self.instances.values_mut() {
682 instance.call(Instance::initialization_complete);
683 }
684 }
685
686 pub async fn ready(&mut self) {
694 if self.stashed_response.is_some() {
695 return;
697 }
698 if self.maintenance_scheduled {
699 return;
701 }
702
703 tokio::select! {
704 resp = self.response_rx.recv() => {
705 let resp = resp.expect("`self.response_tx` not dropped");
706 self.stashed_response = Some(resp);
707 }
708 _ = self.maintenance_ticker.tick() => {
709 self.maintenance_scheduled = true;
710 },
711 }
712 }
713
714 pub fn add_replica_to_instance(
716 &mut self,
717 instance_id: ComputeInstanceId,
718 replica_id: ReplicaId,
719 location: ClusterReplicaLocation,
720 config: ComputeReplicaConfig,
721 ) -> Result<(), ReplicaCreationError> {
722 use ReplicaCreationError::*;
723
724 let instance = self.instance(instance_id)?;
725
726 if instance.replicas.contains(&replica_id) {
728 return Err(ReplicaExists(replica_id));
729 }
730
731 let (enable_logging, interval) = match config.logging.interval {
732 Some(interval) => (true, interval),
733 None => (false, Duration::from_secs(1)),
734 };
735
736 let overrides = self.replica_dyncfg_overrides.get(&replica_id);
744
745 let expiration_offset =
746 COMPUTE_REPLICA_EXPIRATION_OFFSET.get_with_overrides(&self.dyncfg, overrides);
747
748 let arrangement_dictionary_compression = ENABLE_ARRANGEMENT_DICTIONARY_COMPRESSION_ALPHA
755 .get_with_overrides(&self.dyncfg, overrides)
756 && config.arrangement_compression;
757
758 let replica_config = ReplicaConfig {
759 location,
760 logging: LoggingConfig {
761 interval,
762 enable_logging,
763 log_logging: config.logging.log_logging,
764 index_logs: Default::default(),
765 },
766 grpc_client: self.config.grpc_client.clone(),
767 expiration_offset: (!expiration_offset.is_zero()).then_some(expiration_offset),
768 arrangement_dictionary_compression,
769 };
770
771 let instance = self.instance_mut(instance_id).expect("validated");
772 instance.replicas.insert(replica_id);
773
774 instance.call(move |i| {
775 i.add_replica(replica_id, replica_config, None)
776 .expect("validated")
777 });
778
779 Ok(())
780 }
781
782 pub fn drop_replica(
784 &mut self,
785 instance_id: ComputeInstanceId,
786 replica_id: ReplicaId,
787 ) -> Result<(), ReplicaDropError> {
788 use ReplicaDropError::*;
789
790 let instance = self.instance_mut(instance_id)?;
791
792 if !instance.replicas.contains(&replica_id) {
794 return Err(ReplicaMissing(replica_id));
795 }
796
797 instance.replicas.remove(&replica_id);
798
799 self.replica_dyncfg_overrides.remove(&replica_id);
803
804 let instance = self.instance_mut(instance_id).expect("validated");
805 instance.call(move |i| i.remove_replica(replica_id).expect("validated"));
806
807 Ok(())
808 }
809
810 pub fn create_dataflow(
820 &mut self,
821 instance_id: ComputeInstanceId,
822 mut dataflow: DataflowDescription<mz_compute_types::plan::LirRelationExpr, ()>,
823 target_replica: Option<ReplicaId>,
824 ) -> Result<(), DataflowCreationError> {
825 use DataflowCreationError::*;
826
827 let instance = self.instance(instance_id)?;
828
829 if let Some(replica_id) = target_replica {
831 if !instance.replicas.contains(&replica_id) {
832 return Err(ReplicaMissing(replica_id));
833 }
834 assert!(
835 dataflow.exported_index_ids().next().is_none(),
836 "Replica-targeted indexes are not supported"
837 );
838 }
839
840 let as_of = dataflow.as_of.as_ref().ok_or(MissingAsOf)?;
842 if as_of.is_empty() && dataflow.subscribe_ids().next().is_some() {
843 return Err(EmptyAsOfForSubscribe);
844 }
845 if as_of.is_empty() && dataflow.copy_to_ids().next().is_some() {
846 return Err(EmptyAsOfForCopyTo);
847 }
848
849 soft_assert_or_log!(
856 !dataflow.index_exports.is_empty() || !dataflow.sink_exports.is_empty(),
857 "dataflow {} has no exports",
858 dataflow.debug_name,
859 );
860
861 let used_imports = dataflow.used_import_ids();
873
874 soft_assert_or_log!(
881 dataflow.import_ids().all(|id| used_imports.contains(&id)),
882 "dataflow {} imports collections no export reads: imports {:?}, read {:?}",
883 dataflow.debug_name,
884 dataflow.import_ids().collect::<Vec<_>>(),
885 used_imports,
886 );
887
888 let storage_ids = dataflow.imported_source_ids().collect();
890 let mut import_read_holds = self.storage_collections.acquire_read_holds(storage_ids)?;
891 for id in dataflow.imported_index_ids() {
892 let read_hold = instance.acquire_read_hold(id)?;
893 import_read_holds.push(read_hold);
894 }
895 for hold in &import_read_holds {
896 if PartialOrder::less_than(as_of, hold.since()) {
897 return Err(SinceViolation(hold.id()));
898 }
899 }
900
901 for id in dataflow.persist_sink_ids() {
903 if self.storage_collections.check_exists(id).is_err() {
904 return Err(CollectionMissing(id));
905 }
906 }
907 let time_dependence = self
908 .determine_time_dependence(instance_id, &dataflow, &used_imports)
909 .expect("must exist");
910
911 let instance = self.instance_mut(instance_id).expect("validated");
912
913 let mut shared_collection_state = BTreeMap::new();
914 for id in dataflow.export_ids() {
915 let shared = SharedCollectionState::new(as_of.clone());
916 let collection = Collection {
917 write_only: dataflow.sink_exports.contains_key(&id),
918 compute_dependencies: dataflow.imported_index_ids().collect(),
919 shared: shared.clone(),
920 time_dependence: time_dependence.clone(),
921 };
922 instance.collections.insert(id, collection);
923 shared_collection_state.insert(id, shared);
924 }
925
926 dataflow.time_dependence = time_dependence;
927
928 instance.call(move |i| {
929 i.create_dataflow(
930 dataflow,
931 import_read_holds,
932 shared_collection_state,
933 target_replica,
934 )
935 .expect("validated")
936 });
937
938 Ok(())
939 }
940
941 pub fn drop_collections(
944 &mut self,
945 instance_id: ComputeInstanceId,
946 collection_ids: Vec<GlobalId>,
947 ) -> Result<(), CollectionUpdateError> {
948 let instance = self.instance_mut(instance_id)?;
949
950 for id in &collection_ids {
952 instance.collection(*id)?;
953 }
954
955 for id in &collection_ids {
956 instance.collections.remove(id);
957 }
958
959 instance.call(|i| i.drop_collections(collection_ids).expect("validated"));
960
961 Ok(())
962 }
963
964 pub fn peek(
971 &self,
972 instance_id: ComputeInstanceId,
973 peek_target: PeekTarget,
974 literal_constraints: Option<Vec<Row>>,
975 uuid: Uuid,
976 timestamp: Timestamp,
977 result_desc: RelationDesc,
978 finishing: RowSetFinishing,
979 map_filter_project: mz_expr::SafeMfpPlan,
980 read_hold: ReadHold,
981 target_replica: Option<ReplicaId>,
982 peek_response_tx: oneshot::Sender<PeekResponse>,
983 ) -> Result<(), PeekError> {
984 use PeekError::*;
985
986 let instance = self.instance(instance_id)?;
987
988 if let Some(replica_id) = target_replica {
990 if !instance.replicas.contains(&replica_id) {
991 return Err(ReplicaMissing(replica_id));
992 }
993 }
994
995 if read_hold.id() != peek_target.id() {
998 return Err(ReadHoldIdMismatch(read_hold.id()));
999 }
1000 if !read_hold.since().less_equal(×tamp) {
1001 return Err(SinceViolation(peek_target.id()));
1002 }
1003
1004 instance.call(move |i| {
1005 i.peek(
1006 peek_target,
1007 literal_constraints,
1008 uuid,
1009 timestamp,
1010 result_desc,
1011 finishing,
1012 map_filter_project,
1013 read_hold,
1014 target_replica,
1015 peek_response_tx,
1016 )
1017 .expect("validated")
1018 });
1019
1020 Ok(())
1021 }
1022
1023 pub fn cancel_peek(
1033 &self,
1034 instance_id: ComputeInstanceId,
1035 uuid: Uuid,
1036 reason: PeekResponse,
1037 ) -> Result<(), InstanceMissing> {
1038 self.instance(instance_id)?
1039 .call(move |i| i.cancel_peek(uuid, reason));
1040 Ok(())
1041 }
1042
1043 pub fn set_read_policy(
1055 &self,
1056 instance_id: ComputeInstanceId,
1057 policies: Vec<(GlobalId, ReadPolicy)>,
1058 ) -> Result<(), ReadPolicyError> {
1059 use ReadPolicyError::*;
1060
1061 let instance = self.instance(instance_id)?;
1062
1063 for (id, _) in &policies {
1065 let collection = instance.collection(*id)?;
1066 if collection.write_only {
1067 return Err(WriteOnlyCollection(*id));
1068 }
1069 }
1070
1071 self.instance(instance_id)?
1072 .call(|i| i.set_read_policy(policies).expect("validated"));
1073
1074 Ok(())
1075 }
1076
1077 pub fn acquire_read_hold(
1079 &self,
1080 instance_id: ComputeInstanceId,
1081 collection_id: GlobalId,
1082 ) -> Result<ReadHold, CollectionUpdateError> {
1083 let read_hold = self
1084 .instance(instance_id)?
1085 .acquire_read_hold(collection_id)?;
1086 Ok(read_hold)
1087 }
1088
1089 fn determine_time_dependence(
1103 &self,
1104 instance_id: ComputeInstanceId,
1105 dataflow: &DataflowDescription<mz_compute_types::plan::LirRelationExpr, ()>,
1106 used_imports: &BTreeSet<GlobalId>,
1107 ) -> Result<Option<TimeDependence>, TimeDependenceError> {
1108 let instance = self
1109 .instance(instance_id)
1110 .map_err(|err| TimeDependenceError::InstanceMissing(err.0))?;
1111 let mut time_dependencies = Vec::new();
1112
1113 for id in dataflow
1114 .imported_index_ids()
1115 .filter(|id| used_imports.contains(id))
1116 {
1117 let dependence = instance
1118 .get_time_dependence(id)
1119 .map_err(|err| TimeDependenceError::CollectionMissing(err.0))?;
1120 time_dependencies.push(dependence);
1121 }
1122
1123 'source: for id in dataflow
1124 .imported_source_ids()
1125 .filter(|id| used_imports.contains(id))
1126 {
1127 for instance in self.instances.values() {
1130 if let Ok(dependence) = instance.get_time_dependence(id) {
1131 time_dependencies.push(dependence);
1132 continue 'source;
1133 }
1134 }
1135
1136 time_dependencies.push(self.storage_collections.determine_time_dependence(id)?);
1138 }
1139
1140 Ok(TimeDependence::merge(
1141 time_dependencies,
1142 dataflow.refresh_schedule.as_ref(),
1143 ))
1144 }
1145
1146 #[mz_ore::instrument(level = "debug")]
1148 pub fn process(&mut self) -> Option<ComputeControllerResponse> {
1149 if self.maintenance_scheduled {
1151 self.maintain();
1152 self.maintenance_scheduled = false;
1153 }
1154
1155 self.stashed_response.take()
1157 }
1158
1159 #[mz_ore::instrument(level = "debug")]
1160 fn maintain(&mut self) {
1161 for instance in self.instances.values_mut() {
1163 instance.call(Instance::maintain);
1164 }
1165 }
1166
1167 pub fn allow_writes(
1171 &mut self,
1172 instance_id: ComputeInstanceId,
1173 collection_id: GlobalId,
1174 ) -> Result<(), CollectionUpdateError> {
1175 if self.read_only {
1176 tracing::debug!("Skipping allow_writes in read-only mode");
1177 return Ok(());
1178 }
1179
1180 let instance = self.instance_mut(instance_id)?;
1181
1182 instance.collection(collection_id)?;
1184
1185 instance.call(move |i| i.allow_writes(collection_id).expect("validated"));
1186
1187 Ok(())
1188 }
1189}
1190
1191#[derive(Debug)]
1192struct InstanceState {
1193 client: InstanceClient,
1194 replicas: BTreeSet<ReplicaId>,
1195 collections: BTreeMap<GlobalId, Collection>,
1196}
1197
1198impl InstanceState {
1199 fn new(client: InstanceClient, collections: BTreeMap<GlobalId, Collection>) -> Self {
1200 Self {
1201 client,
1202 replicas: Default::default(),
1203 collections,
1204 }
1205 }
1206
1207 fn collection(&self, id: GlobalId) -> Result<&Collection, CollectionMissing> {
1208 self.collections.get(&id).ok_or(CollectionMissing(id))
1209 }
1210
1211 fn call<F>(&self, f: F)
1217 where
1218 F: FnOnce(&mut Instance) + Send + 'static,
1219 {
1220 self.client.call(f).expect("instance not dropped")
1221 }
1222
1223 async fn call_sync<F, R>(&self, f: F) -> R
1229 where
1230 F: FnOnce(&mut Instance) -> R + Send + 'static,
1231 R: Send + 'static,
1232 {
1233 self.client
1234 .call_sync(f)
1235 .await
1236 .expect("instance not dropped")
1237 }
1238
1239 pub fn acquire_read_hold(&self, id: GlobalId) -> Result<ReadHold, CollectionMissing> {
1241 let collection = self.collection(id)?;
1251 let since = collection.shared.lock_read_capabilities(|caps| {
1252 let since = caps.frontier().to_owned();
1253 caps.update_iter(since.iter().map(|t| (t.clone(), 1)));
1254 since
1255 });
1256
1257 let hold = ReadHold::new(id, since, self.client.read_hold_tx());
1258 Ok(hold)
1259 }
1260
1261 fn get_time_dependence(
1263 &self,
1264 id: GlobalId,
1265 ) -> Result<Option<TimeDependence>, CollectionMissing> {
1266 Ok(self.collection(id)?.time_dependence.clone())
1267 }
1268
1269 pub async fn dump(&self) -> Result<serde_json::Value, anyhow::Error> {
1271 let Self {
1273 client: _,
1274 replicas,
1275 collections,
1276 } = self;
1277
1278 let instance = self.call_sync(|i| i.dump()).await?;
1279 let replicas: Vec<_> = replicas.iter().map(|id| id.to_string()).collect();
1280 let collections: BTreeMap<_, _> = collections
1281 .iter()
1282 .map(|(id, c)| (id.to_string(), format!("{c:?}")))
1283 .collect();
1284
1285 Ok(serde_json::json!({
1286 "instance": instance,
1287 "replicas": replicas,
1288 "collections": collections,
1289 }))
1290 }
1291}
1292
1293#[derive(Debug)]
1294struct Collection {
1295 write_only: bool,
1297 compute_dependencies: BTreeSet<GlobalId>,
1298 shared: SharedCollectionState,
1299 time_dependence: Option<TimeDependence>,
1302}
1303
1304impl Collection {
1305 fn new_log() -> Self {
1306 let as_of = Antichain::from_elem(Timestamp::MIN);
1307 Self {
1308 write_only: false,
1309 compute_dependencies: Default::default(),
1310 shared: SharedCollectionState::new(as_of),
1311 time_dependence: Some(TimeDependence::default()),
1312 }
1313 }
1314
1315 fn frontiers(&self) -> CollectionFrontiers {
1316 let read_frontier = self
1317 .shared
1318 .lock_read_capabilities(|c| c.frontier().to_owned());
1319 let write_frontier = self.shared.lock_write_frontier(|f| f.clone());
1320 CollectionFrontiers {
1321 read_frontier,
1322 write_frontier,
1323 }
1324 }
1325}
1326
1327#[derive(Clone, Debug)]
1329pub struct CollectionFrontiers {
1330 pub read_frontier: Antichain<Timestamp>,
1332 pub write_frontier: Antichain<Timestamp>,
1334}
1335
1336impl Default for CollectionFrontiers {
1337 fn default() -> Self {
1338 Self {
1339 read_frontier: Antichain::from_elem(Timestamp::MIN),
1340 write_frontier: Antichain::from_elem(Timestamp::MIN),
1341 }
1342 }
1343}