1use std::borrow::Cow;
70use std::collections::{BTreeMap, BTreeSet, VecDeque};
71use std::net::IpAddr;
72use std::num::NonZeroI64;
73use std::ops::Neg;
74use std::str::FromStr;
75use std::sync::LazyLock;
76use std::sync::{Arc, Mutex};
77use std::thread;
78use std::time::{Duration, Instant};
79use std::{fmt, mem};
80
81use anyhow::Context;
82use chrono::{DateTime, Utc};
83use derivative::Derivative;
84use differential_dataflow::lattice::Lattice;
85use fail::fail_point;
86use futures::StreamExt;
87use futures::future::{BoxFuture, FutureExt, LocalBoxFuture};
88use http::Uri;
89use ipnet::IpNet;
90use itertools::Itertools;
91use mz_adapter_types::bootstrap_builtin_cluster_config::BootstrapBuiltinClusterConfig;
92use mz_adapter_types::compaction::CompactionWindow;
93use mz_adapter_types::connection::ConnectionId;
94use mz_adapter_types::dyncfgs::FRONTEND_READ_THEN_WRITE;
95use mz_adapter_types::dyncfgs::{
96 USER_ID_POOL_BATCH_SIZE, WITH_0DT_DEPLOYMENT_CAUGHT_UP_CHECK_INTERVAL,
97};
98use mz_auth::password::Password;
99use mz_build_info::BuildInfo;
100use mz_catalog::builtin::{
101 BUILTINS, BUILTINS_STATIC, MZ_OBJECT_ARRANGEMENT_SIZE_HISTORY, MZ_STORAGE_USAGE_BY_SHARD,
102};
103use mz_catalog::config::{AwsPrincipalContext, BuiltinItemMigrationConfig, ClusterReplicaSizeMap};
104use mz_catalog::durable::OpenableDurableCatalogState;
105use mz_catalog::expr_cache::{GlobalExpressions, LocalExpressions};
106use mz_catalog::memory::objects::{
107 CatalogEntry, CatalogItem, ClusterReplicaProcessStatus, ClusterVariantManaged, Connection,
108 DataSourceDesc, ReconfigurationTarget, Table, TableDataSource,
109};
110use mz_cloud_resources::{CloudResourceController, VpcEndpointConfig, VpcEndpointEvent};
111use mz_compute_client::as_of_selection;
112use mz_compute_client::controller::error::{
113 CollectionLookupError, CollectionMissing, DataflowCreationError, InstanceMissing,
114};
115use mz_compute_types::ComputeInstanceId;
116use mz_compute_types::dataflows::DataflowDescription;
117use mz_compute_types::plan::LirRelationExpr;
118use mz_controller::clusters::{
119 ClusterConfig, ClusterEvent, ClusterStatus, ProcessId, ReplicaLocation,
120};
121use mz_controller::{ControllerConfig, Readiness};
122use mz_controller_types::{ClusterId, ReplicaId, WatchSetId};
123use mz_dyncfg::{ConfigUpdates, ParameterScope};
124use mz_expr::{MapFilterProject, MirRelationExpr, OptimizedMirRelationExpr, RowSetFinishing};
125use mz_license_keys::{ExpirationBehavior, ValidatedLicenseKey};
126use mz_orchestrator::OfflineReason;
127use mz_ore::cast::{CastFrom, CastInto, CastLossy};
128use mz_ore::channel::trigger::Trigger;
129use mz_ore::future::TimeoutError;
130use mz_ore::metrics::MetricsRegistry;
131use mz_ore::now::{EpochMillis, NowFn};
132use mz_ore::task::{JoinHandle, spawn};
133use mz_ore::thread::JoinHandleExt;
134use mz_ore::tracing::{OpenTelemetryContext, TracingHandle};
135use mz_ore::url::SensitiveUrl;
136use mz_ore::{
137 assert_none, instrument, soft_assert_eq_or_log, soft_assert_or_log, soft_panic_or_log, stack,
138};
139use mz_persist_client::PersistClient;
140use mz_persist_client::batch::ProtoBatch;
141use mz_persist_client::usage::{ShardsUsageReferenced, StorageUsageClient};
142use mz_repr::adt::numeric::Numeric;
143use mz_repr::explain::{ExplainConfig, ExplainFormat};
144use mz_repr::global_id::TransientIdGen;
145use mz_repr::optimize::{OptimizerFeatureOverrides, OptimizerFeatures, OverrideFrom};
146use mz_repr::role_id::RoleId;
147use mz_repr::{CatalogItemId, Diff, GlobalId, RelationDesc, SqlRelationType, Timestamp};
148use mz_secrets::cache::CachingSecretsReader;
149use mz_secrets::{SecretsController, SecretsReader};
150use mz_sql::ast::{Raw, Statement};
151use mz_sql::catalog::{CatalogCluster, EnvironmentId};
152use mz_sql::names::{QualifiedItemName, ResolvedIds, SchemaSpecifier};
153use mz_sql::optimizer_metrics::OptimizerMetrics;
154use mz_sql::plan::{
155 self, AlterSinkPlan, ConnectionDetails, CreateConnectionPlan, HirRelationExpr,
156 NetworkPolicyRule, OnTimeoutAction, Params, QueryWhen,
157};
158use mz_sql::session::user::User;
159use mz_sql::session::vars::{MAX_CREDIT_CONSUMPTION_RATE, SystemVars, Var};
160use mz_sql_parser::ast::ExplainStage;
161use mz_sql_parser::ast::display::AstDisplay;
162use mz_storage_client::client::TableData;
163use mz_storage_client::controller::{CollectionDescription, DataSource, ExportDescription};
164use mz_storage_types::connections::Connection as StorageConnection;
165use mz_storage_types::connections::ConnectionContext;
166use mz_storage_types::connections::inline::{IntoInlineConnection, ReferencedConnection};
167use mz_storage_types::read_holds::ReadHold;
168use mz_storage_types::sinks::{S3SinkFormat, StorageSinkDesc};
169use mz_storage_types::sources::kafka::KAFKA_PROGRESS_DESC;
170use mz_storage_types::sources::{IngestionDescription, SourceExport, Timeline};
171use mz_timestamp_oracle::{TimestampOracleConfig, WriteTimestamp};
172use mz_transform::dataflow::DataflowMetainfo;
173use opentelemetry::trace::TraceContextExt;
174use serde::Serialize;
175use thiserror::Error;
176use timely::progress::{Antichain, Timestamp as _};
177use tokio::runtime::Handle as TokioHandle;
178use tokio::select;
179use tokio::sync::{Notify, OwnedMutexGuard, Semaphore, mpsc, oneshot, watch};
180use tokio::time::{Interval, MissedTickBehavior};
181use tracing::{Instrument, Level, Span, debug, info, info_span, span, warn};
182use tracing_opentelemetry::OpenTelemetrySpanExt;
183use uuid::Uuid;
184
185use crate::active_compute_sink::{ActiveComputeSink, ActiveCopyFrom};
186use crate::catalog::{BuiltinTableUpdate, Catalog, OpenCatalogResult};
187use crate::client::{Client, Handle};
188use crate::command::{Command, ExecuteResponse};
189use crate::config::{
190 ClusterEvalContext, ReplicaEvalContext, ScopedParameters, ScopedParametersScope,
191 SynchronizedParameters, SystemParameterFrontend, SystemParameterSyncConfig,
192};
193use crate::coord::appends::{
194 BuiltinTableAppendCompletion, BuiltinTableAppendNotify, DeferredOp, GroupCommitPermit,
195 PendingWriteTxn,
196};
197use crate::coord::caught_up::CaughtUpCheckContext;
198use crate::coord::id_bundle::CollectionIdBundle;
199use crate::coord::introspection::IntrospectionSubscribe;
200use crate::coord::peek::PendingPeek;
201use crate::coord::statement_logging::StatementLogging;
202use crate::coord::timeline::{TimelineContext, TimelineState};
203use crate::coord::timestamp_selection::{TimestampContext, TimestampDetermination};
204use crate::coord::validity::PlanValidity;
205use crate::error::AdapterError;
206use crate::explain::insights::PlanInsightsContext;
207use crate::explain::optimizer_trace::{DispatchGuard, OptimizerTrace};
208use crate::metrics::Metrics;
209use crate::optimize::dataflows::{ComputeInstanceSnapshot, DataflowBuilder};
210use crate::optimize::{self, Optimize, OptimizerConfig};
211use crate::session::{EndTransactionAction, Session};
212use crate::statement_logging::{
213 StatementEndedExecutionReason, StatementLifecycleEvent, StatementLoggingId,
214};
215use crate::util::{ClientTransmitter, ResultExt, sort_topological};
216use crate::webhook::{WebhookAppenderInvalidator, WebhookConcurrencyLimiter};
217use crate::{AdapterNotice, ReadHolds, flags};
218
219pub(crate) mod appends;
220pub(crate) mod catalog_serving;
221pub(crate) mod cluster_controller;
222pub(crate) mod consistency;
223pub(crate) mod id_bundle;
224pub(crate) mod in_memory_oracle;
225pub(crate) mod peek;
226pub(crate) mod read_policy;
227pub(crate) mod read_then_write;
228pub(crate) mod sequencer;
229pub(crate) mod statement_logging;
230pub(crate) mod timeline;
231pub(crate) mod timestamp_selection;
232
233pub mod catalog_implications;
234mod caught_up;
235mod command_handler;
236mod ddl;
237pub(crate) mod group_sync;
238mod indexes;
239mod info_metrics;
240mod introspection;
241mod message_handler;
242mod privatelink_status;
243mod sql;
244mod validity;
245
246#[derive(Debug)]
272pub(crate) struct IdPool {
273 next: u64,
274 upper: u64,
275}
276
277impl IdPool {
278 pub fn empty() -> Self {
280 IdPool { next: 0, upper: 0 }
281 }
282
283 pub fn allocate(&mut self) -> Option<u64> {
285 if self.next < self.upper {
286 let id = self.next;
287 self.next += 1;
288 Some(id)
289 } else {
290 None
291 }
292 }
293
294 pub fn allocate_many(&mut self, n: u64) -> Option<Vec<u64>> {
297 if self.remaining() >= n {
298 let ids = (self.next..self.next + n).collect();
299 self.next += n;
300 Some(ids)
301 } else {
302 None
303 }
304 }
305
306 pub fn remaining(&self) -> u64 {
308 self.upper - self.next
309 }
310
311 pub fn refill(&mut self, next: u64, upper: u64) {
313 assert!(next <= upper, "invalid pool range: {next}..{upper}");
314 self.next = next;
315 self.upper = upper;
316 }
317}
318
319#[derive(Debug)]
323pub struct ArrangementSizeRecord {
324 pub replica_id: String,
325 pub object_id: String,
326 pub size: i64,
327 pub hydration_complete: bool,
328}
329
330#[derive(Debug)]
331pub enum Message {
332 Command(OpenTelemetryContext, Command),
333 ControllerReady {
334 controller: ControllerReadiness,
335 },
336 PurifiedStatementReady(PurifiedStatementReady),
337 CreateConnectionValidationReady(CreateConnectionValidationReady),
338 AlterConnectionValidationReady(AlterConnectionValidationReady),
339 TryDeferred {
340 conn_id: ConnectionId,
342 acquired_lock: Option<(CatalogItemId, tokio::sync::OwnedMutexGuard<()>)>,
352 },
353 GroupCommitInitiate(Span, Option<GroupCommitPermit>),
355 GroupCommitApplied {
359 responses: Vec<crate::util::CompletedClientTransmitter>,
361 statement_logging_ids: Vec<StatementLoggingId>,
363 internal_results: Vec<crate::coord::appends::InternalWriteResponder>,
365 write_ts: Timestamp,
367 },
368 DeferredStatementReady,
369 AdvanceTimelines,
370 ClusterEvent(ClusterEvent),
371 CancelPendingPeeks {
372 conn_id: ConnectionId,
373 },
374 LinearizeReads,
375 StagedBatches {
376 conn_id: ConnectionId,
377 table_id: CatalogItemId,
378 batches: Vec<Result<ProtoBatch, String>>,
379 },
380 StorageUsageSchedule,
381 StorageUsageFetch,
382 StorageUsageUpdate(ShardsUsageReferenced),
383 StorageUsagePrune(Vec<BuiltinTableUpdate>),
384 ArrangementSizesSchedule,
385 ArrangementSizesSnapshot,
386 ArrangementSizesWrite(Vec<ArrangementSizeRecord>),
387 ArrangementSizesPrune(Vec<BuiltinTableUpdate>),
388 RetireExecute {
391 data: ExecuteContextExtra,
392 otel_ctx: OpenTelemetryContext,
393 reason: StatementEndedExecutionReason,
394 },
395 ExecuteSingleStatementTransaction {
396 ctx: ExecuteContext,
397 otel_ctx: OpenTelemetryContext,
398 stmt: Arc<Statement<Raw>>,
399 params: mz_sql::plan::Params,
400 },
401 PeekStageReady {
402 ctx: ExecuteContext,
403 span: Span,
404 stage: PeekStage,
405 },
406 CreateIndexStageReady {
407 ctx: ExecuteContext,
408 span: Span,
409 stage: CreateIndexStage,
410 },
411 CreateMetricSinkStageReady {
412 ctx: ExecuteContext,
413 span: Span,
414 stage: CreateMetricSinkStage,
415 },
416 CreateViewStageReady {
417 ctx: ExecuteContext,
418 span: Span,
419 stage: CreateViewStage,
420 },
421 CreateMaterializedViewStageReady {
422 ctx: ExecuteContext,
423 span: Span,
424 stage: CreateMaterializedViewStage,
425 },
426 SubscribeStageReady {
427 ctx: ExecuteContext,
428 span: Span,
429 stage: SubscribeStage,
430 },
431 IntrospectionSubscribeStageReady {
432 span: Span,
433 stage: IntrospectionSubscribeStage,
434 },
435 SecretStageReady {
436 ctx: ExecuteContext,
437 span: Span,
438 stage: SecretStage,
439 },
440 ClusterStageReady {
441 ctx: ExecuteContext,
442 span: Span,
443 stage: ClusterStage,
444 },
445 ExplainTimestampStageReady {
446 ctx: ExecuteContext,
447 span: Span,
448 stage: ExplainTimestampStage,
449 },
450 DrainStatementLog,
451 PrivateLinkVpcEndpointEvents(Vec<VpcEndpointEvent>),
452
453 ClusterControllerRequest(cluster_controller::ClusterControllerRequest),
457}
458
459impl Message {
460 pub const fn kind(&self) -> &'static str {
462 match self {
463 Message::Command(_, msg) => match msg {
464 Command::CatalogSnapshot { .. } => "command-catalog_snapshot",
465 Command::Startup { .. } => "command-startup",
466 Command::Execute { .. } => "command-execute",
467 Command::Commit { .. } => "command-commit",
468 Command::CancelRequest { .. } => "command-cancel_request",
469 Command::PrivilegedCancelRequest { .. } => "command-privileged_cancel_request",
470 Command::GetWebhook { .. } => "command-get_webhook",
471 Command::GetSystemVars { .. } => "command-get_system_vars",
472 Command::SetSystemVars { .. } => "command-set_system_vars",
473 Command::UpdateScopedSystemParameters { .. } => {
474 "command-update_scoped_system_parameters"
475 }
476 Command::InstallScopedSystemParameterFrontend { .. } => {
477 "command-install_scoped_system_parameter_frontend"
478 }
479 Command::Terminate { .. } => "command-terminate",
480 Command::RetireExecute { .. } => "command-retire_execute",
481 Command::CheckConsistency { .. } => "command-check_consistency",
482 Command::Dump { .. } => "command-dump",
483 Command::AuthenticatePassword { .. } => "command-auth_check",
484 Command::AuthenticateGetSASLChallenge { .. } => "command-auth_get_sasl_challenge",
485 Command::AuthenticateVerifySASLProof { .. } => "command-auth_verify_sasl_proof",
486 Command::CheckRoleCanLogin { .. } => "command-check_role_can_login",
487 Command::GetComputeInstanceClient { .. } => "get-compute-instance-client",
488 Command::GetOracle { .. } => "get-oracle",
489 Command::DetermineRealTimeRecentTimestamp { .. } => {
490 "determine-real-time-recent-timestamp"
491 }
492 Command::GetTransactionReadHoldsBundle { .. } => {
493 "get-transaction-read-holds-bundle"
494 }
495 Command::StoreTransactionReadHolds { .. } => "store-transaction-read-holds",
496 Command::ExecuteSlowPathPeek { .. } => "execute-slow-path-peek",
497 Command::ExecuteSubscribe { .. } => "execute-subscribe",
498 Command::CopyToPreflight { .. } => "copy-to-preflight",
499 Command::ExecuteCopyTo { .. } => "execute-copy-to",
500 Command::ExecuteSideEffectingFunc { .. } => "execute-side-effecting-func",
501 Command::LookupConnection { .. } => "lookup-connection",
502 Command::RegisterFrontendPeek { .. } => "register-frontend-peek",
503 Command::UnregisterFrontendPeek { .. } => "unregister-frontend-peek",
504 Command::ExplainTimestamp { .. } => "explain-timestamp",
505 Command::FrontendStatementLogging(..) => "frontend-statement-logging",
506 Command::StartCopyFromStdin { .. } => "start-copy-from-stdin",
507 Command::InjectAuditEvents { .. } => "inject-audit-events",
508 Command::RegisterConnectionCancelWatch { .. } => "register-connection-cancel-watch",
509 Command::CreateInternalSubscribe { .. } => "create-internal-subscribe",
510 Command::AttemptWrite { .. } => "attempt-write",
511 Command::DropInternalSubscribe { .. } => "drop-internal-subscribe",
512 },
513 Message::ControllerReady {
514 controller: ControllerReadiness::Compute,
515 } => "controller_ready(compute)",
516 Message::ControllerReady {
517 controller: ControllerReadiness::Storage,
518 } => "controller_ready(storage)",
519 Message::ControllerReady {
520 controller: ControllerReadiness::Metrics,
521 } => "controller_ready(metrics)",
522 Message::ControllerReady {
523 controller: ControllerReadiness::Internal,
524 } => "controller_ready(internal)",
525 Message::PurifiedStatementReady(_) => "purified_statement_ready",
526 Message::CreateConnectionValidationReady(_) => "create_connection_validation_ready",
527 Message::TryDeferred { .. } => "try_deferred",
528 Message::GroupCommitInitiate(..) => "group_commit_initiate",
529 Message::GroupCommitApplied { .. } => "group_commit_applied",
530 Message::AdvanceTimelines => "advance_timelines",
531 Message::ClusterEvent(_) => "cluster_event",
532 Message::CancelPendingPeeks { .. } => "cancel_pending_peeks",
533 Message::LinearizeReads => "linearize_reads",
534 Message::StagedBatches { .. } => "staged_batches",
535 Message::StorageUsageSchedule => "storage_usage_schedule",
536 Message::StorageUsageFetch => "storage_usage_fetch",
537 Message::StorageUsageUpdate(_) => "storage_usage_update",
538 Message::StorageUsagePrune(_) => "storage_usage_prune",
539 Message::ArrangementSizesSchedule => "arrangement_sizes_schedule",
540 Message::ArrangementSizesSnapshot => "arrangement_sizes_snapshot",
541 Message::ArrangementSizesWrite(_) => "arrangement_sizes_write",
542 Message::ArrangementSizesPrune(_) => "arrangement_sizes_prune",
543 Message::RetireExecute { .. } => "retire_execute",
544 Message::ExecuteSingleStatementTransaction { .. } => {
545 "execute_single_statement_transaction"
546 }
547 Message::PeekStageReady { .. } => "peek_stage_ready",
548 Message::ExplainTimestampStageReady { .. } => "explain_timestamp_stage_ready",
549 Message::CreateIndexStageReady { .. } => "create_index_stage_ready",
550 Message::CreateMetricSinkStageReady { .. } => "create_metric_sink_stage_ready",
551 Message::CreateViewStageReady { .. } => "create_view_stage_ready",
552 Message::CreateMaterializedViewStageReady { .. } => {
553 "create_materialized_view_stage_ready"
554 }
555 Message::SubscribeStageReady { .. } => "subscribe_stage_ready",
556 Message::IntrospectionSubscribeStageReady { .. } => {
557 "introspection_subscribe_stage_ready"
558 }
559 Message::SecretStageReady { .. } => "secret_stage_ready",
560 Message::ClusterStageReady { .. } => "cluster_stage_ready",
561 Message::DrainStatementLog => "drain_statement_log",
562 Message::AlterConnectionValidationReady(..) => "alter_connection_validation_ready",
563 Message::PrivateLinkVpcEndpointEvents(_) => "private_link_vpc_endpoint_events",
564 Message::ClusterControllerRequest(_) => "cluster_controller_request",
565 Message::DeferredStatementReady => "deferred_statement_ready",
566 }
567 }
568}
569
570#[derive(Debug)]
572pub enum ControllerReadiness {
573 Storage,
575 Compute,
577 Metrics,
579 Internal,
581}
582
583#[derive(Derivative)]
584#[derivative(Debug)]
585pub struct BackgroundWorkResult<T> {
586 #[derivative(Debug = "ignore")]
587 pub ctx: ExecuteContext,
588 pub result: Result<T, AdapterError>,
589 pub params: Params,
590 pub plan_validity: PlanValidity,
591 pub original_stmt: Arc<Statement<Raw>>,
592 pub otel_ctx: OpenTelemetryContext,
593}
594
595pub type PurifiedStatementReady = BackgroundWorkResult<mz_sql::pure::PurifiedStatement>;
596
597#[derive(Derivative)]
598#[derivative(Debug)]
599pub struct ValidationReady<T> {
600 #[derivative(Debug = "ignore")]
601 pub ctx: ExecuteContext,
602 pub result: Result<T, AdapterError>,
603 pub resolved_ids: ResolvedIds,
604 pub connection_id: CatalogItemId,
605 pub connection_gid: GlobalId,
606 pub plan_validity: PlanValidity,
607 pub otel_ctx: OpenTelemetryContext,
608}
609
610pub type CreateConnectionValidationReady = ValidationReady<CreateConnectionPlan>;
611pub type AlterConnectionValidationReady = ValidationReady<Connection>;
612
613#[derive(Debug)]
614pub enum PeekStage {
615 LinearizeTimestamp(PeekStageLinearizeTimestamp),
617 RealTimeRecency(PeekStageRealTimeRecency),
618 TimestampReadHold(PeekStageTimestampReadHold),
619 Optimize(PeekStageOptimize),
620 Finish(PeekStageFinish),
622 ExplainPlan(PeekStageExplainPlan),
624 ExplainPushdown(PeekStageExplainPushdown),
625 CopyToPreflight(PeekStageCopyTo),
627 CopyToDataflow(PeekStageCopyTo),
629}
630
631#[derive(Debug)]
632pub struct CopyToContext {
633 pub desc: RelationDesc,
635 pub uri: Uri,
637 pub connection: StorageConnection<ReferencedConnection>,
639 pub connection_id: CatalogItemId,
641 pub format: S3SinkFormat,
643 pub max_file_size: u64,
645 pub output_batch_count: Option<u64>,
650}
651
652#[derive(Debug)]
653pub struct PeekStageLinearizeTimestamp {
654 validity: PlanValidity,
655 plan: mz_sql::plan::SelectPlan,
656 max_query_result_size: Option<u64>,
657 source_ids: BTreeSet<GlobalId>,
658 target_replica: Option<ReplicaId>,
659 timeline_context: TimelineContext,
660 optimizer: optimize::PeekOptimizer,
661 explain_ctx: ExplainContext,
664}
665
666#[derive(Debug)]
667pub struct PeekStageRealTimeRecency {
668 validity: PlanValidity,
669 plan: mz_sql::plan::SelectPlan,
670 max_query_result_size: Option<u64>,
671 source_ids: BTreeSet<GlobalId>,
672 target_replica: Option<ReplicaId>,
673 timeline_context: TimelineContext,
674 oracle_read_ts: Option<Timestamp>,
675 optimizer: optimize::PeekOptimizer,
676 explain_ctx: ExplainContext,
679}
680
681#[derive(Debug)]
682pub struct PeekStageTimestampReadHold {
683 validity: PlanValidity,
684 plan: mz_sql::plan::SelectPlan,
685 max_query_result_size: Option<u64>,
686 source_ids: BTreeSet<GlobalId>,
687 target_replica: Option<ReplicaId>,
688 timeline_context: TimelineContext,
689 oracle_read_ts: Option<Timestamp>,
690 real_time_recency_ts: Option<mz_repr::Timestamp>,
691 optimizer: optimize::PeekOptimizer,
692 explain_ctx: ExplainContext,
695}
696
697#[derive(Debug)]
698pub struct PeekStageOptimize {
699 validity: PlanValidity,
700 plan: mz_sql::plan::SelectPlan,
701 max_query_result_size: Option<u64>,
702 source_ids: BTreeSet<GlobalId>,
703 id_bundle: CollectionIdBundle,
704 target_replica: Option<ReplicaId>,
705 determination: TimestampDetermination,
706 optimizer: optimize::PeekOptimizer,
707 explain_ctx: ExplainContext,
710}
711
712#[derive(Debug)]
713pub struct PeekStageFinish {
714 validity: PlanValidity,
715 plan: mz_sql::plan::SelectPlan,
716 max_query_result_size: Option<u64>,
717 id_bundle: CollectionIdBundle,
718 target_replica: Option<ReplicaId>,
719 source_ids: BTreeSet<GlobalId>,
720 determination: TimestampDetermination,
721 cluster_id: ComputeInstanceId,
722 finishing: RowSetFinishing,
723 plan_insights_optimizer_trace: Option<OptimizerTrace>,
726 insights_ctx: Option<Box<PlanInsightsContext>>,
727 global_lir_plan: optimize::peek::GlobalLirPlan,
728 optimization_finished_at: EpochMillis,
729}
730
731#[derive(Debug)]
732pub struct PeekStageCopyTo {
733 validity: PlanValidity,
734 optimizer: optimize::copy_to::Optimizer,
735 global_lir_plan: optimize::copy_to::GlobalLirPlan,
736 optimization_finished_at: EpochMillis,
737 target_replica: Option<ReplicaId>,
738 source_ids: BTreeSet<GlobalId>,
739}
740
741#[derive(Debug)]
742pub struct PeekStageExplainPlan {
743 validity: PlanValidity,
744 optimizer: optimize::peek::Optimizer,
745 df_meta: DataflowMetainfo,
746 explain_ctx: ExplainPlanContext,
747 insights_ctx: Option<Box<PlanInsightsContext>>,
748}
749
750#[derive(Debug)]
751pub struct PeekStageExplainPushdown {
752 validity: PlanValidity,
753 determination: TimestampDetermination,
754 imports: BTreeMap<GlobalId, MapFilterProject>,
755}
756
757#[derive(Debug)]
758pub enum CreateIndexStage {
759 Optimize(CreateIndexOptimize),
760 Finish(CreateIndexFinish),
761 Explain(CreateIndexExplain),
762}
763
764#[derive(Debug)]
765pub struct CreateIndexOptimize {
766 validity: PlanValidity,
767 plan: plan::CreateIndexPlan,
768 resolved_ids: ResolvedIds,
769 explain_ctx: ExplainContext,
772}
773
774#[derive(Debug)]
775pub struct CreateIndexFinish {
776 validity: PlanValidity,
777 item_id: CatalogItemId,
778 global_id: GlobalId,
779 plan: plan::CreateIndexPlan,
780 resolved_ids: ResolvedIds,
781 global_mir_plan: optimize::index::GlobalMirPlan,
782 global_lir_plan: optimize::index::GlobalLirPlan,
783 optimizer_features: OptimizerFeatures,
784}
785
786#[derive(Debug)]
787pub struct CreateIndexExplain {
788 validity: PlanValidity,
789 exported_index_id: GlobalId,
790 plan: plan::CreateIndexPlan,
791 df_meta: DataflowMetainfo,
792 explain_ctx: ExplainPlanContext,
793}
794
795#[derive(Debug)]
796pub enum CreateMetricSinkStage {
797 Optimize(CreateMetricSinkOptimize),
798 Finish(CreateMetricSinkFinish),
799}
800
801#[derive(Debug)]
802pub struct CreateMetricSinkOptimize {
803 validity: PlanValidity,
804 plan: plan::CreateMetricSinkPlan,
805 resolved_ids: ResolvedIds,
806}
807
808#[derive(Debug)]
809pub struct CreateMetricSinkFinish {
810 validity: PlanValidity,
811 item_id: CatalogItemId,
812 global_id: GlobalId,
813 plan: plan::CreateMetricSinkPlan,
814 resolved_ids: ResolvedIds,
815 global_mir_plan: optimize::metric_sink::GlobalMirPlan,
816 global_lir_plan: optimize::metric_sink::GlobalLirPlan,
817 optimizer_features: OptimizerFeatures,
818}
819
820#[derive(Debug)]
821pub enum CreateViewStage {
822 Optimize(CreateViewOptimize),
823 Finish(CreateViewFinish),
824 Explain(CreateViewExplain),
825}
826
827#[derive(Debug)]
828pub struct CreateViewOptimize {
829 validity: PlanValidity,
830 plan: plan::CreateViewPlan,
831 resolved_ids: ResolvedIds,
832 explain_ctx: ExplainContext,
835}
836
837#[derive(Debug)]
838pub struct CreateViewFinish {
839 validity: PlanValidity,
840 item_id: CatalogItemId,
842 global_id: GlobalId,
844 plan: plan::CreateViewPlan,
845 resolved_ids: ResolvedIds,
847 optimized_expr: OptimizedMirRelationExpr,
848}
849
850#[derive(Debug)]
851pub struct CreateViewExplain {
852 validity: PlanValidity,
853 id: GlobalId,
854 plan: plan::CreateViewPlan,
855 explain_ctx: ExplainPlanContext,
856}
857
858#[derive(Debug)]
859pub enum ExplainTimestampStage {
860 Optimize(ExplainTimestampOptimize),
861 RealTimeRecency(ExplainTimestampRealTimeRecency),
862 LinearizeTimestamp(ExplainTimestampLinearizeTimestamp),
863 Finish(ExplainTimestampFinish),
864}
865
866#[derive(Debug)]
867pub struct ExplainTimestampOptimize {
868 validity: PlanValidity,
869 plan: plan::ExplainTimestampPlan,
870 cluster_id: ClusterId,
871}
872
873#[derive(Debug)]
874pub struct ExplainTimestampRealTimeRecency {
875 validity: PlanValidity,
876 format: ExplainFormat,
877 optimized_plan: OptimizedMirRelationExpr,
878 cluster_id: ClusterId,
879 when: QueryWhen,
880}
881
882#[derive(Debug)]
883pub struct ExplainTimestampLinearizeTimestamp {
884 validity: PlanValidity,
885 format: ExplainFormat,
886 optimized_plan: OptimizedMirRelationExpr,
887 cluster_id: ClusterId,
888 source_ids: BTreeSet<GlobalId>,
889 when: QueryWhen,
890 real_time_recency_ts: Option<Timestamp>,
891}
892
893#[derive(Debug)]
894pub struct ExplainTimestampFinish {
895 validity: PlanValidity,
896 format: ExplainFormat,
897 cluster_id: ClusterId,
898 source_ids: BTreeSet<GlobalId>,
899 when: QueryWhen,
900 real_time_recency_ts: Option<Timestamp>,
901 timeline_context: TimelineContext,
904 oracle_read_ts: Option<Timestamp>,
908}
909
910#[derive(Debug)]
911pub enum ClusterStage {
912 Alter(AlterCluster),
913 WaitForHydrated(AlterClusterWaitForHydrated),
914 Finalize(AlterClusterFinalize),
915 AwaitReconfiguration(AlterClusterAwaitReconfiguration),
920}
921
922#[derive(Debug)]
923pub struct AlterCluster {
924 validity: PlanValidity,
925 plan: plan::AlterClusterPlan,
926}
927
928#[derive(Debug)]
929pub struct AlterClusterWaitForHydrated {
930 validity: PlanValidity,
931 plan: plan::AlterClusterPlan,
932 new_config: ClusterVariantManaged,
933 workload_class: Option<String>,
934 timeout_time: Instant,
935 on_timeout: OnTimeoutAction,
936}
937
938#[derive(Debug)]
939pub struct AlterClusterFinalize {
940 validity: PlanValidity,
941 plan: plan::AlterClusterPlan,
942 new_config: ClusterVariantManaged,
943 workload_class: Option<String>,
944}
945
946#[derive(Debug)]
947pub struct AlterClusterAwaitReconfiguration {
948 validity: PlanValidity,
949 cluster_id: ClusterId,
950 target: ReconfigurationTarget,
954}
955
956#[derive(Debug)]
957pub enum ExplainContext {
958 None,
960 Plan(ExplainPlanContext),
962 PlanInsightsNotice(OptimizerTrace),
965 Pushdown,
967}
968
969impl ExplainContext {
970 pub(crate) fn dispatch_guard(&self) -> Option<DispatchGuard<'_>> {
974 let optimizer_trace = match self {
975 ExplainContext::Plan(explain_ctx) => Some(&explain_ctx.optimizer_trace),
976 ExplainContext::PlanInsightsNotice(optimizer_trace) => Some(optimizer_trace),
977 _ => None,
978 };
979 optimizer_trace.map(|optimizer_trace| optimizer_trace.as_guard())
980 }
981
982 pub(crate) fn needs_cluster(&self) -> bool {
983 match self {
984 ExplainContext::None => true,
985 ExplainContext::Plan(..) => false,
986 ExplainContext::PlanInsightsNotice(..) => true,
987 ExplainContext::Pushdown => false,
988 }
989 }
990
991 pub(crate) fn needs_plan_insights(&self) -> bool {
992 matches!(
993 self,
994 ExplainContext::Plan(ExplainPlanContext {
995 stage: ExplainStage::PlanInsights,
996 ..
997 }) | ExplainContext::PlanInsightsNotice(_)
998 )
999 }
1000}
1001
1002#[derive(Debug)]
1003pub struct ExplainPlanContext {
1004 pub broken: bool,
1009 pub config: ExplainConfig,
1010 pub format: ExplainFormat,
1011 pub stage: ExplainStage,
1012 pub replan: Option<GlobalId>,
1013 pub desc: Option<RelationDesc>,
1014 pub optimizer_trace: OptimizerTrace,
1015}
1016
1017#[derive(Debug)]
1018pub enum CreateMaterializedViewStage {
1019 Optimize(CreateMaterializedViewOptimize),
1020 Finish(CreateMaterializedViewFinish),
1021 Explain(CreateMaterializedViewExplain),
1022}
1023
1024#[derive(Debug)]
1025pub struct CreateMaterializedViewOptimize {
1026 validity: PlanValidity,
1027 plan: plan::CreateMaterializedViewPlan,
1028 resolved_ids: ResolvedIds,
1029 explain_ctx: ExplainContext,
1032}
1033
1034#[derive(Debug)]
1035pub struct CreateMaterializedViewFinish {
1036 item_id: CatalogItemId,
1038 global_id: GlobalId,
1040 validity: PlanValidity,
1041 plan: plan::CreateMaterializedViewPlan,
1042 resolved_ids: ResolvedIds,
1043 local_mir_plan: optimize::materialized_view::LocalMirPlan,
1044 global_mir_plan: optimize::materialized_view::GlobalMirPlan,
1045 global_lir_plan: optimize::materialized_view::GlobalLirPlan,
1046 optimizer_features: OptimizerFeatures,
1047}
1048
1049#[derive(Debug)]
1050pub struct CreateMaterializedViewExplain {
1051 global_id: GlobalId,
1052 validity: PlanValidity,
1053 plan: plan::CreateMaterializedViewPlan,
1054 df_meta: DataflowMetainfo,
1055 explain_ctx: ExplainPlanContext,
1056}
1057
1058#[derive(Debug)]
1059pub enum SubscribeStage {
1060 OptimizeMir(SubscribeOptimizeMir),
1061 LinearizeTimestamp(SubscribeLinearizeTimestamp),
1062 TimestampOptimizeLir(SubscribeTimestampOptimizeLir),
1063 Finish(SubscribeFinish),
1064 Explain(SubscribeExplain),
1065}
1066
1067#[derive(Debug)]
1068pub struct SubscribeOptimizeMir {
1069 validity: PlanValidity,
1070 plan: plan::SubscribePlan,
1071 timeline: TimelineContext,
1072 dependency_ids: BTreeSet<GlobalId>,
1073 cluster_id: ComputeInstanceId,
1074 replica_id: Option<ReplicaId>,
1075 explain_ctx: ExplainContext,
1078}
1079
1080#[derive(Debug)]
1081pub struct SubscribeLinearizeTimestamp {
1082 validity: PlanValidity,
1083 plan: plan::SubscribePlan,
1084 timeline: TimelineContext,
1085 optimizer: optimize::subscribe::Optimizer,
1086 global_mir_plan: optimize::subscribe::GlobalMirPlan<optimize::subscribe::Unresolved>,
1087 dependency_ids: BTreeSet<GlobalId>,
1088 replica_id: Option<ReplicaId>,
1089 explain_ctx: ExplainContext,
1092}
1093
1094#[derive(Debug)]
1095pub struct SubscribeTimestampOptimizeLir {
1096 validity: PlanValidity,
1097 plan: plan::SubscribePlan,
1098 timeline: TimelineContext,
1099 optimizer: optimize::subscribe::Optimizer,
1100 global_mir_plan: optimize::subscribe::GlobalMirPlan<optimize::subscribe::Unresolved>,
1101 dependency_ids: BTreeSet<GlobalId>,
1102 replica_id: Option<ReplicaId>,
1103 oracle_read_ts: Option<Timestamp>,
1107 explain_ctx: ExplainContext,
1110}
1111
1112#[derive(Debug)]
1113pub struct SubscribeFinish {
1114 validity: PlanValidity,
1115 cluster_id: ComputeInstanceId,
1116 replica_id: Option<ReplicaId>,
1117 plan: plan::SubscribePlan,
1118 global_lir_plan: optimize::subscribe::GlobalLirPlan,
1119 dependency_ids: BTreeSet<GlobalId>,
1120}
1121
1122#[derive(Debug)]
1123pub struct SubscribeExplain {
1124 validity: PlanValidity,
1125 optimizer: optimize::subscribe::Optimizer,
1126 df_meta: DataflowMetainfo,
1127 cluster_id: ComputeInstanceId,
1128 explain_ctx: ExplainPlanContext,
1129}
1130
1131#[derive(Debug)]
1132pub enum IntrospectionSubscribeStage {
1133 OptimizeMir(IntrospectionSubscribeOptimizeMir),
1134 TimestampOptimizeLir(IntrospectionSubscribeTimestampOptimizeLir),
1135 Finish(IntrospectionSubscribeFinish),
1136}
1137
1138#[derive(Debug)]
1139pub struct IntrospectionSubscribeOptimizeMir {
1140 validity: PlanValidity,
1141 plan: plan::SubscribePlan,
1142 subscribe_id: GlobalId,
1143 cluster_id: ComputeInstanceId,
1144 replica_id: ReplicaId,
1145}
1146
1147#[derive(Debug)]
1148pub struct IntrospectionSubscribeTimestampOptimizeLir {
1149 validity: PlanValidity,
1150 optimizer: optimize::subscribe::Optimizer,
1151 global_mir_plan: optimize::subscribe::GlobalMirPlan<optimize::subscribe::Unresolved>,
1152 cluster_id: ComputeInstanceId,
1153 replica_id: ReplicaId,
1154}
1155
1156#[derive(Debug)]
1157pub struct IntrospectionSubscribeFinish {
1158 validity: PlanValidity,
1159 global_lir_plan: optimize::subscribe::GlobalLirPlan,
1160 read_holds: ReadHolds,
1161 cluster_id: ComputeInstanceId,
1162 replica_id: ReplicaId,
1163}
1164
1165#[derive(Debug)]
1166pub enum SecretStage {
1167 CreateEnsure(CreateSecretEnsure),
1168 CreateFinish(CreateSecretFinish),
1169 RotateKeysEnsure(RotateKeysSecretEnsure),
1170 RotateKeysFinish(RotateKeysSecretFinish),
1171 Alter(AlterSecret),
1172}
1173
1174#[derive(Debug)]
1175pub struct CreateSecretEnsure {
1176 validity: PlanValidity,
1177 plan: plan::CreateSecretPlan,
1178}
1179
1180#[derive(Debug)]
1181pub struct CreateSecretFinish {
1182 validity: PlanValidity,
1183 item_id: CatalogItemId,
1184 global_id: GlobalId,
1185 plan: plan::CreateSecretPlan,
1186}
1187
1188#[derive(Debug)]
1189pub struct RotateKeysSecretEnsure {
1190 validity: PlanValidity,
1191 id: CatalogItemId,
1192}
1193
1194#[derive(Debug)]
1195pub struct RotateKeysSecretFinish {
1196 validity: PlanValidity,
1197 ops: Vec<crate::catalog::Op>,
1198}
1199
1200#[derive(Debug)]
1201pub struct AlterSecret {
1202 validity: PlanValidity,
1203 plan: plan::AlterSecretPlan,
1204}
1205
1206#[derive(Debug, Copy, Clone, PartialEq, Eq)]
1211pub enum TargetCluster {
1212 CatalogServer,
1214 Active,
1216 Transaction(ClusterId),
1218}
1219
1220pub(crate) enum StageResult<T> {
1222 Handle(JoinHandle<Result<T, AdapterError>>),
1224 HandleRetire(JoinHandle<Result<ExecuteResponse, AdapterError>>),
1226 Immediate(T),
1228 Response(ExecuteResponse),
1230}
1231
1232pub(crate) trait Staged: Send {
1234 type Ctx: StagedContext;
1235
1236 fn validity(&mut self) -> &mut PlanValidity;
1237
1238 async fn stage(
1240 self,
1241 coord: &mut Coordinator,
1242 ctx: &mut Self::Ctx,
1243 ) -> Result<StageResult<Box<Self>>, AdapterError>;
1244
1245 fn message(self, ctx: Self::Ctx, span: Span) -> Message;
1247
1248 fn cancel_enabled(&self) -> bool;
1250}
1251
1252pub trait StagedContext {
1253 fn retire(self, result: Result<ExecuteResponse, AdapterError>);
1254 fn session(&self) -> Option<&Session>;
1255}
1256
1257impl StagedContext for ExecuteContext {
1258 fn retire(self, result: Result<ExecuteResponse, AdapterError>) {
1259 self.retire(result);
1260 }
1261
1262 fn session(&self) -> Option<&Session> {
1263 Some(self.session())
1264 }
1265}
1266
1267impl StagedContext for () {
1268 fn retire(self, _result: Result<ExecuteResponse, AdapterError>) {}
1269
1270 fn session(&self) -> Option<&Session> {
1271 None
1272 }
1273}
1274
1275pub struct Config {
1277 pub controller_config: ControllerConfig,
1278 pub controller_envd_epoch: NonZeroI64,
1279 pub storage: Box<dyn mz_catalog::durable::DurableCatalogState>,
1280 pub timestamp_oracle_url: Option<SensitiveUrl>,
1281 pub unsafe_mode: bool,
1282 pub all_features: bool,
1283 pub build_info: &'static BuildInfo,
1284 pub environment_id: EnvironmentId,
1285 pub metrics_registry: MetricsRegistry,
1286 pub now: NowFn,
1287 pub secrets_controller: Arc<dyn SecretsController>,
1288 pub cloud_resource_controller: Option<Arc<dyn CloudResourceController>>,
1289 pub availability_zones: Vec<String>,
1290 pub cluster_replica_sizes: ClusterReplicaSizeMap,
1291 pub builtin_system_cluster_config: BootstrapBuiltinClusterConfig,
1292 pub builtin_catalog_server_cluster_config: BootstrapBuiltinClusterConfig,
1293 pub builtin_probe_cluster_config: BootstrapBuiltinClusterConfig,
1294 pub builtin_support_cluster_config: BootstrapBuiltinClusterConfig,
1295 pub builtin_analytics_cluster_config: BootstrapBuiltinClusterConfig,
1296 pub system_parameter_defaults: BTreeMap<String, String>,
1297 pub storage_usage_client: StorageUsageClient,
1298 pub storage_usage_collection_interval: Duration,
1299 pub storage_usage_retention_period: Option<Duration>,
1300 pub segment_client: Option<mz_segment::Client>,
1301 pub egress_addresses: Vec<IpNet>,
1302 pub remote_system_parameters: Option<BTreeMap<String, String>>,
1303 pub aws_account_id: Option<String>,
1304 pub aws_privatelink_availability_zones: Option<Vec<String>>,
1305 pub connection_context: ConnectionContext,
1306 pub connection_limit_callback: Box<dyn Fn(u64, u64) -> () + Send + Sync + 'static>,
1307 pub webhook_concurrency_limit: WebhookConcurrencyLimiter,
1308 pub http_host_name: Option<String>,
1309 pub tracing_handle: TracingHandle,
1310 pub read_only_controllers: bool,
1314
1315 pub caught_up_trigger: Option<Trigger>,
1319
1320 pub helm_chart_version: Option<String>,
1321 pub license_key: ValidatedLicenseKey,
1322 pub external_login_password_mz_system: Option<Password>,
1323 pub force_builtin_schema_migration: Option<String>,
1324}
1325
1326#[derive(Debug, Serialize)]
1328pub struct ConnMeta {
1329 secret_key: u32,
1334 connected_at: EpochMillis,
1336 user: User,
1337 application_name: String,
1338 uuid: Uuid,
1339 conn_id: ConnectionId,
1340 client_ip: Option<IpAddr>,
1341
1342 drop_sinks: BTreeSet<GlobalId>,
1345
1346 #[serde(skip)]
1348 deferred_lock: Option<OwnedMutexGuard<()>>,
1349
1350 pending_cluster_alters: BTreeSet<ClusterId>,
1353
1354 #[serde(skip)]
1356 notice_tx: mpsc::UnboundedSender<AdapterNotice>,
1357
1358 authenticated_role: RoleId,
1362}
1363
1364impl ConnMeta {
1365 pub fn conn_id(&self) -> &ConnectionId {
1366 &self.conn_id
1367 }
1368
1369 pub fn user(&self) -> &User {
1370 &self.user
1371 }
1372
1373 pub fn application_name(&self) -> &str {
1374 &self.application_name
1375 }
1376
1377 pub fn authenticated_role_id(&self) -> &RoleId {
1378 &self.authenticated_role
1379 }
1380
1381 pub fn uuid(&self) -> Uuid {
1382 self.uuid
1383 }
1384
1385 pub fn client_ip(&self) -> Option<IpAddr> {
1386 self.client_ip
1387 }
1388
1389 pub fn connected_at(&self) -> EpochMillis {
1390 self.connected_at
1391 }
1392}
1393
1394#[derive(Debug)]
1395pub struct PendingTxn {
1397 ctx: ExecuteContext,
1399 response: Result<PendingTxnResponse, AdapterError>,
1401 action: EndTransactionAction,
1403}
1404
1405#[derive(Debug)]
1406pub enum PendingTxnResponse {
1408 Committed {
1410 params: BTreeMap<&'static str, String>,
1412 },
1413 Rolledback {
1415 params: BTreeMap<&'static str, String>,
1417 },
1418}
1419
1420impl PendingTxnResponse {
1421 pub fn extend_params(&mut self, p: impl IntoIterator<Item = (&'static str, String)>) {
1422 match self {
1423 PendingTxnResponse::Committed { params }
1424 | PendingTxnResponse::Rolledback { params } => params.extend(p),
1425 }
1426 }
1427}
1428
1429impl From<PendingTxnResponse> for ExecuteResponse {
1430 fn from(value: PendingTxnResponse) -> Self {
1431 match value {
1432 PendingTxnResponse::Committed { params } => {
1433 ExecuteResponse::TransactionCommitted { params }
1434 }
1435 PendingTxnResponse::Rolledback { params } => {
1436 ExecuteResponse::TransactionRolledBack { params }
1437 }
1438 }
1439 }
1440}
1441
1442#[derive(Debug)]
1443pub struct PendingReadTxn {
1445 txn: PendingRead,
1447 timestamp_context: TimestampContext,
1449 created: Instant,
1451 num_requeues: u64,
1455 otel_ctx: OpenTelemetryContext,
1457}
1458
1459impl PendingReadTxn {
1460 pub fn timestamp_context(&self) -> &TimestampContext {
1462 &self.timestamp_context
1463 }
1464
1465 pub(crate) fn take_context(self) -> ExecuteContext {
1466 self.txn.take_context()
1467 }
1468}
1469
1470#[derive(Debug)]
1471enum PendingRead {
1473 Read {
1474 txn: PendingTxn,
1476 },
1477 ReadThenWrite {
1478 ctx: ExecuteContext,
1480 tx: oneshot::Sender<Option<ExecuteContext>>,
1483 },
1484}
1485
1486impl PendingRead {
1487 #[instrument(level = "debug")]
1492 pub fn finish(self) -> Option<(ExecuteContext, Result<ExecuteResponse, AdapterError>)> {
1493 match self {
1494 PendingRead::Read {
1495 txn:
1496 PendingTxn {
1497 mut ctx,
1498 response,
1499 action,
1500 },
1501 ..
1502 } => {
1503 let changed = ctx.session_mut().vars_mut().end_transaction(action);
1504 let response = response.map(|mut r| {
1506 r.extend_params(changed);
1507 ExecuteResponse::from(r)
1508 });
1509
1510 Some((ctx, response))
1511 }
1512 PendingRead::ReadThenWrite { ctx, tx, .. } => {
1513 let _ = tx.send(Some(ctx));
1515 None
1516 }
1517 }
1518 }
1519
1520 fn label(&self) -> &'static str {
1521 match self {
1522 PendingRead::Read { .. } => "read",
1523 PendingRead::ReadThenWrite { .. } => "read_then_write",
1524 }
1525 }
1526
1527 pub(crate) fn take_context(self) -> ExecuteContext {
1528 match self {
1529 PendingRead::Read { txn, .. } => txn.ctx,
1530 PendingRead::ReadThenWrite { ctx, tx, .. } => {
1531 let _ = tx.send(None);
1534 ctx
1535 }
1536 }
1537 }
1538}
1539
1540#[derive(Debug, Default)]
1550#[must_use]
1551pub struct ExecuteContextExtra {
1552 statement_uuid: Option<StatementLoggingId>,
1553}
1554
1555impl ExecuteContextExtra {
1556 pub(crate) fn new(statement_uuid: Option<StatementLoggingId>) -> Self {
1557 Self { statement_uuid }
1558 }
1559 pub fn is_trivial(&self) -> bool {
1560 self.statement_uuid.is_none()
1561 }
1562 pub fn contents(&self) -> Option<StatementLoggingId> {
1563 self.statement_uuid
1564 }
1565 #[must_use]
1569 pub(crate) fn retire(self) -> Option<StatementLoggingId> {
1570 self.statement_uuid
1571 }
1572}
1573
1574#[derive(Debug)]
1584#[must_use]
1585pub struct ExecuteContextGuard {
1586 extra: ExecuteContextExtra,
1587 coordinator_tx: mpsc::UnboundedSender<Message>,
1592}
1593
1594impl Default for ExecuteContextGuard {
1595 fn default() -> Self {
1596 let (tx, _rx) = mpsc::unbounded_channel();
1600 Self {
1601 extra: ExecuteContextExtra::default(),
1602 coordinator_tx: tx,
1603 }
1604 }
1605}
1606
1607impl ExecuteContextGuard {
1608 pub(crate) fn new(
1609 statement_uuid: Option<StatementLoggingId>,
1610 coordinator_tx: mpsc::UnboundedSender<Message>,
1611 ) -> Self {
1612 Self {
1613 extra: ExecuteContextExtra::new(statement_uuid),
1614 coordinator_tx,
1615 }
1616 }
1617 pub fn is_trivial(&self) -> bool {
1618 self.extra.is_trivial()
1619 }
1620 pub fn contents(&self) -> Option<StatementLoggingId> {
1621 self.extra.contents()
1622 }
1623 pub(crate) fn defuse(mut self) -> ExecuteContextExtra {
1630 std::mem::take(&mut self.extra)
1632 }
1633}
1634
1635impl Drop for ExecuteContextGuard {
1636 fn drop(&mut self) {
1637 if let Some(statement_uuid) = self.extra.statement_uuid.take() {
1638 let msg = Message::RetireExecute {
1641 data: ExecuteContextExtra {
1642 statement_uuid: Some(statement_uuid),
1643 },
1644 otel_ctx: OpenTelemetryContext::obtain(),
1645 reason: StatementEndedExecutionReason::Aborted,
1646 };
1647 let _ = self.coordinator_tx.send(msg);
1650 }
1651 }
1652}
1653
1654#[derive(Debug)]
1659pub struct ExecuteContext {
1660 inner: Option<Box<ExecuteContextInner>>,
1662}
1663
1664impl std::ops::Deref for ExecuteContext {
1665 type Target = ExecuteContextInner;
1666 fn deref(&self) -> &Self::Target {
1667 self.inner.as_ref().expect("only consumed by value")
1668 }
1669}
1670
1671impl std::ops::DerefMut for ExecuteContext {
1672 fn deref_mut(&mut self) -> &mut Self::Target {
1673 self.inner.as_mut().expect("only consumed by value")
1674 }
1675}
1676
1677impl Drop for ExecuteContext {
1678 fn drop(&mut self) {
1679 let Some(inner) = self.inner.take() else {
1680 return;
1681 };
1682 tracing::warn!("execute context dropped without retirement, failing the client");
1685 let ExecuteContextInner { tx, session, .. } = *inner;
1686 tx.send(
1687 Err(AdapterError::Internal(
1688 "statement execution abandoned, outcome unknown (server shutting down)".into(),
1689 )),
1690 session,
1691 );
1692 }
1693}
1694
1695#[derive(Derivative)]
1696#[derivative(Debug)]
1697pub struct ExecuteContextInner {
1698 tx: ClientTransmitter<ExecuteResponse>,
1699 internal_cmd_tx: mpsc::UnboundedSender<Message>,
1700 session: Session,
1701 extra: ExecuteContextGuard,
1702 #[derivative(Debug = "ignore")]
1703 response_barriers: Vec<BuiltinTableAppendNotify>,
1704}
1705
1706impl ExecuteContext {
1707 pub fn session(&self) -> &Session {
1708 &self.session
1709 }
1710
1711 pub fn session_mut(&mut self) -> &mut Session {
1712 &mut self.session
1713 }
1714
1715 pub fn tx(&self) -> &ClientTransmitter<ExecuteResponse> {
1716 &self.tx
1717 }
1718
1719 pub fn tx_mut(&mut self) -> &mut ClientTransmitter<ExecuteResponse> {
1720 &mut self.tx
1721 }
1722
1723 pub fn from_parts(
1724 tx: ClientTransmitter<ExecuteResponse>,
1725 internal_cmd_tx: mpsc::UnboundedSender<Message>,
1726 session: Session,
1727 extra: ExecuteContextGuard,
1728 ) -> Self {
1729 Self::from_parts_with_response_barriers(tx, internal_cmd_tx, session, extra, Vec::new())
1730 }
1731
1732 pub fn from_parts_with_response_barriers(
1733 tx: ClientTransmitter<ExecuteResponse>,
1734 internal_cmd_tx: mpsc::UnboundedSender<Message>,
1735 session: Session,
1736 extra: ExecuteContextGuard,
1737 response_barriers: Vec<BuiltinTableAppendNotify>,
1738 ) -> Self {
1739 Self {
1740 inner: Some(
1741 ExecuteContextInner {
1742 tx,
1743 session,
1744 extra,
1745 response_barriers,
1746 internal_cmd_tx,
1747 }
1748 .into(),
1749 ),
1750 }
1751 }
1752
1753 pub fn into_parts(
1767 mut self,
1768 ) -> (
1769 ClientTransmitter<ExecuteResponse>,
1770 mpsc::UnboundedSender<Message>,
1771 Session,
1772 ExecuteContextGuard,
1773 Vec<BuiltinTableAppendNotify>,
1774 ) {
1775 let ExecuteContextInner {
1776 tx,
1777 internal_cmd_tx,
1778 session,
1779 extra,
1780 response_barriers,
1781 } = *self.inner.take().expect("only consumed by value");
1782 (tx, internal_cmd_tx, session, extra, response_barriers)
1783 }
1784
1785 #[instrument(level = "debug")]
1787 pub fn retire(mut self, result: Result<ExecuteResponse, AdapterError>) {
1788 let response_barriers = std::mem::take(&mut self.response_barriers);
1789 if response_barriers.is_empty() {
1790 let (tx, internal_cmd_tx, session, extra, _) = self.into_parts();
1791 retire_execution_context(tx, internal_cmd_tx, session, extra, result);
1792 return;
1793 }
1794 spawn(
1797 || "execute_context::retire_after_response_barriers",
1798 async move {
1799 for barrier in response_barriers {
1800 barrier.await;
1801 }
1802 self.retire(result);
1803 },
1804 );
1805 }
1806
1807 pub(crate) fn delay_response_until(&mut self, barrier: BuiltinTableAppendCompletion) {
1809 self.response_barriers.push(barrier.into_notify());
1810 }
1811
1812 pub fn extra(&self) -> &ExecuteContextGuard {
1813 &self.extra
1814 }
1815
1816 pub fn extra_mut(&mut self) -> &mut ExecuteContextGuard {
1817 &mut self.extra
1818 }
1819}
1820
1821fn retire_execution_context(
1822 tx: ClientTransmitter<ExecuteResponse>,
1823 internal_cmd_tx: mpsc::UnboundedSender<Message>,
1824 session: Session,
1825 extra: ExecuteContextGuard,
1826 result: Result<ExecuteResponse, AdapterError>,
1827) {
1828 let reason = if extra.is_trivial() {
1829 None
1830 } else {
1831 Some((&result).into())
1832 };
1833 tx.send(result, session);
1834 if let Some(reason) = reason {
1835 let extra = extra.defuse();
1836 if let Err(e) = internal_cmd_tx.send(Message::RetireExecute {
1837 otel_ctx: OpenTelemetryContext::obtain(),
1838 data: extra,
1839 reason,
1840 }) {
1841 warn!("internal_cmd_rx dropped before we could send: {:?}", e);
1842 }
1843 }
1844}
1845
1846#[derive(Debug)]
1847struct ClusterReplicaStatuses(
1848 BTreeMap<ClusterId, BTreeMap<ReplicaId, BTreeMap<ProcessId, ClusterReplicaProcessStatus>>>,
1849);
1850
1851impl ClusterReplicaStatuses {
1852 pub(crate) fn new() -> ClusterReplicaStatuses {
1853 ClusterReplicaStatuses(BTreeMap::new())
1854 }
1855
1856 pub(crate) fn initialize_cluster_statuses(&mut self, cluster_id: ClusterId) {
1860 let prev = self.0.insert(cluster_id, BTreeMap::new());
1861 assert_eq!(
1862 prev, None,
1863 "cluster {cluster_id} statuses already initialized"
1864 );
1865 }
1866
1867 pub(crate) fn initialize_cluster_replica_statuses(
1871 &mut self,
1872 cluster_id: ClusterId,
1873 replica_id: ReplicaId,
1874 num_processes: usize,
1875 time: DateTime<Utc>,
1876 ) {
1877 tracing::info!(
1878 ?cluster_id,
1879 ?replica_id,
1880 ?time,
1881 "initializing cluster replica status"
1882 );
1883 let replica_statuses = self.0.entry(cluster_id).or_default();
1884 let process_statuses = (0..num_processes)
1885 .map(|process_id| {
1886 let status = ClusterReplicaProcessStatus {
1887 status: ClusterStatus::Offline(Some(OfflineReason::Initializing)),
1888 restart_count: 0,
1889 time: time.clone(),
1890 };
1891 (u64::cast_from(process_id), status)
1892 })
1893 .collect();
1894 let prev = replica_statuses.insert(replica_id, process_statuses);
1895 assert_none!(
1896 prev,
1897 "cluster replica {cluster_id}.{replica_id} statuses already initialized"
1898 );
1899 }
1900
1901 pub(crate) fn remove_cluster_statuses(
1905 &mut self,
1906 cluster_id: &ClusterId,
1907 ) -> BTreeMap<ReplicaId, BTreeMap<ProcessId, ClusterReplicaProcessStatus>> {
1908 let prev = self.0.remove(cluster_id);
1909 prev.unwrap_or_else(|| panic!("unknown cluster: {cluster_id}"))
1910 }
1911
1912 pub(crate) fn remove_cluster_replica_statuses(
1916 &mut self,
1917 cluster_id: &ClusterId,
1918 replica_id: &ReplicaId,
1919 ) -> BTreeMap<ProcessId, ClusterReplicaProcessStatus> {
1920 let replica_statuses = self
1921 .0
1922 .get_mut(cluster_id)
1923 .unwrap_or_else(|| panic!("unknown cluster: {cluster_id}"));
1924 let prev = replica_statuses.remove(replica_id);
1925 prev.unwrap_or_else(|| panic!("unknown cluster replica: {cluster_id}.{replica_id}"))
1926 }
1927
1928 pub(crate) fn ensure_cluster_status(
1932 &mut self,
1933 cluster_id: ClusterId,
1934 replica_id: ReplicaId,
1935 process_id: ProcessId,
1936 status: ClusterReplicaProcessStatus,
1937 ) {
1938 let replica_statuses = self
1939 .0
1940 .get_mut(&cluster_id)
1941 .unwrap_or_else(|| panic!("unknown cluster: {cluster_id}"))
1942 .get_mut(&replica_id)
1943 .unwrap_or_else(|| panic!("unknown cluster replica: {cluster_id}.{replica_id}"));
1944 replica_statuses.insert(process_id, status);
1945 }
1946
1947 pub fn get_cluster_replica_status(
1951 &self,
1952 cluster_id: ClusterId,
1953 replica_id: ReplicaId,
1954 ) -> ClusterStatus {
1955 let process_status = self.get_cluster_replica_statuses(cluster_id, replica_id);
1956 Self::cluster_replica_status(process_status)
1957 }
1958
1959 pub fn cluster_replica_status(
1961 process_status: &BTreeMap<ProcessId, ClusterReplicaProcessStatus>,
1962 ) -> ClusterStatus {
1963 process_status
1964 .values()
1965 .fold(ClusterStatus::Online, |s, p| match (s, p.status) {
1966 (ClusterStatus::Online, ClusterStatus::Online) => ClusterStatus::Online,
1967 (x, y) => {
1968 let reason_x = match x {
1969 ClusterStatus::Offline(reason) => reason,
1970 ClusterStatus::Online => None,
1971 };
1972 let reason_y = match y {
1973 ClusterStatus::Offline(reason) => reason,
1974 ClusterStatus::Online => None,
1975 };
1976 ClusterStatus::Offline(reason_x.or(reason_y))
1978 }
1979 })
1980 }
1981
1982 pub(crate) fn get_cluster_replica_statuses(
1986 &self,
1987 cluster_id: ClusterId,
1988 replica_id: ReplicaId,
1989 ) -> &BTreeMap<ProcessId, ClusterReplicaProcessStatus> {
1990 self.try_get_cluster_replica_statuses(cluster_id, replica_id)
1991 .unwrap_or_else(|| panic!("unknown cluster replica: {cluster_id}.{replica_id}"))
1992 }
1993
1994 pub(crate) fn try_get_cluster_replica_statuses(
1996 &self,
1997 cluster_id: ClusterId,
1998 replica_id: ReplicaId,
1999 ) -> Option<&BTreeMap<ProcessId, ClusterReplicaProcessStatus>> {
2000 self.try_get_cluster_statuses(cluster_id)
2001 .and_then(|statuses| statuses.get(&replica_id))
2002 }
2003
2004 pub(crate) fn try_get_cluster_statuses(
2006 &self,
2007 cluster_id: ClusterId,
2008 ) -> Option<&BTreeMap<ReplicaId, BTreeMap<ProcessId, ClusterReplicaProcessStatus>>> {
2009 self.0.get(&cluster_id)
2010 }
2011}
2012
2013#[derive(Derivative)]
2015#[derivative(Debug)]
2016pub struct Coordinator {
2017 #[derivative(Debug = "ignore")]
2019 controller: mz_controller::Controller,
2020 catalog: Arc<Catalog>,
2028
2029 persist_client: PersistClient,
2032
2033 internal_cmd_tx: mpsc::UnboundedSender<Message>,
2035 group_commit_tx: appends::GroupCommitNotifier,
2037 reconcile_now: Arc<Notify>,
2041 group_committer_tx: mpsc::UnboundedSender<appends::TableWriteCmd>,
2042
2043 strict_serializable_reads_tx: mpsc::UnboundedSender<(ConnectionId, PendingReadTxn)>,
2045
2046 linearize_reads_notify: Arc<Notify>,
2050
2051 global_timelines: BTreeMap<Timeline, TimelineState>,
2054
2055 transient_id_gen: Arc<TransientIdGen>,
2057 active_conns: BTreeMap<ConnectionId, ConnMeta>,
2060
2061 txn_read_holds: BTreeMap<ConnectionId, read_policy::ReadHolds>,
2065
2066 pending_peeks: BTreeMap<Uuid, PendingPeek>,
2070 client_pending_peeks: BTreeMap<ConnectionId, BTreeMap<Uuid, ClusterId>>,
2072
2073 pending_linearize_read_txns: BTreeMap<ConnectionId, PendingReadTxn>,
2075
2076 active_compute_sinks: BTreeMap<GlobalId, ActiveComputeSink>,
2078 active_webhooks: BTreeMap<CatalogItemId, WebhookAppenderInvalidator>,
2080 active_copies: BTreeMap<ConnectionId, ActiveCopyFrom>,
2083
2084 connection_cancel_watches: BTreeMap<ConnectionId, (watch::Sender<bool>, watch::Receiver<bool>)>,
2095 introspection_subscribes: BTreeMap<GlobalId, IntrospectionSubscribe>,
2097
2098 write_locks: BTreeMap<CatalogItemId, Arc<tokio::sync::Mutex<()>>>,
2100 deferred_write_ops: BTreeMap<ConnectionId, DeferredOp>,
2102
2103 pending_writes: Vec<PendingWriteTxn>,
2105
2106 occ_write_semaphore: Arc<Semaphore>,
2117
2118 frontend_read_then_write_enabled: bool,
2123
2124 advance_timelines_interval: Interval,
2134
2135 serialized_ddl: LockedVecDeque<DeferredPlanStatement>,
2144
2145 secrets_controller: Arc<dyn SecretsController>,
2148 caching_secrets_reader: CachingSecretsReader,
2150
2151 cloud_resource_controller: Option<Arc<dyn CloudResourceController>>,
2154
2155 storage_usage_client: StorageUsageClient,
2157 storage_usage_collection_interval: Duration,
2159
2160 #[derivative(Debug = "ignore")]
2162 segment_client: Option<mz_segment::Client>,
2163
2164 metrics: Metrics,
2166 optimizer_metrics: OptimizerMetrics,
2168
2169 tracing_handle: TracingHandle,
2171
2172 statement_logging: StatementLogging,
2174
2175 webhook_concurrency_limit: WebhookConcurrencyLimiter,
2177
2178 timestamp_oracle_config: Option<TimestampOracleConfig>,
2181
2182 caught_up_check_interval: Interval,
2185
2186 caught_up_check: Option<CaughtUpCheckContext>,
2189
2190 catalog_info_metrics_registry: MetricsRegistry,
2193
2194 scoped_frontend: Option<Arc<SystemParameterFrontend>>,
2203
2204 installed_watch_sets: BTreeMap<WatchSetId, (ConnectionId, WatchSetResponse)>,
2206
2207 connection_watch_sets: BTreeMap<ConnectionId, BTreeSet<WatchSetId>>,
2209
2210 cluster_replica_statuses: ClusterReplicaStatuses,
2212
2213 read_only_controllers: bool,
2217
2218 buffered_builtin_table_updates: Option<Vec<BuiltinTableUpdate>>,
2226
2227 license_key: ValidatedLicenseKey,
2228
2229 user_id_pool: IdPool,
2231}
2232
2233impl Coordinator {
2234 pub(crate) async fn reconcile_scoped_system_parameters(
2253 &mut self,
2254 scoped: ScopedParameters,
2255 prune_scope: ScopedParametersScope,
2256 ) {
2257 if self.catalog().state().scoped_system_parameters() == &scoped {
2260 return;
2261 }
2262
2263 if let Err(e) = self
2271 .catalog_transact(
2272 None,
2273 vec![crate::catalog::Op::UpdateScopedSystemParameters {
2274 scoped,
2275 prune_scope,
2276 }],
2277 )
2278 .await
2279 {
2280 tracing::warn!("failed to persist scoped system parameters: {e}");
2281 }
2282 }
2283
2284 fn scoped_overrides_create_op(
2304 &self,
2305 clusters: &[ClusterEvalContext],
2306 replicas: &[ReplicaEvalContext],
2307 ) -> Option<crate::catalog::Op> {
2308 let frontend = self.scoped_frontend.clone()?;
2309 let catalog = self.catalog();
2310 let system_config = catalog.system_config();
2311
2312 let replica_param_names: Vec<&'static str> = system_config
2315 .iter_synced()
2316 .filter(|var| var.scope() == ParameterScope::Replica)
2317 .map(|var| var.name())
2318 .collect();
2319 let cluster_param_names: Vec<&'static str> = system_config
2320 .iter_synced()
2321 .filter(|var| var.scope() == ParameterScope::Cluster)
2322 .map(|var| var.name())
2323 .collect();
2324
2325 let params = SynchronizedParameters::new(system_config.clone());
2326 let mut evaluated = ScopedParameters::default();
2327 if !cluster_param_names.is_empty() && !clusters.is_empty() {
2328 evaluated.cluster =
2329 frontend.pull_cluster_overrides(¶ms, &cluster_param_names, clusters);
2330 }
2331 if !replica_param_names.is_empty() && !replicas.is_empty() {
2332 evaluated.replica =
2333 frontend.pull_replica_overrides(¶ms, &replica_param_names, replicas);
2334 }
2335 if evaluated.is_empty() {
2336 return None;
2337 }
2338
2339 let prune_scope = ScopedParametersScope {
2343 clusters: clusters.iter().map(|cluster| cluster.cluster_id).collect(),
2344 replicas: replicas.iter().map(|replica| replica.replica_id).collect(),
2345 };
2346 Some(crate::catalog::Op::UpdateScopedSystemParameters {
2347 scoped: evaluated,
2348 prune_scope,
2349 })
2350 }
2351
2352 pub(crate) fn push_replica_dyncfg_overrides(&mut self) {
2358 let replica_overrides = self
2361 .catalog()
2362 .state()
2363 .scoped_system_parameters()
2364 .replica
2365 .clone();
2366
2367 let dyncfgs = self.catalog().system_config().dyncfgs();
2368 let mut instance_overrides: BTreeMap<
2369 ComputeInstanceId,
2370 BTreeMap<ReplicaId, ConfigUpdates>,
2371 > = BTreeMap::new();
2372 for cluster in self.catalog().clusters() {
2373 for replica in cluster.replicas() {
2374 let Some(values) = replica_overrides.get(&replica.replica_id) else {
2375 continue;
2376 };
2377 let mut updates = ConfigUpdates::default();
2378 for (name, value) in values {
2379 let Some(entry) = dyncfgs.entry(name) else {
2380 continue;
2383 };
2384 match entry.parse_val(value) {
2385 Ok(val) => updates.add_dynamic(name, val),
2386 Err(e) => {
2387 tracing::warn!(%name, %value, "cannot parse scoped override: {e}")
2388 }
2389 }
2390 }
2391 if !updates.updates.is_empty() {
2392 instance_overrides
2393 .entry(cluster.id)
2394 .or_default()
2395 .insert(replica.replica_id, updates);
2396 }
2397 }
2398 }
2399
2400 self.controller
2412 .compute
2413 .update_replica_dyncfg_overrides(instance_overrides.clone());
2414 self.controller
2415 .storage
2416 .update_replica_dyncfg_overrides(instance_overrides);
2417 let compute_config = crate::flags::compute_config(self.catalog().system_config());
2423 self.controller.compute.update_configuration(compute_config);
2424 let storage_config = crate::flags::storage_config(self.catalog().system_config());
2425 self.controller.storage.update_parameters(storage_config);
2426 }
2427
2428 pub(crate) fn cluster_scoped_optimizer_overrides(
2432 &self,
2433 cluster_id: ClusterId,
2434 ) -> OptimizerFeatureOverrides {
2435 self.catalog()
2436 .state()
2437 .cluster_scoped_optimizer_overrides(cluster_id)
2438 }
2439
2440 #[instrument(name = "coord::bootstrap")]
2444 pub(crate) async fn bootstrap(
2445 &mut self,
2446 boot_ts: Timestamp,
2447 migrated_storage_collections_0dt: BTreeSet<CatalogItemId>,
2448 mut builtin_table_updates: Vec<BuiltinTableUpdate>,
2449 cached_global_exprs: BTreeMap<GlobalId, GlobalExpressions>,
2450 uncached_local_exprs: BTreeMap<GlobalId, LocalExpressions>,
2451 ) -> Result<(), AdapterError> {
2452 let bootstrap_start = Instant::now();
2453 info!("startup: coordinator init: bootstrap beginning");
2454 info!("startup: coordinator init: bootstrap: preamble beginning");
2455
2456 let cluster_statuses: Vec<(_, Vec<_>)> = self
2459 .catalog()
2460 .clusters()
2461 .map(|cluster| {
2462 (
2463 cluster.id(),
2464 cluster
2465 .replicas()
2466 .map(|replica| {
2467 (replica.replica_id, replica.config.location.num_processes())
2468 })
2469 .collect(),
2470 )
2471 })
2472 .collect();
2473 let now = self.now_datetime();
2474 for (cluster_id, replica_statuses) in cluster_statuses {
2475 self.cluster_replica_statuses
2476 .initialize_cluster_statuses(cluster_id);
2477 for (replica_id, num_processes) in replica_statuses {
2478 self.cluster_replica_statuses
2479 .initialize_cluster_replica_statuses(
2480 cluster_id,
2481 replica_id,
2482 num_processes,
2483 now,
2484 );
2485 }
2486 }
2487
2488 let system_config = self.catalog().system_config();
2489
2490 mz_metrics::update_dyncfg(&system_config.dyncfg_updates());
2492
2493 let compute_config = flags::compute_config(system_config);
2495 let storage_config = flags::storage_config(system_config);
2496 let scheduling_config = flags::orchestrator_scheduling_config(system_config);
2497 let dyncfg_updates = system_config.dyncfg_updates();
2498 self.controller.compute.update_configuration(compute_config);
2499 self.controller.storage.update_parameters(storage_config);
2500 self.controller
2501 .update_orchestrator_scheduling_config(scheduling_config);
2502 self.controller.update_configuration(dyncfg_updates);
2503
2504 let enforce_credit_limit_at_bootstrap = !matches!(
2509 self.license_key.expiration_behavior,
2510 ExpirationBehavior::DisableClusterCreation,
2511 );
2512 if enforce_credit_limit_at_bootstrap {
2513 self.validate_resource_limit_numeric(
2514 Numeric::zero(),
2515 self.current_credit_consumption_rate(None),
2516 |system_vars| {
2517 self.license_key
2518 .max_credit_consumption_rate()
2519 .map_or_else(|| system_vars.max_credit_consumption_rate(), Numeric::from)
2520 },
2521 "cluster replica",
2522 MAX_CREDIT_CONSUMPTION_RATE.name(),
2523 )?;
2524 }
2525
2526 let mut policies_to_set: BTreeMap<CompactionWindow, CollectionIdBundle> =
2527 Default::default();
2528
2529 let enable_worker_core_affinity =
2530 self.catalog().system_config().enable_worker_core_affinity();
2531 let enable_storage_introspection_logs = self
2532 .catalog()
2533 .system_config()
2534 .enable_storage_introspection_logs();
2535 for instance in self.catalog.clusters() {
2536 self.controller.create_cluster(
2537 instance.id,
2538 ClusterConfig {
2539 arranged_logs: instance.log_indexes.clone(),
2540 workload_class: instance.config.workload_class.clone(),
2541 },
2542 )?;
2543 for replica in instance.replicas() {
2544 let role = instance.role();
2545 self.controller.create_replica(
2546 instance.id,
2547 replica.replica_id,
2548 instance.name.clone(),
2549 replica.name.clone(),
2550 role,
2551 replica.config.clone(),
2552 enable_worker_core_affinity,
2553 enable_storage_introspection_logs,
2554 )?;
2555 }
2556 }
2557
2558 self.push_replica_dyncfg_overrides();
2570
2571 info!(
2572 "startup: coordinator init: bootstrap: preamble complete in {:?}",
2573 bootstrap_start.elapsed()
2574 );
2575
2576 let init_storage_collections_start = Instant::now();
2577 info!("startup: coordinator init: bootstrap: storage collections init beginning");
2578 self.bootstrap_storage_collections(&migrated_storage_collections_0dt)
2579 .await;
2580 info!(
2581 "startup: coordinator init: bootstrap: storage collections init complete in {:?}",
2582 init_storage_collections_start.elapsed()
2583 );
2584
2585 self.controller.start_compute_introspection_sink();
2590
2591 let sorting_start = Instant::now();
2592 info!("startup: coordinator init: bootstrap: sorting catalog entries");
2593 let entries = self.bootstrap_sort_catalog_entries();
2594 info!(
2595 "startup: coordinator init: bootstrap: sorting catalog entries complete in {:?}",
2596 sorting_start.elapsed()
2597 );
2598
2599 let optimize_dataflows_start = Instant::now();
2600 info!("startup: coordinator init: bootstrap: optimize dataflow plans beginning");
2601 let uncached_global_exps = self.bootstrap_dataflow_plans(&entries, cached_global_exprs)?;
2602 info!(
2603 "startup: coordinator init: bootstrap: optimize dataflow plans complete in {:?}",
2604 optimize_dataflows_start.elapsed()
2605 );
2606
2607 let _fut = self.catalog().update_expression_cache(
2609 uncached_local_exprs.into_iter().collect(),
2610 uncached_global_exps.into_iter().collect(),
2611 Default::default(),
2612 );
2613
2614 let bootstrap_as_ofs_start = Instant::now();
2618 info!("startup: coordinator init: bootstrap: dataflow as-of bootstrapping beginning");
2619 let dataflow_read_holds = self.bootstrap_dataflow_as_ofs().await;
2620 info!(
2621 "startup: coordinator init: bootstrap: dataflow as-of bootstrapping complete in {:?}",
2622 bootstrap_as_ofs_start.elapsed()
2623 );
2624
2625 let postamble_start = Instant::now();
2626 info!("startup: coordinator init: bootstrap: postamble beginning");
2627
2628 let logs: BTreeSet<_> = BUILTINS::logs()
2629 .map(|log| self.catalog().resolve_builtin_log(log))
2630 .flat_map(|item_id| self.catalog().get_global_ids(&item_id))
2631 .collect();
2632
2633 let mut privatelink_connections = BTreeMap::new();
2634
2635 for entry in &entries {
2636 debug!(
2637 "coordinator init: installing {} {}",
2638 entry.item().typ(),
2639 entry.id()
2640 );
2641 let mut policy = entry.item().initial_logical_compaction_window();
2642 match entry.item() {
2643 CatalogItem::Source(source) => {
2649 if source.custom_logical_compaction_window.is_none() {
2651 if let DataSourceDesc::IngestionExport { ingestion_id, .. } =
2652 source.data_source
2653 {
2654 policy = Some(
2655 self.catalog()
2656 .get_entry(&ingestion_id)
2657 .source()
2658 .expect("must be source")
2659 .custom_logical_compaction_window
2660 .unwrap_or_default(),
2661 );
2662 }
2663 }
2664 policies_to_set
2665 .entry(policy.expect("sources have a compaction window"))
2666 .or_insert_with(Default::default)
2667 .storage_ids
2668 .insert(source.global_id());
2669 }
2670 CatalogItem::Table(table) => {
2671 policies_to_set
2672 .entry(policy.expect("tables have a compaction window"))
2673 .or_insert_with(Default::default)
2674 .storage_ids
2675 .extend(table.global_ids());
2676 }
2677 CatalogItem::Index(idx) => {
2678 let policy_entry = policies_to_set
2679 .entry(policy.expect("indexes have a compaction window"))
2680 .or_insert_with(Default::default);
2681
2682 if logs.contains(&idx.on) {
2683 policy_entry
2684 .compute_ids
2685 .entry(idx.cluster_id)
2686 .or_insert_with(BTreeSet::new)
2687 .insert(idx.global_id());
2688 } else {
2689 let df_desc = self
2690 .catalog()
2691 .try_get_physical_plan(&idx.global_id())
2692 .expect("added in `bootstrap_dataflow_plans`")
2693 .clone();
2694
2695 let df_meta = self
2696 .catalog()
2697 .try_get_dataflow_metainfo(&idx.global_id())
2698 .expect("added in `bootstrap_dataflow_plans`");
2699
2700 if self.catalog().state().system_config().enable_mz_notices() {
2701 self.catalog().state().pack_optimizer_notices(
2703 &mut builtin_table_updates,
2704 df_meta.optimizer_notices.iter(),
2705 Diff::ONE,
2706 );
2707 }
2708
2709 policy_entry
2712 .compute_ids
2713 .entry(idx.cluster_id)
2714 .or_insert_with(Default::default)
2715 .extend(df_desc.export_ids());
2716
2717 self.controller
2718 .compute
2719 .create_dataflow(idx.cluster_id, df_desc, None)
2720 .unwrap_or_terminate("cannot fail to create dataflows");
2721 }
2722 }
2723 CatalogItem::View(_) => (),
2724 CatalogItem::MaterializedView(mview) => {
2725 policies_to_set
2731 .entry(policy.expect("materialized views have a compaction window"))
2732 .or_insert_with(Default::default)
2733 .storage_ids
2734 .extend(mview.global_ids());
2735
2736 let mut df_desc = self
2737 .catalog()
2738 .try_get_physical_plan(&mview.global_id_writes())
2739 .expect("added in `bootstrap_dataflow_plans`")
2740 .clone();
2741
2742 if let Some(initial_as_of) = mview.initial_as_of.clone() {
2743 df_desc.set_initial_as_of(initial_as_of);
2744 }
2745
2746 let until = mview
2748 .refresh_schedule
2749 .as_ref()
2750 .and_then(|s| s.last_refresh())
2751 .and_then(|r| r.try_step_forward());
2752 if let Some(until) = until {
2753 df_desc.until.meet_assign(&Antichain::from_elem(until));
2754 }
2755
2756 let df_meta = self
2757 .catalog()
2758 .try_get_dataflow_metainfo(&mview.global_id_writes())
2759 .expect("added in `bootstrap_dataflow_plans`");
2760
2761 if self.catalog().state().system_config().enable_mz_notices() {
2762 self.catalog().state().pack_optimizer_notices(
2764 &mut builtin_table_updates,
2765 df_meta.optimizer_notices.iter(),
2766 Diff::ONE,
2767 );
2768 }
2769
2770 self.ship_dataflow(df_desc, mview.cluster_id, mview.target_replica)
2771 .await;
2772
2773 if mview.replacement_target.is_none() {
2776 self.allow_writes(mview.cluster_id, mview.global_id_writes());
2777 }
2778 }
2779 CatalogItem::MetricSink(metric_sink) => {
2780 let df_desc = self
2781 .catalog()
2782 .try_get_physical_plan(&metric_sink.global_id)
2783 .expect("added in `bootstrap_dataflow_plans`")
2784 .clone();
2785
2786 let df_meta = self
2787 .catalog()
2788 .try_get_dataflow_metainfo(&metric_sink.global_id)
2789 .expect("added in `bootstrap_dataflow_plans`");
2790
2791 if self.catalog().state().system_config().enable_mz_notices() {
2792 self.catalog().state().pack_optimizer_notices(
2794 &mut builtin_table_updates,
2795 df_meta.optimizer_notices.iter(),
2796 Diff::ONE,
2797 );
2798 }
2799
2800 self.ship_dataflow(df_desc, metric_sink.cluster_id, None)
2803 .await;
2804 }
2805 CatalogItem::Sink(sink) => {
2806 policies_to_set
2807 .entry(CompactionWindow::Default)
2808 .or_insert_with(Default::default)
2809 .storage_ids
2810 .insert(sink.global_id());
2811 }
2812 CatalogItem::Connection(catalog_connection) => {
2813 if let ConnectionDetails::AwsPrivatelink(conn) = &catalog_connection.details {
2814 privatelink_connections.insert(
2815 entry.id(),
2816 VpcEndpointConfig {
2817 aws_service_name: conn.service_name.clone(),
2818 availability_zone_ids: conn.availability_zones.clone(),
2819 },
2820 );
2821 }
2822 }
2823 CatalogItem::Log(_)
2825 | CatalogItem::Type(_)
2826 | CatalogItem::Func(_)
2827 | CatalogItem::Secret(_) => {}
2828 }
2829 }
2830
2831 if let Some(cloud_resource_controller) = &self.cloud_resource_controller {
2832 let existing_vpc_endpoints = cloud_resource_controller
2834 .list_vpc_endpoints()
2835 .await
2836 .context("list vpc endpoints")?;
2837 let existing_vpc_endpoints = BTreeSet::from_iter(existing_vpc_endpoints.into_keys());
2838 let desired_vpc_endpoints = privatelink_connections.keys().cloned().collect();
2839 let vpc_endpoints_to_remove = existing_vpc_endpoints.difference(&desired_vpc_endpoints);
2840 for id in vpc_endpoints_to_remove {
2841 cloud_resource_controller
2842 .delete_vpc_endpoint(*id)
2843 .await
2844 .context("deleting extraneous vpc endpoint")?;
2845 }
2846
2847 for (id, spec) in privatelink_connections {
2849 cloud_resource_controller
2850 .ensure_vpc_endpoint(id, spec)
2851 .await
2852 .context("ensuring vpc endpoint")?;
2853 }
2854 }
2855
2856 drop(dataflow_read_holds);
2859 for (cw, policies) in policies_to_set {
2861 self.initialize_read_policies(&policies, cw).await;
2862 }
2863
2864 builtin_table_updates.extend(
2866 self.catalog().state().resolve_builtin_table_updates(
2867 self.catalog().state().pack_all_replica_size_updates(),
2868 ),
2869 );
2870
2871 debug!("startup: coordinator init: bootstrap: initializing migrated builtin tables");
2872 let migrated_updates_fut = if self.controller.read_only() {
2878 let min_timestamp = Timestamp::minimum();
2879 let migrated_builtin_table_updates: Vec<_> = builtin_table_updates
2880 .extract_if(.., |update| {
2881 let gid = self.catalog().get_entry(&update.id).latest_global_id();
2882 migrated_storage_collections_0dt.contains(&update.id)
2883 && self
2884 .controller
2885 .storage_collections
2886 .collection_frontiers(gid)
2887 .expect("all tables are registered")
2888 .write_frontier
2889 .elements()
2890 == &[min_timestamp]
2891 })
2892 .collect();
2893 if migrated_builtin_table_updates.is_empty() {
2894 futures::future::ready(()).boxed()
2895 } else {
2896 let mut grouped_appends: BTreeMap<GlobalId, Vec<TableData>> = BTreeMap::new();
2898 for update in migrated_builtin_table_updates {
2899 let gid = self.catalog().get_entry(&update.id).latest_global_id();
2900 grouped_appends.entry(gid).or_default().push(update.data);
2901 }
2902 info!(
2903 "coordinator init: rehydrating migrated builtin tables in read-only mode: {:?}",
2904 grouped_appends.keys().collect::<Vec<_>>()
2905 );
2906
2907 let mut all_appends = Vec::with_capacity(grouped_appends.len());
2909 for (item_id, table_data) in grouped_appends.into_iter() {
2910 let mut all_rows = Vec::new();
2911 let mut all_data = Vec::new();
2912 for data in table_data {
2913 match data {
2914 TableData::Rows(rows) => all_rows.extend(rows),
2915 TableData::Batches(_) => all_data.push(data),
2916 }
2917 }
2918 differential_dataflow::consolidation::consolidate(&mut all_rows);
2919 all_data.push(TableData::Rows(all_rows));
2920
2921 all_appends.push((item_id, all_data));
2923 }
2924
2925 let fut = self
2926 .controller
2927 .storage
2928 .append_table(min_timestamp, boot_ts.step_forward(), all_appends)
2929 .expect("cannot fail to append");
2930 async {
2931 fut.await
2932 .expect("One-shot shouldn't be dropped during bootstrap")
2933 .unwrap_or_terminate("cannot fail to append")
2934 }
2935 .boxed()
2936 }
2937 } else {
2938 futures::future::ready(()).boxed()
2939 };
2940
2941 info!(
2942 "startup: coordinator init: bootstrap: postamble complete in {:?}",
2943 postamble_start.elapsed()
2944 );
2945
2946 let builtin_update_start = Instant::now();
2947 info!("startup: coordinator init: bootstrap: generate builtin updates beginning");
2948
2949 if self.controller.read_only() {
2950 info!(
2951 "coordinator init: bootstrap: stashing builtin table updates while in read-only mode"
2952 );
2953
2954 self.buffered_builtin_table_updates
2955 .as_mut()
2956 .expect("in read-only mode")
2957 .append(&mut builtin_table_updates);
2958 } else {
2959 self.bootstrap_tables(&entries, builtin_table_updates).await;
2960 };
2961 info!(
2962 "startup: coordinator init: bootstrap: generate builtin updates complete in {:?}",
2963 builtin_update_start.elapsed()
2964 );
2965
2966 let cleanup_secrets_start = Instant::now();
2967 info!("startup: coordinator init: bootstrap: generate secret cleanup beginning");
2968 {
2972 let Self {
2975 secrets_controller,
2976 catalog,
2977 ..
2978 } = self;
2979
2980 let next_user_item_id = catalog.get_next_user_item_id().await?;
2981 let next_system_item_id = catalog.get_next_system_item_id().await?;
2982 let read_only = self.controller.read_only();
2983 let catalog_ids: BTreeSet<CatalogItemId> =
2988 catalog.entries().map(|entry| entry.id()).collect();
2989 let secrets_controller = Arc::clone(secrets_controller);
2990
2991 spawn(|| "cleanup-orphaned-secrets", async move {
2992 if read_only {
2993 info!(
2994 "coordinator init: not cleaning up orphaned secrets while in read-only mode"
2995 );
2996 return;
2997 }
2998 info!("coordinator init: cleaning up orphaned secrets");
2999
3000 match secrets_controller.list().await {
3001 Ok(controller_secrets) => {
3002 let controller_secrets: BTreeSet<CatalogItemId> =
3003 controller_secrets.into_iter().collect();
3004 let orphaned = controller_secrets.difference(&catalog_ids);
3005 for id in orphaned {
3006 let id_too_large = match id {
3007 CatalogItemId::System(id) => *id >= next_system_item_id,
3008 CatalogItemId::User(id) => *id >= next_user_item_id,
3009 CatalogItemId::IntrospectionSourceIndex(_)
3010 | CatalogItemId::Transient(_) => false,
3011 };
3012 if id_too_large {
3013 info!(
3014 %next_user_item_id, %next_system_item_id,
3015 "coordinator init: not deleting orphaned secret {id} that was likely created by a newer deploy generation"
3016 );
3017 } else {
3018 info!("coordinator init: deleting orphaned secret {id}");
3019 fail_point!("orphan_secrets");
3020 if let Err(e) = secrets_controller.delete(*id).await {
3021 warn!(
3022 "Dropping orphaned secret has encountered an error: {}",
3023 e
3024 );
3025 }
3026 }
3027 }
3028 }
3029 Err(e) => warn!("Failed to list secrets during orphan cleanup: {:?}", e),
3030 }
3031 });
3032 }
3033 info!(
3034 "startup: coordinator init: bootstrap: generate secret cleanup complete in {:?}",
3035 cleanup_secrets_start.elapsed()
3036 );
3037
3038 let final_steps_start = Instant::now();
3040 info!(
3041 "startup: coordinator init: bootstrap: migrate builtin tables in read-only mode beginning"
3042 );
3043 migrated_updates_fut
3044 .instrument(info_span!("coord::bootstrap::final"))
3045 .await;
3046
3047 debug!(
3048 "startup: coordinator init: bootstrap: announcing completion of initialization to controller"
3049 );
3050 self.controller.initialization_complete();
3052
3053 self.bootstrap_introspection_subscribes().await;
3055
3056 info!(
3057 "startup: coordinator init: bootstrap: migrate builtin tables in read-only mode complete in {:?}",
3058 final_steps_start.elapsed()
3059 );
3060
3061 info!(
3062 "startup: coordinator init: bootstrap complete in {:?}",
3063 bootstrap_start.elapsed()
3064 );
3065 Ok(())
3066 }
3067
3068 #[allow(clippy::async_yields_async)]
3073 #[instrument]
3074 async fn bootstrap_tables(
3075 &mut self,
3076 entries: &[CatalogEntry],
3077 mut builtin_table_updates: Vec<BuiltinTableUpdate>,
3078 ) {
3079 struct TableMetadata<'a> {
3081 id: CatalogItemId,
3082 name: &'a QualifiedItemName,
3083 table: &'a Table,
3084 }
3085
3086 let table_metas: Vec<_> = entries
3088 .into_iter()
3089 .filter_map(|entry| {
3090 entry.table().map(|table| TableMetadata {
3091 id: entry.id(),
3092 name: entry.name(),
3093 table,
3094 })
3095 })
3096 .collect();
3097
3098 debug!("coordinator init: advancing all tables to current timestamp");
3100 let WriteTimestamp {
3101 timestamp: write_ts,
3102 advance_to,
3103 } = self.get_local_write_ts().await;
3104 let appends = table_metas
3105 .iter()
3106 .map(|meta| (meta.table.global_id_writes(), Vec::new()))
3107 .collect();
3108 let table_fence_rx = self
3112 .controller
3113 .storage
3114 .append_table(write_ts.clone(), advance_to, appends)
3115 .expect("invalid updates");
3116
3117 self.apply_local_write(write_ts).await;
3118
3119 debug!("coordinator init: resetting system tables");
3121 let read_ts = self.get_local_read_ts().await;
3122
3123 let mz_storage_usage_by_shard_schema: SchemaSpecifier = self
3128 .catalog()
3129 .resolve_system_schema(MZ_STORAGE_USAGE_BY_SHARD.schema)
3130 .into();
3131 let arrangement_size_history_schema: SchemaSpecifier = self
3132 .catalog()
3133 .resolve_system_schema(MZ_OBJECT_ARRANGEMENT_SIZE_HISTORY.schema)
3134 .into();
3135 let is_retained_across_restarts = |meta: &TableMetadata| -> bool {
3136 (meta.name.item == MZ_STORAGE_USAGE_BY_SHARD.name
3137 && meta.name.qualifiers.schema_spec == mz_storage_usage_by_shard_schema)
3138 || (meta.name.item == MZ_OBJECT_ARRANGEMENT_SIZE_HISTORY.name
3139 && meta.name.qualifiers.schema_spec == arrangement_size_history_schema)
3140 };
3141
3142 let mut retraction_tasks = Vec::new();
3143 let system_tables: Vec<_> = table_metas
3144 .iter()
3145 .filter(|meta| meta.id.is_system() && !is_retained_across_restarts(meta))
3146 .collect();
3147
3148 for system_table in system_tables {
3149 let table_id = system_table.id;
3150 let full_name = self.catalog().resolve_full_name(system_table.name, None);
3151 debug!("coordinator init: resetting system table {full_name} ({table_id})");
3152
3153 let snapshot_fut = self
3155 .controller
3156 .storage_collections
3157 .snapshot_cursor(system_table.table.global_id_writes(), read_ts);
3158 let batch_fut = self
3159 .controller
3160 .storage_collections
3161 .create_update_builder(system_table.table.global_id_writes());
3162
3163 let task = spawn(|| format!("snapshot-{table_id}"), async move {
3164 let mut batch = batch_fut
3166 .await
3167 .unwrap_or_terminate("cannot fail to create a batch for a BuiltinTable");
3168 tracing::info!(?table_id, "starting snapshot");
3169 let mut snapshot_cursor = snapshot_fut
3171 .await
3172 .unwrap_or_terminate("cannot fail to snapshot");
3173
3174 while let Some(values) = snapshot_cursor.next().await {
3176 for (key, _t, d) in values {
3177 let d_invert = d.neg();
3178 batch.add(&key, &(), &d_invert).await;
3179 }
3180 }
3181 tracing::info!(?table_id, "finished snapshot");
3182
3183 let batch = batch.finish().await;
3184 BuiltinTableUpdate::batch(table_id, batch)
3185 });
3186 retraction_tasks.push(task);
3187 }
3188
3189 let retractions_res = futures::future::join_all(retraction_tasks).await;
3190 for retractions in retractions_res {
3191 builtin_table_updates.push(retractions);
3192 }
3193
3194 table_fence_rx
3196 .await
3197 .expect("One-shot shouldn't be dropped during bootstrap")
3198 .unwrap_or_terminate("cannot fail to append");
3199
3200 info!("coordinator init: sending builtin table updates");
3201 let builtin_updates_fut = self.builtin_table_update().execute(builtin_table_updates);
3202 builtin_updates_fut.await;
3205 }
3206
3207 #[instrument]
3220 async fn bootstrap_storage_collections(
3221 &mut self,
3222 migrated_storage_collections: &BTreeSet<CatalogItemId>,
3223 ) {
3224 let catalog = self.catalog();
3225
3226 let source_desc = |object_id: GlobalId,
3227 data_source: &DataSourceDesc,
3228 desc: &RelationDesc,
3229 timeline: &Timeline| {
3230 let data_source = match data_source.clone() {
3231 DataSourceDesc::Ingestion { desc, cluster_id } => {
3233 let desc = desc.into_inline_connection(catalog.state());
3234 let ingestion = IngestionDescription::new(desc, cluster_id, object_id);
3235 DataSource::Ingestion(ingestion)
3236 }
3237 DataSourceDesc::OldSyntaxIngestion {
3238 desc,
3239 progress_subsource,
3240 data_config,
3241 details,
3242 cluster_id,
3243 } => {
3244 let desc = desc.into_inline_connection(catalog.state());
3245 let data_config = data_config.into_inline_connection(catalog.state());
3246 let progress_subsource =
3249 catalog.get_entry(&progress_subsource).latest_global_id();
3250 let mut ingestion =
3251 IngestionDescription::new(desc, cluster_id, progress_subsource);
3252 let legacy_export = SourceExport {
3253 storage_metadata: (),
3254 data_config,
3255 details,
3256 };
3257 ingestion.source_exports.insert(object_id, legacy_export);
3258
3259 DataSource::Ingestion(ingestion)
3260 }
3261 DataSourceDesc::IngestionExport {
3262 ingestion_id,
3263 external_reference: _,
3264 details,
3265 data_config,
3266 } => {
3267 let ingestion_id = catalog.get_entry(&ingestion_id).latest_global_id();
3270
3271 DataSource::IngestionExport {
3272 ingestion_id,
3273 details,
3274 data_config: data_config.into_inline_connection(catalog.state()),
3275 }
3276 }
3277 DataSourceDesc::Webhook { .. } => DataSource::Webhook,
3278 DataSourceDesc::Progress => DataSource::Progress,
3279 DataSourceDesc::Introspection(introspection) => {
3280 DataSource::Introspection(introspection)
3281 }
3282 DataSourceDesc::Catalog => DataSource::Other,
3283 };
3284 CollectionDescription {
3285 desc: desc.clone(),
3286 data_source,
3287 since: None,
3288 timeline: Some(timeline.clone()),
3289 primary: None,
3290 }
3291 };
3292
3293 let mut compute_collections = vec![];
3294 let mut collections = vec![];
3295 for entry in catalog.entries() {
3296 match entry.item() {
3297 CatalogItem::Source(source) => {
3298 collections.push((
3299 source.global_id(),
3300 source_desc(
3301 source.global_id(),
3302 &source.data_source,
3303 &source.desc,
3304 &source.timeline,
3305 ),
3306 ));
3307 }
3308 CatalogItem::Table(table) => {
3309 match &table.data_source {
3310 TableDataSource::TableWrites { defaults: _ } => {
3311 let versions: BTreeMap<_, _> = table
3312 .collection_descs()
3313 .map(|(gid, version, desc)| (version, (gid, desc)))
3314 .collect();
3315 let collection_descs = versions.iter().map(|(version, (gid, desc))| {
3316 let next_version = version.bump();
3317 let primary_collection =
3318 versions.get(&next_version).map(|(gid, _desc)| gid).copied();
3319 let mut collection_desc =
3320 CollectionDescription::for_table(desc.clone());
3321 collection_desc.primary = primary_collection;
3322
3323 (*gid, collection_desc)
3324 });
3325 collections.extend(collection_descs);
3326 }
3327 TableDataSource::DataSource {
3328 desc: data_source_desc,
3329 timeline,
3330 } => {
3331 soft_assert_eq_or_log!(table.collections.len(), 1);
3333 let collection_descs =
3334 table.collection_descs().map(|(gid, _version, desc)| {
3335 (
3336 gid,
3337 source_desc(
3338 entry.latest_global_id(),
3339 data_source_desc,
3340 &desc,
3341 timeline,
3342 ),
3343 )
3344 });
3345 collections.extend(collection_descs);
3346 }
3347 };
3348 }
3349 CatalogItem::MaterializedView(mv) => {
3350 let mut primary = mv
3358 .replacement_target
3359 .map(|target_id| catalog.get_entry(&target_id).latest_global_id());
3360 let collection_descs = mv.collection_descs().map(|(gid, _version, desc)| {
3361 let mut collection_desc =
3362 CollectionDescription::for_other(desc, mv.initial_as_of.clone());
3363 collection_desc.primary = primary;
3364 primary = Some(gid);
3365 (gid, collection_desc)
3366 });
3367
3368 collections.extend(collection_descs);
3369 compute_collections.push((mv.global_id_writes(), mv.desc.latest()));
3370 }
3371 CatalogItem::Sink(sink) => {
3372 let storage_sink_from_entry = self.catalog().get_entry_by_global_id(&sink.from);
3373 let from_desc = storage_sink_from_entry
3374 .relation_desc()
3375 .expect("sinks can only be built on items with descs")
3376 .into_owned();
3377 let collection_desc = CollectionDescription {
3378 desc: KAFKA_PROGRESS_DESC.clone(),
3380 data_source: DataSource::Sink {
3381 desc: ExportDescription {
3382 sink: StorageSinkDesc {
3383 from: sink.from,
3384 from_desc,
3385 connection: sink
3386 .connection
3387 .clone()
3388 .into_inline_connection(self.catalog().state()),
3389 envelope: sink.envelope,
3390 as_of: Antichain::from_elem(Timestamp::minimum()),
3391 with_snapshot: sink.with_snapshot,
3392 version: sink.version,
3393 from_storage_metadata: (),
3394 to_storage_metadata: (),
3395 commit_interval: sink.commit_interval,
3396 },
3397 instance_id: sink.cluster_id,
3398 },
3399 },
3400 since: None,
3401 timeline: None,
3402 primary: None,
3403 };
3404 collections.push((sink.global_id, collection_desc));
3405 }
3406 CatalogItem::Log(_)
3407 | CatalogItem::View(_)
3408 | CatalogItem::Index(_)
3409 | CatalogItem::Type(_)
3410 | CatalogItem::Func(_)
3411 | CatalogItem::Secret(_)
3412 | CatalogItem::Connection(_)
3413 | CatalogItem::MetricSink(_) => (),
3416 }
3417 }
3418
3419 let register_ts = if self.controller.read_only() {
3420 self.get_local_read_ts().await
3421 } else {
3422 self.get_local_write_ts().await.timestamp
3425 };
3426
3427 let storage_metadata = self.catalog.state().storage_metadata();
3428 let migrated_storage_collections = migrated_storage_collections
3429 .into_iter()
3430 .flat_map(|item_id| self.catalog.get_entry(item_id).global_ids())
3431 .collect();
3432
3433 self.controller
3438 .storage
3439 .evolve_nullability_for_bootstrap(storage_metadata, compute_collections)
3440 .await
3441 .unwrap_or_terminate("cannot fail to evolve collections");
3442
3443 let mut pending: BTreeMap<_, _> = collections.into_iter().collect();
3456
3457 let transitive_dep_gids: BTreeMap<_, _> = pending
3459 .keys()
3460 .map(|gid| {
3461 let entry = self.catalog.get_entry_by_global_id(gid);
3462 let item_id = entry.id();
3463 let deps = self.catalog.state().transitive_uses(item_id);
3464 let dep_gids: BTreeSet<_> = deps
3465 .filter(|dep_id| *dep_id != item_id)
3468 .map(|dep_id| self.catalog.get_entry(&dep_id).latest_global_id())
3469 .filter(|dep_gid| pending.contains_key(dep_gid))
3471 .collect();
3472 (*gid, dep_gids)
3473 })
3474 .collect();
3475
3476 let mut created_gids = Vec::new();
3477
3478 while !pending.is_empty() {
3479 let ready_gids: BTreeSet<_> = pending
3482 .keys()
3483 .filter(|gid| {
3484 let mut deps = transitive_dep_gids[gid].iter();
3485 !deps.any(|dep_gid| pending.contains_key(dep_gid))
3486 })
3487 .copied()
3488 .collect();
3489 let mut ready: Vec<_> = pending
3490 .extract_if(.., |gid, _| ready_gids.contains(gid))
3491 .collect();
3492
3493 for (gid, collection) in &mut ready {
3495 if !gid.is_system() || collection.since.is_some() {
3497 continue;
3498 }
3499
3500 let mut derived_since = Antichain::from_elem(Timestamp::MIN);
3501 for dep_gid in &transitive_dep_gids[gid] {
3502 let (since, _) = self
3503 .controller
3504 .storage
3505 .collection_frontiers(*dep_gid)
3506 .expect("previously registered");
3507 derived_since.join_assign(&since);
3508 }
3509 collection.since = Some(derived_since);
3510 }
3511
3512 if ready.is_empty() {
3513 soft_panic_or_log!(
3514 "cycle in storage collections: {:?}",
3515 pending.keys().collect::<Vec<_>>(),
3516 );
3517 ready = mem::take(&mut pending).into_iter().collect();
3521 }
3522
3523 created_gids.extend(ready.iter().map(|(gid, _collection)| *gid));
3524
3525 self.controller
3526 .storage
3527 .create_collections_for_bootstrap(
3528 storage_metadata,
3529 Some(register_ts),
3530 ready,
3531 &migrated_storage_collections,
3532 )
3533 .await
3534 .unwrap_or_terminate("cannot fail to create collections");
3535 }
3536
3537 self.controller
3539 .storage
3540 .register_table_collections(register_ts, created_gids)
3541 .await
3542 .unwrap_or_terminate("cannot fail to register tables");
3543
3544 if !self.controller.read_only() {
3545 self.apply_local_write(register_ts).await;
3546 }
3547 }
3548
3549 fn bootstrap_sort_catalog_entries(&self) -> Vec<CatalogEntry> {
3556 let mut indexes_on = BTreeMap::<_, Vec<_>>::new();
3557 let mut non_indexes = Vec::new();
3558 for entry in self.catalog().entries().cloned() {
3559 if let Some(index) = entry.index() {
3560 let on = self.catalog().get_entry_by_global_id(&index.on);
3561 indexes_on.entry(on.id()).or_default().push(entry);
3562 } else {
3563 non_indexes.push(entry);
3564 }
3565 }
3566
3567 let key_fn = |entry: &CatalogEntry| entry.id;
3568 let dependencies_fn = |entry: &CatalogEntry| entry.uses();
3569 sort_topological(&mut non_indexes, key_fn, dependencies_fn);
3570
3571 let mut result = Vec::new();
3572 for entry in non_indexes {
3573 let id = entry.id();
3574 result.push(entry);
3575 if let Some(mut indexes) = indexes_on.remove(&id) {
3576 result.append(&mut indexes);
3577 }
3578 }
3579
3580 soft_assert_or_log!(
3581 indexes_on.is_empty(),
3582 "indexes with missing dependencies: {indexes_on:?}",
3583 );
3584
3585 result
3586 }
3587
3588 #[instrument]
3599 fn bootstrap_dataflow_plans(
3600 &mut self,
3601 ordered_catalog_entries: &[CatalogEntry],
3602 mut cached_global_exprs: BTreeMap<GlobalId, GlobalExpressions>,
3603 ) -> Result<BTreeMap<GlobalId, GlobalExpressions>, AdapterError> {
3604 let mut instance_snapshots = BTreeMap::new();
3610 let mut uncached_expressions = BTreeMap::new();
3611
3612 let optimizer_config = |catalog: &Catalog, cluster_id| {
3613 let system_config = catalog.system_config();
3614 let overrides = catalog.get_cluster(cluster_id).config.features();
3615 OptimizerConfig::from(system_config)
3616 .override_from(&overrides)
3617 .override_from(
3620 &catalog
3621 .state()
3622 .cluster_scoped_optimizer_overrides(cluster_id),
3623 )
3624 };
3625
3626 for entry in ordered_catalog_entries {
3627 match entry.item() {
3628 CatalogItem::Index(idx) => {
3629 let compute_instance =
3631 instance_snapshots.entry(idx.cluster_id).or_insert_with(|| {
3632 self.instance_snapshot(idx.cluster_id)
3633 .expect("compute instance exists")
3634 });
3635 let global_id = idx.global_id();
3636
3637 if compute_instance.contains_collection(&global_id) {
3640 continue;
3641 }
3642
3643 let optimizer_config = optimizer_config(&self.catalog, idx.cluster_id);
3644
3645 let (optimized_plan, physical_plan, metainfo) =
3646 match cached_global_exprs.remove(&global_id) {
3647 Some(global_expressions)
3648 if global_expressions.optimizer_features
3649 == optimizer_config.features =>
3650 {
3651 debug!("global expression cache hit for {global_id:?}");
3652 (
3653 global_expressions.global_mir,
3654 global_expressions.physical_plan,
3655 global_expressions.dataflow_metainfos,
3656 )
3657 }
3658 Some(_) | None => {
3659 let (optimized_plan, global_lir_plan) = {
3660 let mut optimizer = optimize::index::Optimizer::new(
3662 self.owned_catalog(),
3663 compute_instance.clone(),
3664 global_id,
3665 optimizer_config.clone(),
3666 self.optimizer_metrics(),
3667 );
3668
3669 let index_plan = optimize::index::Index::new(
3671 entry.name().clone(),
3672 idx.on,
3673 idx.keys.to_vec(),
3674 );
3675 let global_mir_plan = optimizer.optimize(index_plan)?;
3676 let optimized_plan = global_mir_plan.df_desc().clone();
3677
3678 let global_lir_plan = optimizer.optimize(global_mir_plan)?;
3680
3681 (optimized_plan, global_lir_plan)
3682 };
3683
3684 let (physical_plan, metainfo) = global_lir_plan.unapply();
3685 let metainfo = {
3686 let notice_ids =
3688 std::iter::repeat_with(|| self.allocate_transient_id())
3689 .map(|(_item_id, gid)| gid)
3690 .take(metainfo.optimizer_notices.len())
3691 .collect::<Vec<_>>();
3692 self.catalog().render_notices(
3694 metainfo,
3695 notice_ids,
3696 Some(idx.global_id()),
3697 )
3698 };
3699 uncached_expressions.insert(
3700 global_id,
3701 GlobalExpressions {
3702 global_mir: optimized_plan.clone(),
3703 physical_plan: physical_plan.clone(),
3704 dataflow_metainfos: metainfo.clone(),
3705 optimizer_features: optimizer_config.features.clone(),
3706 },
3707 );
3708 (optimized_plan, physical_plan, metainfo)
3709 }
3710 };
3711
3712 let catalog = self.catalog_mut();
3713 catalog.set_optimized_plan(idx.global_id(), optimized_plan);
3714 catalog.set_physical_plan(idx.global_id(), physical_plan);
3715 catalog.set_dataflow_metainfo(idx.global_id(), metainfo);
3716
3717 compute_instance.insert_collection(idx.global_id());
3718 }
3719 CatalogItem::MaterializedView(mv) => {
3720 let compute_instance =
3722 instance_snapshots.entry(mv.cluster_id).or_insert_with(|| {
3723 self.instance_snapshot(mv.cluster_id)
3724 .expect("compute instance exists")
3725 });
3726 let global_id = mv.global_id_writes();
3727
3728 let optimizer_config = optimizer_config(&self.catalog, mv.cluster_id);
3729
3730 let (optimized_plan, physical_plan, metainfo) = match cached_global_exprs
3731 .remove(&global_id)
3732 {
3733 Some(global_expressions)
3734 if global_expressions.optimizer_features
3735 == optimizer_config.features =>
3736 {
3737 debug!("global expression cache hit for {global_id:?}");
3738 (
3739 global_expressions.global_mir,
3740 global_expressions.physical_plan,
3741 global_expressions.dataflow_metainfos,
3742 )
3743 }
3744 Some(_) | None => {
3745 let (_, internal_view_id) = self.allocate_transient_id();
3746 let debug_name = self
3747 .catalog()
3748 .resolve_full_name(entry.name(), None)
3749 .to_string();
3750
3751 let (optimized_plan, global_lir_plan) = {
3752 let mut optimizer = optimize::materialized_view::Optimizer::new(
3754 self.owned_catalog().as_optimizer_catalog(),
3755 compute_instance.clone(),
3756 global_id,
3757 internal_view_id,
3758 mv.desc.latest().iter_names().cloned().collect(),
3759 mv.non_null_assertions.clone(),
3760 mv.refresh_schedule.clone(),
3761 debug_name,
3762 optimizer_config.clone(),
3763 self.optimizer_metrics(),
3764 );
3765
3766 let typ = infer_sql_type_for_catalog(
3769 &mv.raw_expr,
3770 &mv.locally_optimized_expr.as_ref().clone(),
3771 );
3772 let global_mir_plan = optimizer
3773 .optimize((mv.locally_optimized_expr.as_ref().clone(), typ))?;
3774 let optimized_plan = global_mir_plan.df_desc().clone();
3775
3776 let global_lir_plan = optimizer.optimize(global_mir_plan)?;
3778
3779 (optimized_plan, global_lir_plan)
3780 };
3781
3782 let (physical_plan, metainfo) = global_lir_plan.unapply();
3783 let metainfo = {
3784 let notice_ids =
3786 std::iter::repeat_with(|| self.allocate_transient_id())
3787 .map(|(_item_id, global_id)| global_id)
3788 .take(metainfo.optimizer_notices.len())
3789 .collect::<Vec<_>>();
3790 self.catalog().render_notices(
3792 metainfo,
3793 notice_ids,
3794 Some(mv.global_id_writes()),
3795 )
3796 };
3797 uncached_expressions.insert(
3798 global_id,
3799 GlobalExpressions {
3800 global_mir: optimized_plan.clone(),
3801 physical_plan: physical_plan.clone(),
3802 dataflow_metainfos: metainfo.clone(),
3803 optimizer_features: optimizer_config.features.clone(),
3804 },
3805 );
3806 (optimized_plan, physical_plan, metainfo)
3807 }
3808 };
3809
3810 let catalog = self.catalog_mut();
3811 catalog.set_optimized_plan(mv.global_id_writes(), optimized_plan);
3812 catalog.set_physical_plan(mv.global_id_writes(), physical_plan);
3813 catalog.set_dataflow_metainfo(mv.global_id_writes(), metainfo);
3814
3815 compute_instance.insert_collection(mv.global_id_writes());
3816 }
3817 CatalogItem::MetricSink(metric_sink) => {
3818 let compute_instance = instance_snapshots
3820 .entry(metric_sink.cluster_id)
3821 .or_insert_with(|| {
3822 self.instance_snapshot(metric_sink.cluster_id)
3823 .expect("compute instance exists")
3824 });
3825 let global_id = metric_sink.global_id;
3826 let optimizer_config = optimizer_config(&self.catalog, metric_sink.cluster_id);
3827
3828 let (optimized_plan, physical_plan, metainfo) = match cached_global_exprs
3829 .remove(&global_id)
3830 {
3831 Some(global_expressions)
3832 if global_expressions.optimizer_features
3833 == optimizer_config.features =>
3834 {
3835 debug!("global expression cache hit for {global_id:?}");
3836 (
3837 global_expressions.global_mir,
3838 global_expressions.physical_plan,
3839 global_expressions.dataflow_metainfos,
3840 )
3841 }
3842 Some(_) | None => {
3843 let (_, view_id) = self.allocate_transient_id();
3850
3851 let (optimized_plan, global_lir_plan) = {
3852 let mut optimizer = optimize::metric_sink::Optimizer::new(
3853 self.owned_catalog(),
3854 compute_instance.clone(),
3855 view_id,
3856 global_id,
3857 optimizer_config.clone(),
3858 self.optimizer_metrics(),
3859 );
3860
3861 let metric_sink_plan = optimize::metric_sink::MetricSink::new(
3863 entry.name().clone(),
3864 metric_sink.from,
3865 metric_sink.prefix.clone(),
3866 );
3867 let global_mir_plan = optimizer.optimize(metric_sink_plan)?;
3868 let optimized_plan = global_mir_plan.df_desc().clone();
3869
3870 let global_lir_plan = optimizer.optimize(global_mir_plan)?;
3872
3873 (optimized_plan, global_lir_plan)
3874 };
3875
3876 let (physical_plan, metainfo) = global_lir_plan.unapply();
3877 let metainfo = {
3878 let notice_ids =
3880 std::iter::repeat_with(|| self.allocate_transient_id())
3881 .map(|(_item_id, gid)| gid)
3882 .take(metainfo.optimizer_notices.len())
3883 .collect::<Vec<_>>();
3884 self.catalog()
3886 .render_notices(metainfo, notice_ids, Some(global_id))
3887 };
3888 uncached_expressions.insert(
3889 global_id,
3890 GlobalExpressions {
3891 global_mir: optimized_plan.clone(),
3892 physical_plan: physical_plan.clone(),
3893 dataflow_metainfos: metainfo.clone(),
3894 optimizer_features: optimizer_config.features.clone(),
3895 },
3896 );
3897 (optimized_plan, physical_plan, metainfo)
3898 }
3899 };
3900
3901 let catalog = self.catalog_mut();
3902 catalog.set_optimized_plan(global_id, optimized_plan);
3903 catalog.set_physical_plan(global_id, physical_plan);
3904 catalog.set_dataflow_metainfo(global_id, metainfo);
3905
3906 }
3910 CatalogItem::Table(_)
3911 | CatalogItem::Source(_)
3912 | CatalogItem::Log(_)
3913 | CatalogItem::View(_)
3914 | CatalogItem::Sink(_)
3915 | CatalogItem::Type(_)
3916 | CatalogItem::Func(_)
3917 | CatalogItem::Secret(_)
3918 | CatalogItem::Connection(_) => (),
3919 }
3920 }
3921
3922 Ok(uncached_expressions)
3923 }
3924
3925 async fn bootstrap_dataflow_as_ofs(&mut self) -> BTreeMap<GlobalId, ReadHold> {
3935 let mut catalog_ids = Vec::new();
3936 let mut dataflows = Vec::new();
3937 let mut read_policies = BTreeMap::new();
3938 for entry in self.catalog.entries() {
3939 let gid = match entry.item() {
3940 CatalogItem::Index(idx) => idx.global_id(),
3941 CatalogItem::MaterializedView(mv) => mv.global_id_writes(),
3942 CatalogItem::MetricSink(metric_sink) => metric_sink.global_id,
3943 CatalogItem::Table(_)
3944 | CatalogItem::Source(_)
3945 | CatalogItem::Log(_)
3946 | CatalogItem::View(_)
3947 | CatalogItem::Sink(_)
3948 | CatalogItem::Type(_)
3949 | CatalogItem::Func(_)
3950 | CatalogItem::Secret(_)
3951 | CatalogItem::Connection(_) => continue,
3952 };
3953 if let Some(plan) = self.catalog.try_get_physical_plan(&gid) {
3954 catalog_ids.push(gid);
3955 dataflows.push(plan.clone());
3956
3957 if let Some(compaction_window) = entry.item().initial_logical_compaction_window() {
3958 read_policies.insert(gid, compaction_window.into());
3959 }
3960 }
3961 }
3962
3963 let read_ts = self.get_local_read_ts().await;
3964 let read_holds = as_of_selection::run(
3965 &mut dataflows,
3966 &read_policies,
3967 &*self.controller.storage_collections,
3968 read_ts,
3969 self.controller.read_only(),
3970 );
3971
3972 let catalog = self.catalog_mut();
3973 for (id, plan) in catalog_ids.into_iter().zip_eq(dataflows) {
3974 catalog.set_physical_plan(id, plan);
3975 }
3976
3977 read_holds
3978 }
3979
3980 fn serve(
3989 mut self,
3990 mut internal_cmd_rx: mpsc::UnboundedReceiver<Message>,
3991 mut strict_serializable_reads_rx: mpsc::UnboundedReceiver<(ConnectionId, PendingReadTxn)>,
3992 mut cmd_rx: mpsc::UnboundedReceiver<(OpenTelemetryContext, Command)>,
3993 group_commit_rx: appends::GroupCommitWaiter,
3994 ) -> LocalBoxFuture<'static, ()> {
3995 async move {
3996 let mut cluster_events = self.controller.events_stream();
3998 let last_message = Arc::new(Mutex::new(LastMessage {
3999 kind: "none",
4000 stmt: None,
4001 }));
4002
4003 let (idle_tx, mut idle_rx) = tokio::sync::mpsc::channel(1);
4004 let idle_metric = self.metrics.queue_busy_seconds.clone();
4005 let last_message_watchdog = Arc::clone(&last_message);
4006
4007 spawn(|| "coord watchdog", async move {
4008 let mut interval = tokio::time::interval(Duration::from_secs(5));
4013 interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
4017
4018 let mut coord_stuck = false;
4020
4021 loop {
4022 interval.tick().await;
4023
4024 let duration = tokio::time::Duration::from_secs(30);
4026 let timeout = tokio::time::timeout(duration, idle_tx.reserve()).await;
4027 let Ok(maybe_permit) = timeout else {
4028 if !coord_stuck {
4030 let last_message = last_message_watchdog.lock().expect("poisoned");
4031 tracing::warn!(
4032 last_message_kind = %last_message.kind,
4033 last_message_sql = %last_message.stmt_to_string(),
4034 "coordinator stuck for {duration:?}",
4035 );
4036 }
4037 coord_stuck = true;
4038
4039 continue;
4040 };
4041
4042 if coord_stuck {
4044 tracing::info!("Coordinator became unstuck");
4045 }
4046 coord_stuck = false;
4047
4048 let Ok(permit) = maybe_permit else {
4050 break;
4051 };
4052
4053 permit.send(idle_metric.start_timer());
4054 }
4055 });
4056
4057 self.schedule_storage_usage_collection().await;
4058 self.schedule_arrangement_sizes_collection().await;
4059 self.spawn_privatelink_vpc_endpoints_watch_task();
4060 self.spawn_statement_logging_task();
4061 self.spawn_catalog_info_metrics_task();
4062 self.spawn_cluster_controller_task();
4063 flags::tracing_config(self.catalog.system_config()).apply(&self.tracing_handle);
4064
4065 let warn_threshold = self
4067 .catalog()
4068 .system_config()
4069 .coord_slow_message_warn_threshold();
4070
4071 const MESSAGE_BATCH: usize = 64;
4073 let mut messages = Vec::with_capacity(MESSAGE_BATCH);
4074 let mut cmd_messages = Vec::with_capacity(MESSAGE_BATCH);
4075
4076 let message_batch = self.metrics.message_batch.clone();
4077
4078 let linearize_reads_notify = Arc::clone(&self.linearize_reads_notify);
4085 let linearize_reads_notified = linearize_reads_notify.notified();
4086 tokio::pin!(linearize_reads_notified);
4087
4088 loop {
4089 select! {
4093 biased;
4098
4099 _ = internal_cmd_rx.recv_many(&mut messages, MESSAGE_BATCH) => {},
4103 Some(event) = cluster_events.next() => {
4107 messages.push(Message::ClusterEvent(event))
4108 },
4109 () = self.controller.ready() => {
4113 let controller = match self.controller.get_readiness() {
4117 Readiness::Storage => ControllerReadiness::Storage,
4118 Readiness::Compute => ControllerReadiness::Compute,
4119 Readiness::Metrics(_) => ControllerReadiness::Metrics,
4120 Readiness::Internal(_) => ControllerReadiness::Internal,
4121 Readiness::NotReady => unreachable!("just signaled as ready"),
4122 };
4123 messages.push(Message::ControllerReady { controller });
4124 }
4125 permit = group_commit_rx.ready() => {
4128 let user_write_spans = self.pending_writes.iter().flat_map(|x| match x {
4134 PendingWriteTxn::User { span, .. } => Some(span),
4135 PendingWriteTxn::System { .. } => None,
4136 });
4137 let span = match user_write_spans.exactly_one() {
4138 Ok(span) => span.clone(),
4139 Err(user_write_spans) => {
4140 let span = info_span!(parent: None, "group_commit_notify");
4141 for s in user_write_spans {
4142 span.follows_from(s);
4143 }
4144 span
4145 }
4146 };
4147 messages.push(Message::GroupCommitInitiate(span, Some(permit)));
4148 },
4149 count = cmd_rx.recv_many(&mut cmd_messages, MESSAGE_BATCH) => {
4153 if count == 0 {
4154 break;
4155 } else {
4156 messages.extend(cmd_messages.drain(..).map(
4157 |(otel_ctx, cmd)| Message::Command(otel_ctx, cmd),
4158 ));
4159 }
4160 },
4161 Some(pending_read_txn) = strict_serializable_reads_rx.recv() => {
4165 let mut pending_read_txns = vec![pending_read_txn];
4166 while let Ok(pending_read_txn) = strict_serializable_reads_rx.try_recv() {
4167 pending_read_txns.push(pending_read_txn);
4168 }
4169 for (conn_id, pending_read_txn) in pending_read_txns {
4170 let prev = self
4171 .pending_linearize_read_txns
4172 .insert(conn_id, pending_read_txn);
4173 soft_assert_or_log!(
4174 prev.is_none(),
4175 "connections can not have multiple concurrent reads, prev: {prev:?}"
4176 )
4177 }
4178 messages.push(Message::LinearizeReads);
4179 }
4180 _ = self.advance_timelines_interval.tick() => {
4184 if self.controller.read_only() {
4188 messages.push(Message::AdvanceTimelines);
4189 } else {
4190 self.group_commit_tx.notify();
4191 }
4192 },
4193 () = linearize_reads_notified.as_mut() => {
4204 linearize_reads_notified.set(linearize_reads_notify.notified());
4205 messages.push(Message::LinearizeReads);
4206 }
4207 _ = self.caught_up_check_interval.tick() => {
4211 self.maybe_check_caught_up().await;
4216
4217 continue;
4218 },
4219
4220 timer = idle_rx.recv() => {
4225 timer.expect("does not drop").observe_duration();
4226 self.metrics
4227 .message_handling
4228 .with_label_values(&["watchdog"])
4229 .observe(0.0);
4230 continue;
4231 }
4232 };
4233
4234 message_batch.observe(f64::cast_lossy(messages.len()));
4236
4237 for msg in messages.drain(..) {
4238 let msg_kind = msg.kind();
4241 let span = span!(
4242 target: "mz_adapter::coord::handle_message_loop",
4243 Level::INFO,
4244 "coord::handle_message",
4245 kind = msg_kind
4246 );
4247 let otel_context = span.context().span().span_context().clone();
4248
4249 *last_message.lock().expect("poisoned") = LastMessage {
4253 kind: msg_kind,
4254 stmt: match &msg {
4255 Message::Command(
4256 _,
4257 Command::Execute {
4258 portal_name,
4259 session,
4260 ..
4261 },
4262 ) => session
4263 .get_portal_unverified(portal_name)
4264 .and_then(|p| p.stmt.as_ref().map(Arc::clone)),
4265 _ => None,
4266 },
4267 };
4268
4269 let start = Instant::now();
4270 self.handle_message(msg).instrument(span).await;
4271 let duration = start.elapsed();
4272
4273 self.metrics
4274 .message_handling
4275 .with_label_values(&[msg_kind])
4276 .observe(duration.as_secs_f64());
4277
4278 if duration > warn_threshold {
4280 let trace_id = otel_context.is_valid().then(|| otel_context.trace_id());
4281 tracing::error!(
4282 ?msg_kind,
4283 ?trace_id,
4284 ?duration,
4285 "very slow coordinator message"
4286 );
4287 }
4288 }
4289 }
4290 if let Some(catalog) = Arc::into_inner(self.catalog) {
4293 catalog.expire().await;
4294 }
4295 }
4296 .boxed_local()
4297 }
4298
4299 fn catalog(&self) -> &Catalog {
4301 &self.catalog
4302 }
4303
4304 fn owned_catalog(&self) -> Arc<Catalog> {
4307 Arc::clone(&self.catalog)
4308 }
4309
4310 fn optimizer_metrics(&self) -> OptimizerMetrics {
4313 self.optimizer_metrics.clone()
4314 }
4315
4316 fn catalog_mut(&mut self) -> &mut Catalog {
4318 Arc::make_mut(&mut self.catalog)
4326 }
4327
4328 async fn refill_user_id_pool(&mut self, min_count: u64) -> Result<(), AdapterError> {
4333 let batch_size = USER_ID_POOL_BATCH_SIZE.get(self.catalog().system_config().dyncfgs());
4334 let to_allocate = min_count.max(u64::from(batch_size));
4335 let id_ts = self.get_catalog_write_ts().await;
4336 let ids = self.catalog().allocate_user_ids(to_allocate, id_ts).await?;
4337 if let (Some((first_id, _)), Some((last_id, _))) = (ids.first(), ids.last()) {
4338 let start = match first_id {
4339 CatalogItemId::User(id) => *id,
4340 other => {
4341 return Err(AdapterError::Internal(format!(
4342 "expected User CatalogItemId, got {other:?}"
4343 )));
4344 }
4345 };
4346 let end = match last_id {
4347 CatalogItemId::User(id) => *id + 1, other => {
4349 return Err(AdapterError::Internal(format!(
4350 "expected User CatalogItemId, got {other:?}"
4351 )));
4352 }
4353 };
4354 self.user_id_pool.refill(start, end);
4355 } else {
4356 return Err(AdapterError::Internal(
4357 "catalog returned no user IDs".into(),
4358 ));
4359 }
4360 Ok(())
4361 }
4362
4363 async fn allocate_user_id(&mut self) -> Result<(CatalogItemId, GlobalId), AdapterError> {
4365 if let Some(id) = self.user_id_pool.allocate() {
4366 return Ok((CatalogItemId::User(id), GlobalId::User(id)));
4367 }
4368 self.refill_user_id_pool(1).await?;
4369 let id = self.user_id_pool.allocate().expect("ID pool just refilled");
4370 Ok((CatalogItemId::User(id), GlobalId::User(id)))
4371 }
4372
4373 async fn allocate_user_ids(
4375 &mut self,
4376 count: u64,
4377 ) -> Result<Vec<(CatalogItemId, GlobalId)>, AdapterError> {
4378 if self.user_id_pool.remaining() < count {
4379 self.refill_user_id_pool(count).await?;
4380 }
4381 let raw_ids = self
4382 .user_id_pool
4383 .allocate_many(count)
4384 .expect("pool has enough IDs after refill");
4385 Ok(raw_ids
4386 .into_iter()
4387 .map(|id| (CatalogItemId::User(id), GlobalId::User(id)))
4388 .collect())
4389 }
4390
4391 fn connection_context(&self) -> &ConnectionContext {
4393 self.controller.connection_context()
4394 }
4395
4396 fn secrets_reader(&self) -> &Arc<dyn SecretsReader> {
4398 &self.connection_context().secrets_reader
4399 }
4400
4401 #[allow(dead_code)]
4406 pub(crate) fn broadcast_notice(&self, notice: AdapterNotice) {
4407 for meta in self.active_conns.values() {
4408 let _ = meta.notice_tx.send(notice.clone());
4409 }
4410 }
4411
4412 pub(crate) fn broadcast_notice_tx(
4415 &self,
4416 ) -> Box<dyn FnOnce(AdapterNotice) -> () + Send + 'static> {
4417 let senders: Vec<_> = self
4418 .active_conns
4419 .values()
4420 .map(|meta| meta.notice_tx.clone())
4421 .collect();
4422 Box::new(move |notice| {
4423 for tx in senders {
4424 let _ = tx.send(notice.clone());
4425 }
4426 })
4427 }
4428
4429 pub(crate) fn active_conns(&self) -> &BTreeMap<ConnectionId, ConnMeta> {
4430 &self.active_conns
4431 }
4432
4433 #[instrument(level = "debug")]
4434 pub(crate) fn retire_execution(
4435 &mut self,
4436 reason: StatementEndedExecutionReason,
4437 ctx_extra: ExecuteContextExtra,
4438 ) {
4439 if let Some(uuid) = ctx_extra.retire() {
4440 let ended_at = self.now();
4441 self.end_statement_execution(uuid, reason, ended_at);
4442 }
4443 }
4444
4445 #[instrument(level = "debug")]
4447 pub fn dataflow_builder(&self, instance: ComputeInstanceId) -> DataflowBuilder<'_> {
4448 let compute = self
4449 .instance_snapshot(instance)
4450 .expect("compute instance does not exist");
4451 DataflowBuilder::new(self.catalog().state(), compute)
4452 }
4453
4454 pub fn instance_snapshot(
4456 &self,
4457 id: ComputeInstanceId,
4458 ) -> Result<ComputeInstanceSnapshot, InstanceMissing> {
4459 ComputeInstanceSnapshot::new(&self.controller, id)
4460 }
4461
4462 pub(crate) async fn ship_dataflow(
4469 &mut self,
4470 dataflow: DataflowDescription<LirRelationExpr>,
4471 instance: ComputeInstanceId,
4472 target_replica: Option<ReplicaId>,
4473 ) {
4474 self.try_ship_dataflow(dataflow, instance, target_replica)
4475 .await
4476 .unwrap_or_terminate("dataflow creation cannot fail");
4477 }
4478
4479 pub(crate) async fn try_ship_dataflow(
4482 &mut self,
4483 dataflow: DataflowDescription<LirRelationExpr>,
4484 instance: ComputeInstanceId,
4485 target_replica: Option<ReplicaId>,
4486 ) -> Result<(), DataflowCreationError> {
4487 let export_ids = dataflow.exported_index_ids().collect();
4490
4491 self.controller
4492 .compute
4493 .create_dataflow(instance, dataflow, target_replica)?;
4494
4495 self.initialize_compute_read_policies(export_ids, instance, CompactionWindow::Default)
4496 .await;
4497
4498 Ok(())
4499 }
4500
4501 pub(crate) fn allow_writes(&mut self, instance: ComputeInstanceId, id: GlobalId) {
4505 self.controller
4506 .compute
4507 .allow_writes(instance, id)
4508 .unwrap_or_terminate("allow_writes cannot fail");
4509 }
4510
4511 pub(crate) async fn ship_dataflow_and_notice_builtin_table_updates(
4513 &mut self,
4514 dataflow: DataflowDescription<LirRelationExpr>,
4515 instance: ComputeInstanceId,
4516 notice_builtin_updates_fut: Option<BuiltinTableAppendNotify>,
4517 target_replica: Option<ReplicaId>,
4518 ) {
4519 if let Some(notice_builtin_updates_fut) = notice_builtin_updates_fut {
4520 let ship_dataflow_fut = self.ship_dataflow(dataflow, instance, target_replica);
4521 let ((), ()) =
4522 futures::future::join(notice_builtin_updates_fut, ship_dataflow_fut).await;
4523 } else {
4524 self.ship_dataflow(dataflow, instance, target_replica).await;
4525 }
4526 }
4527
4528 pub fn install_compute_watch_set(
4532 &mut self,
4533 conn_id: ConnectionId,
4534 objects: BTreeSet<GlobalId>,
4535 t: Timestamp,
4536 state: WatchSetResponse,
4537 ) -> Result<(), CollectionLookupError> {
4538 let ws_id = self.controller.install_compute_watch_set(objects, t)?;
4539 self.connection_watch_sets
4540 .entry(conn_id.clone())
4541 .or_default()
4542 .insert(ws_id);
4543 self.installed_watch_sets.insert(ws_id, (conn_id, state));
4544 Ok(())
4545 }
4546
4547 pub fn install_storage_watch_set(
4551 &mut self,
4552 conn_id: ConnectionId,
4553 objects: BTreeSet<GlobalId>,
4554 t: Timestamp,
4555 state: WatchSetResponse,
4556 ) -> Result<(), CollectionMissing> {
4557 let ws_id = self.controller.install_storage_watch_set(objects, t)?;
4558 self.connection_watch_sets
4559 .entry(conn_id.clone())
4560 .or_default()
4561 .insert(ws_id);
4562 self.installed_watch_sets.insert(ws_id, (conn_id, state));
4563 Ok(())
4564 }
4565
4566 pub fn cancel_pending_watchsets(&mut self, conn_id: &ConnectionId) {
4568 if let Some(ws_ids) = self.connection_watch_sets.remove(conn_id) {
4569 for ws_id in ws_ids {
4570 self.installed_watch_sets.remove(&ws_id);
4571 }
4572 }
4573 }
4574
4575 pub async fn dump(&self) -> Result<serde_json::Value, anyhow::Error> {
4579 let global_timelines: BTreeMap<_, _> = self
4585 .global_timelines
4586 .iter()
4587 .map(|(timeline, state)| (timeline.to_string(), format!("{state:?}")))
4588 .collect();
4589 let active_conns: BTreeMap<_, _> = self
4590 .active_conns
4591 .iter()
4592 .map(|(id, meta)| (id.unhandled().to_string(), format!("{meta:?}")))
4593 .collect();
4594 let txn_read_holds: BTreeMap<_, _> = self
4595 .txn_read_holds
4596 .iter()
4597 .map(|(id, capability)| (id.unhandled().to_string(), format!("{capability:?}")))
4598 .collect();
4599 let pending_peeks: BTreeMap<_, _> = self
4600 .pending_peeks
4601 .iter()
4602 .map(|(id, peek)| (id.to_string(), format!("{peek:?}")))
4603 .collect();
4604 let client_pending_peeks: BTreeMap<_, _> = self
4605 .client_pending_peeks
4606 .iter()
4607 .map(|(id, peek)| {
4608 let peek: BTreeMap<_, _> = peek
4609 .iter()
4610 .map(|(uuid, storage_id)| (uuid.to_string(), storage_id))
4611 .collect();
4612 (id.to_string(), peek)
4613 })
4614 .collect();
4615 let pending_linearize_read_txns: BTreeMap<_, _> = self
4616 .pending_linearize_read_txns
4617 .iter()
4618 .map(|(id, read_txn)| (id.unhandled().to_string(), format!("{read_txn:?}")))
4619 .collect();
4620
4621 Ok(serde_json::json!({
4622 "global_timelines": global_timelines,
4623 "active_conns": active_conns,
4624 "txn_read_holds": txn_read_holds,
4625 "pending_peeks": pending_peeks,
4626 "client_pending_peeks": client_pending_peeks,
4627 "pending_linearize_read_txns": pending_linearize_read_txns,
4628 "controller": self.controller.dump().await?,
4629 }))
4630 }
4631
4632 async fn prune_storage_usage_events_on_startup(&self, retention_period: Duration) {
4646 let item_id = self
4647 .catalog()
4648 .resolve_builtin_table(&MZ_STORAGE_USAGE_BY_SHARD);
4649 let global_id = self.catalog.get_entry(&item_id).latest_global_id();
4650 let read_ts = self.get_local_read_ts().await;
4651 let current_contents_fut = self
4652 .controller
4653 .storage_collections
4654 .snapshot(global_id, read_ts);
4655 let internal_cmd_tx = self.internal_cmd_tx.clone();
4656 spawn(|| "storage_usage_prune", async move {
4657 let mut current_contents = current_contents_fut
4658 .await
4659 .unwrap_or_terminate("cannot fail to fetch snapshot");
4660 differential_dataflow::consolidation::consolidate(&mut current_contents);
4661
4662 let cutoff_ts = u128::from(read_ts).saturating_sub(retention_period.as_millis());
4663 let mut expired = Vec::new();
4664 for (row, diff) in current_contents {
4665 assert_eq!(
4666 diff, 1,
4667 "consolidated contents should not contain retractions: ({row:#?}, {diff:#?})"
4668 );
4669 let collection_timestamp = row
4671 .unpack()
4672 .get(3)
4673 .expect("definition of mz_storage_by_shard changed")
4674 .unwrap_timestamptz();
4675 let collection_timestamp = collection_timestamp.timestamp_millis();
4676 let collection_timestamp: u128 = collection_timestamp
4677 .try_into()
4678 .expect("all collections happen after Jan 1 1970");
4679 if collection_timestamp < cutoff_ts {
4680 debug!("pruning storage event {row:?}");
4681 let builtin_update = BuiltinTableUpdate::row(item_id, row, Diff::MINUS_ONE);
4682 expired.push(builtin_update);
4683 }
4684 }
4685
4686 let _ = internal_cmd_tx.send(Message::StorageUsagePrune(expired));
4688 });
4689 }
4690
4691 async fn prune_arrangement_sizes_history_on_startup(&self) {
4700 if self.controller.read_only() {
4702 return;
4703 }
4704
4705 let retention_period = mz_adapter_types::dyncfgs::ARRANGEMENT_SIZE_HISTORY_RETENTION_PERIOD
4706 .get(self.catalog().system_config().dyncfgs());
4707 let item_id = self
4708 .catalog()
4709 .resolve_builtin_table(&mz_catalog::builtin::MZ_OBJECT_ARRANGEMENT_SIZE_HISTORY);
4710 let global_id = self.catalog.get_entry(&item_id).latest_global_id();
4711 let read_ts = self.get_local_read_ts().await;
4712 let current_contents_fut = self
4713 .controller
4714 .storage_collections
4715 .snapshot(global_id, read_ts);
4716 let internal_cmd_tx = self.internal_cmd_tx.clone();
4717 spawn(|| "arrangement_sizes_history_prune", async move {
4718 let mut current_contents = current_contents_fut
4719 .await
4720 .unwrap_or_terminate("cannot fail to fetch snapshot");
4721 differential_dataflow::consolidation::consolidate(&mut current_contents);
4722
4723 let cutoff_ts = u128::from(read_ts).saturating_sub(retention_period.as_millis());
4724 let expired =
4725 arrangement_sizes_expired_retractions(current_contents, cutoff_ts, item_id);
4726
4727 let _ = internal_cmd_tx.send(Message::ArrangementSizesPrune(expired));
4731 });
4732 }
4733
4734 fn current_credit_consumption_rate(&self, exclude_cluster: Option<ClusterId>) -> Numeric {
4737 self.catalog()
4738 .user_cluster_replicas()
4739 .filter(|replica| Some(replica.cluster_id) != exclude_cluster)
4740 .filter_map(|replica| match &replica.config.location {
4741 ReplicaLocation::Managed(location) => Some(location.size_for_billing()),
4742 ReplicaLocation::Unmanaged(_) => None,
4743 })
4744 .map(|size| {
4745 self.catalog()
4746 .cluster_replica_sizes()
4747 .0
4748 .get(size)
4749 .expect("location size is validated against the cluster replica sizes")
4750 .credits_per_hour
4751 })
4752 .sum()
4753 }
4754}
4755
4756fn arrangement_sizes_expired_retractions(
4764 rows: impl IntoIterator<Item = (mz_repr::Row, i64)>,
4765 cutoff_ts: u128,
4766 item_id: CatalogItemId,
4767) -> Vec<BuiltinTableUpdate> {
4768 let mut expired = Vec::new();
4769 for (row, diff) in rows {
4770 assert_eq!(
4771 diff, 1,
4772 "consolidated contents should not contain retractions: ({row:#?}, {diff:#?})"
4773 );
4774 let collection_timestamp = row
4775 .unpack()
4776 .get(3)
4777 .expect("definition of mz_object_arrangement_size_history changed")
4778 .unwrap_timestamptz()
4779 .timestamp_millis();
4780 let collection_timestamp: u128 = collection_timestamp
4781 .try_into()
4782 .expect("all collections happen after Jan 1 1970");
4783 if collection_timestamp < cutoff_ts {
4784 expired.push(BuiltinTableUpdate::row(item_id, row, Diff::MINUS_ONE));
4785 }
4786 }
4787 expired
4788}
4789
4790#[cfg(test)]
4791impl Coordinator {
4792 #[allow(dead_code)]
4793 async fn verify_ship_dataflow_no_error(
4794 &mut self,
4795 dataflow: DataflowDescription<LirRelationExpr>,
4796 ) {
4797 let compute_instance = ComputeInstanceId::user(1).expect("1 is a valid ID");
4805
4806 let _: () = self.ship_dataflow(dataflow, compute_instance, None).await;
4807 }
4808}
4809
4810struct LastMessage {
4812 kind: &'static str,
4813 stmt: Option<Arc<Statement<Raw>>>,
4814}
4815
4816impl LastMessage {
4817 fn stmt_to_string(&self) -> Cow<'static, str> {
4819 self.stmt
4820 .as_ref()
4821 .map(|stmt| stmt.to_ast_string_redacted().into())
4822 .unwrap_or(Cow::Borrowed("<none>"))
4823 }
4824}
4825
4826impl fmt::Debug for LastMessage {
4827 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
4828 f.debug_struct("LastMessage")
4829 .field("kind", &self.kind)
4830 .field("stmt", &self.stmt_to_string())
4831 .finish()
4832 }
4833}
4834
4835impl Drop for LastMessage {
4836 fn drop(&mut self) {
4837 if std::thread::panicking() {
4839 eprintln!("Coordinator panicking, dumping last message\n{self:?}",);
4841 }
4842 }
4843}
4844
4845pub fn serve(
4857 Config {
4858 controller_config,
4859 controller_envd_epoch,
4860 mut storage,
4861 timestamp_oracle_url,
4862 unsafe_mode,
4863 all_features,
4864 build_info,
4865 environment_id,
4866 metrics_registry,
4867 now,
4868 secrets_controller,
4869 cloud_resource_controller,
4870 cluster_replica_sizes,
4871 builtin_system_cluster_config,
4872 builtin_catalog_server_cluster_config,
4873 builtin_probe_cluster_config,
4874 builtin_support_cluster_config,
4875 builtin_analytics_cluster_config,
4876 system_parameter_defaults,
4877 availability_zones,
4878 storage_usage_client,
4879 storage_usage_collection_interval,
4880 storage_usage_retention_period,
4881 segment_client,
4882 egress_addresses,
4883 aws_account_id,
4884 aws_privatelink_availability_zones,
4885 connection_context,
4886 connection_limit_callback,
4887 remote_system_parameters,
4888 webhook_concurrency_limit,
4889 http_host_name,
4890 tracing_handle,
4891 read_only_controllers,
4892 caught_up_trigger: clusters_caught_up_trigger,
4893 helm_chart_version,
4894 license_key,
4895 external_login_password_mz_system,
4896 force_builtin_schema_migration,
4897 }: Config,
4898) -> BoxFuture<'static, Result<(Handle, Client), AdapterError>> {
4899 async move {
4900 let coord_start = Instant::now();
4901 info!("startup: coordinator init: beginning");
4902 info!("startup: coordinator init: preamble beginning");
4903
4904 let _builtins = LazyLock::force(&BUILTINS_STATIC);
4908
4909 let (cmd_tx, cmd_rx) = mpsc::unbounded_channel();
4910 let (internal_cmd_tx, internal_cmd_rx) = mpsc::unbounded_channel();
4911 let (strict_serializable_reads_tx, strict_serializable_reads_rx) =
4912 mpsc::unbounded_channel();
4913
4914 if !availability_zones.iter().all_unique() {
4916 coord_bail!("availability zones must be unique");
4917 }
4918
4919 let aws_principal_context = match (
4920 aws_account_id,
4921 connection_context.aws_external_id_prefix.clone(),
4922 ) {
4923 (Some(aws_account_id), Some(aws_external_id_prefix)) => Some(AwsPrincipalContext {
4924 aws_account_id,
4925 aws_external_id_prefix,
4926 }),
4927 _ => None,
4928 };
4929
4930 let aws_privatelink_availability_zones = aws_privatelink_availability_zones
4931 .map(|azs_vec| BTreeSet::from_iter(azs_vec.iter().cloned()));
4932
4933 info!(
4934 "startup: coordinator init: preamble complete in {:?}",
4935 coord_start.elapsed()
4936 );
4937 let oracle_init_start = Instant::now();
4938 info!("startup: coordinator init: timestamp oracle init beginning");
4939
4940 let timestamp_oracle_config = timestamp_oracle_url
4941 .map(|url| TimestampOracleConfig::from_url(&url, &metrics_registry))
4942 .transpose()?;
4943 let mut initial_timestamps =
4944 get_initial_oracle_timestamps(×tamp_oracle_config).await?;
4945
4946 initial_timestamps
4950 .entry(Timeline::EpochMilliseconds)
4951 .or_insert_with(mz_repr::Timestamp::minimum);
4952 let mut timestamp_oracles = BTreeMap::new();
4953 for (timeline, initial_timestamp) in initial_timestamps {
4954 Coordinator::ensure_timeline_state_with_initial_time(
4955 &timeline,
4956 initial_timestamp,
4957 now.clone(),
4958 timestamp_oracle_config.clone(),
4959 &mut timestamp_oracles,
4960 read_only_controllers,
4961 )
4962 .await;
4963 }
4964
4965 let catalog_upper = storage.current_upper().await;
4969 let epoch_millis_oracle = ×tamp_oracles
4975 .get(&Timeline::EpochMilliseconds)
4976 .expect("inserted above")
4977 .oracle;
4978
4979 let boot_now: mz_repr::Timestamp = (now)().into();
4984 if catalog_upper > timeline::write_ts_upper_bound(&boot_now) {
4985 tracing::error!(
4986 %catalog_upper, %boot_now,
4987 "catalog upper is far ahead of the wall clock, so writes and \
4988 strict-serializable reads on the EpochMilliseconds timeline will block \
4989 until the clock catches up",
4990 );
4991 }
4992
4993 let mut boot_ts = if read_only_controllers {
4994 let read_ts = epoch_millis_oracle.read_ts().await;
4995 std::cmp::max(read_ts, catalog_upper)
4996 } else {
4997 epoch_millis_oracle.apply_write(catalog_upper).await;
5000 epoch_millis_oracle.write_ts().await.timestamp
5001 };
5002
5003 info!(
5004 "startup: coordinator init: timestamp oracle init complete in {:?}",
5005 oracle_init_start.elapsed()
5006 );
5007
5008 let catalog_open_start = Instant::now();
5009 info!("startup: coordinator init: catalog open beginning");
5010 let persist_client = controller_config
5011 .persist_clients
5012 .open(controller_config.persist_location.clone())
5013 .await
5014 .context("opening persist client")?;
5015 let builtin_item_migration_config =
5016 BuiltinItemMigrationConfig {
5017 persist_client: persist_client.clone(),
5018 read_only: read_only_controllers,
5019 force_migration: force_builtin_schema_migration,
5020 }
5021 ;
5022 let OpenCatalogResult {
5023 mut catalog,
5024 migrated_storage_collections_0dt,
5025 new_builtin_collections,
5026 builtin_table_updates,
5027 cached_global_exprs,
5028 uncached_local_exprs,
5029 } = Catalog::open(mz_catalog::config::Config {
5030 storage,
5031 metrics_registry: &metrics_registry,
5032 state: mz_catalog::config::StateConfig {
5033 unsafe_mode,
5034 all_features,
5035 build_info,
5036 environment_id: environment_id.clone(),
5037 read_only: read_only_controllers,
5038 now: now.clone(),
5039 boot_ts: boot_ts.clone(),
5040 skip_migrations: false,
5041 cluster_replica_sizes,
5042 builtin_system_cluster_config,
5043 builtin_catalog_server_cluster_config,
5044 builtin_probe_cluster_config,
5045 builtin_support_cluster_config,
5046 builtin_analytics_cluster_config,
5047 system_parameter_defaults,
5048 remote_system_parameters,
5049 availability_zones,
5050 egress_addresses,
5051 aws_principal_context,
5052 aws_privatelink_availability_zones,
5053 connection_context,
5054 http_host_name,
5055 builtin_item_migration_config,
5056 persist_client: persist_client.clone(),
5057 enable_expression_cache_override: None,
5058 helm_chart_version,
5059 external_login_password_mz_system,
5060 license_key: license_key.clone(),
5061 },
5062 })
5063 .await?;
5064
5065 let catalog_upper = catalog.current_upper().await;
5068 boot_ts = std::cmp::max(boot_ts, catalog_upper);
5069
5070 if !read_only_controllers {
5071 epoch_millis_oracle.apply_write(boot_ts).await;
5072 }
5073
5074 info!(
5075 "startup: coordinator init: catalog open complete in {:?}",
5076 catalog_open_start.elapsed()
5077 );
5078
5079 let coord_thread_start = Instant::now();
5080 info!("startup: coordinator init: coordinator thread start beginning");
5081
5082 let session_id = catalog.config().session_id;
5083 let start_instant = catalog.config().start_instant;
5084
5085 let (bootstrap_tx, bootstrap_rx) = oneshot::channel();
5089 let handle = TokioHandle::current();
5090
5091 let metrics = Metrics::register_into(&metrics_registry);
5092 let metrics_clone = metrics.clone();
5093 let optimizer_metrics = OptimizerMetrics::register_into(
5094 &metrics_registry,
5095 catalog.system_config().optimizer_e2e_latency_warning_threshold(),
5096 );
5097 let segment_client_clone = segment_client.clone();
5098 let coord_now = now.clone();
5099 let advance_timelines_interval =
5100 tokio::time::interval(catalog.system_config().default_timestamp_interval());
5101
5102 let clusters_caught_up_check_interval = if read_only_controllers {
5103 let dyncfgs = catalog.system_config().dyncfgs();
5104 let interval = WITH_0DT_DEPLOYMENT_CAUGHT_UP_CHECK_INTERVAL.get(dyncfgs);
5105
5106 let mut interval = tokio::time::interval(interval);
5107 interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
5108 interval
5109 } else {
5110 let mut interval = tokio::time::interval(Duration::from_secs(60 * 60));
5118 interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
5119 interval
5120 };
5121
5122 let clusters_caught_up_check =
5123 clusters_caught_up_trigger.map(|trigger| {
5124 let mut exclude_collections: BTreeSet<GlobalId> =
5125 new_builtin_collections.iter().copied().collect();
5126
5127 let new_builtin_mvs = new_builtin_collections
5142 .iter()
5143 .map(|global_id| {
5144 catalog
5145 .state()
5146 .try_get_entry_by_global_id(global_id)
5147 .expect("new builtin collections have catalog entries")
5148 })
5149 .filter(|entry| entry.is_materialized_view())
5150 .map(|entry| entry.id());
5151 let mut todo: Vec<_> = migrated_storage_collections_0dt
5152 .iter()
5153 .copied()
5154 .filter(|id| catalog.state().get_entry(id).is_materialized_view())
5155 .chain(new_builtin_mvs)
5156 .collect();
5157 while let Some(item_id) = todo.pop() {
5158 let entry = catalog.state().get_entry(&item_id);
5159 exclude_collections.extend(entry.global_ids());
5160 todo.extend_from_slice(entry.used_by());
5161 }
5162
5163 CaughtUpCheckContext {
5164 trigger,
5165 exclude_collections,
5166 cluster_stability: BTreeMap::new(),
5167 }
5168 });
5169
5170 if let Some(TimestampOracleConfig::Postgres(pg_config)) =
5171 timestamp_oracle_config.as_ref()
5172 {
5173 let pg_timestamp_oracle_params =
5176 flags::timestamp_oracle_config(catalog.system_config());
5177 pg_timestamp_oracle_params.apply(pg_config);
5178 }
5179
5180 let connection_limit_callback: Arc<dyn Fn(&SystemVars) + Send + Sync> =
5183 Arc::new(move |system_vars: &SystemVars| {
5184 let limit: u64 = system_vars.max_connections().cast_into();
5185 let superuser_reserved: u64 =
5186 system_vars.superuser_reserved_connections().cast_into();
5187
5188 let superuser_reserved = if superuser_reserved >= limit {
5193 tracing::warn!(
5194 "superuser_reserved ({superuser_reserved}) is greater than max connections ({limit})!"
5195 );
5196 limit
5197 } else {
5198 superuser_reserved
5199 };
5200
5201 (connection_limit_callback)(limit, superuser_reserved);
5202 });
5203 catalog.system_config_mut().register_callback(
5204 &mz_sql::session::vars::MAX_CONNECTIONS,
5205 Arc::clone(&connection_limit_callback),
5206 );
5207 catalog.system_config_mut().register_callback(
5208 &mz_sql::session::vars::SUPERUSER_RESERVED_CONNECTIONS,
5209 connection_limit_callback,
5210 );
5211
5212 let (group_commit_tx, group_commit_rx) = appends::notifier();
5213
5214 let parent_span = tracing::Span::current();
5215 let thread = thread::Builder::new()
5216 .stack_size(3 * stack::STACK_SIZE)
5220 .name("coordinator".to_string())
5221 .spawn(move || {
5222 let span = info_span!(parent: parent_span, "coord::coordinator").entered();
5223
5224 let controller = handle
5225 .block_on({
5226 catalog.initialize_controller(
5227 controller_config,
5228 controller_envd_epoch,
5229 read_only_controllers,
5230 )
5231 })
5232 .unwrap_or_terminate("failed to initialize storage_controller");
5233 let catalog_upper = handle.block_on(catalog.current_upper());
5236 boot_ts = std::cmp::max(boot_ts, catalog_upper);
5237 if !read_only_controllers {
5238 let epoch_millis_oracle = ×tamp_oracles
5239 .get(&Timeline::EpochMilliseconds)
5240 .expect("inserted above")
5241 .oracle;
5242 handle.block_on(epoch_millis_oracle.apply_write(boot_ts));
5243 }
5244
5245 let catalog = Arc::new(catalog);
5246 let max_concurrent_occ_writes =
5249 usize::cast_from(catalog.system_config().max_concurrent_occ_writes());
5250 let frontend_read_then_write_enabled = {
5251 FRONTEND_READ_THEN_WRITE.get(catalog.system_config().dyncfgs())
5252 };
5253
5254 let caching_secrets_reader = CachingSecretsReader::new(secrets_controller.reader());
5255 let (group_committer_tx, group_committer_rx) = mpsc::unbounded_channel();
5256 let mut coord = Coordinator {
5257 controller,
5258 catalog,
5259 internal_cmd_tx,
5260 group_commit_tx,
5261 reconcile_now: Arc::new(Notify::new()),
5262 group_committer_tx,
5263 strict_serializable_reads_tx,
5264 linearize_reads_notify: Arc::new(Notify::new()),
5265 global_timelines: timestamp_oracles,
5266 transient_id_gen: Arc::new(TransientIdGen::new()),
5267 active_conns: BTreeMap::new(),
5268 txn_read_holds: Default::default(),
5269 pending_peeks: BTreeMap::new(),
5270 client_pending_peeks: BTreeMap::new(),
5271 pending_linearize_read_txns: BTreeMap::new(),
5272 serialized_ddl: LockedVecDeque::new(),
5273 active_compute_sinks: BTreeMap::new(),
5274 active_webhooks: BTreeMap::new(),
5275 active_copies: BTreeMap::new(),
5276 connection_cancel_watches: BTreeMap::new(),
5277 introspection_subscribes: BTreeMap::new(),
5278 write_locks: BTreeMap::new(),
5279 deferred_write_ops: BTreeMap::new(),
5280 pending_writes: Vec::new(),
5281 occ_write_semaphore: Arc::new(Semaphore::new(max_concurrent_occ_writes)),
5282 frontend_read_then_write_enabled,
5283 advance_timelines_interval,
5284 secrets_controller,
5285 caching_secrets_reader,
5286 cloud_resource_controller,
5287 storage_usage_client,
5288 storage_usage_collection_interval,
5289 segment_client,
5290 metrics,
5291 catalog_info_metrics_registry: metrics_registry.clone(),
5292 scoped_frontend: None,
5293 optimizer_metrics,
5294 tracing_handle,
5295 statement_logging: StatementLogging::new(coord_now.clone()),
5296 webhook_concurrency_limit,
5297 timestamp_oracle_config,
5298 caught_up_check_interval: clusters_caught_up_check_interval,
5299 caught_up_check: clusters_caught_up_check,
5300 installed_watch_sets: BTreeMap::new(),
5301 connection_watch_sets: BTreeMap::new(),
5302 cluster_replica_statuses: ClusterReplicaStatuses::new(),
5303 read_only_controllers,
5304 buffered_builtin_table_updates: Some(Vec::new()),
5305 license_key,
5306 user_id_pool: IdPool::empty(),
5307 persist_client,
5308 };
5309
5310 handle.block_on(async {
5312 appends::spawn_group_committer(
5313 group_committer_rx,
5314 coord.get_local_timestamp_oracle(),
5315 coord.controller.storage.table_write_handle(),
5316 coord.catalog().upper_handle(),
5317 coord.internal_cmd_tx.clone(),
5318 coord.catalog().config().now.clone(),
5319 coord.metrics.clone(),
5320 coord.catalog().system_config().dyncfgs(),
5321 );
5322 });
5323
5324 let bootstrap = handle.block_on(async {
5325 coord
5326 .bootstrap(
5327 boot_ts,
5328 migrated_storage_collections_0dt,
5329 builtin_table_updates,
5330 cached_global_exprs,
5331 uncached_local_exprs,
5332 )
5333 .await?;
5334 coord
5335 .controller
5336 .remove_orphaned_replicas(
5337 coord.catalog().get_next_user_replica_id().await?,
5338 coord.catalog().get_next_system_replica_id().await?,
5339 )
5340 .await
5341 .map_err(AdapterError::Orchestrator)?;
5342
5343 if let Some(retention_period) = storage_usage_retention_period {
5344 coord
5345 .prune_storage_usage_events_on_startup(retention_period)
5346 .await;
5347 }
5348
5349 coord.prune_arrangement_sizes_history_on_startup().await;
5350
5351 Ok(())
5352 });
5353 let ok = bootstrap.is_ok();
5354 drop(span);
5355 bootstrap_tx
5356 .send(bootstrap)
5357 .expect("bootstrap_rx is not dropped until it receives this message");
5358 if ok {
5359 handle.block_on(coord.serve(
5360 internal_cmd_rx,
5361 strict_serializable_reads_rx,
5362 cmd_rx,
5363 group_commit_rx,
5364 ));
5365 }
5366 })
5367 .expect("failed to create coordinator thread");
5368 match bootstrap_rx
5369 .await
5370 .expect("bootstrap_tx always sends a message or panics/halts")
5371 {
5372 Ok(()) => {
5373 info!(
5374 "startup: coordinator init: coordinator thread start complete in {:?}",
5375 coord_thread_start.elapsed()
5376 );
5377 info!(
5378 "startup: coordinator init: complete in {:?}",
5379 coord_start.elapsed()
5380 );
5381 let handle = Handle {
5382 session_id,
5383 start_instant,
5384 _thread: thread.join_on_drop(),
5385 };
5386 let client = Client::new(
5387 build_info,
5388 cmd_tx,
5389 metrics_clone,
5390 now,
5391 environment_id,
5392 segment_client_clone,
5393 );
5394 Ok((handle, client))
5395 }
5396 Err(e) => Err(e),
5397 }
5398 }
5399 .boxed()
5400}
5401
5402async fn get_initial_oracle_timestamps(
5416 timestamp_oracle_config: &Option<TimestampOracleConfig>,
5417) -> Result<BTreeMap<Timeline, Timestamp>, AdapterError> {
5418 let mut initial_timestamps = BTreeMap::new();
5419
5420 if let Some(config) = timestamp_oracle_config {
5421 let oracle_timestamps = config.get_all_timelines().await?;
5422
5423 let debug_msg = || {
5424 oracle_timestamps
5425 .iter()
5426 .map(|(timeline, ts)| format!("{:?} -> {}", timeline, ts))
5427 .join(", ")
5428 };
5429 info!(
5430 "current timestamps from the timestamp oracle: {}",
5431 debug_msg()
5432 );
5433
5434 for (timeline, ts) in oracle_timestamps {
5435 let entry = initial_timestamps
5436 .entry(Timeline::from_str(&timeline).expect("could not parse timeline"));
5437
5438 entry
5439 .and_modify(|current_ts| *current_ts = std::cmp::max(*current_ts, ts))
5440 .or_insert(ts);
5441 }
5442 } else {
5443 info!("no timestamp oracle configured!");
5444 };
5445
5446 let debug_msg = || {
5447 initial_timestamps
5448 .iter()
5449 .map(|(timeline, ts)| format!("{:?}: {}", timeline, ts))
5450 .join(", ")
5451 };
5452 info!("initial oracle timestamps: {}", debug_msg());
5453
5454 Ok(initial_timestamps)
5455}
5456
5457#[instrument]
5458pub async fn load_remote_system_parameters(
5459 storage: &mut Box<dyn OpenableDurableCatalogState>,
5460 system_parameter_sync_config: Option<SystemParameterSyncConfig>,
5461 system_parameter_sync_timeout: Duration,
5462) -> Result<Option<BTreeMap<String, String>>, AdapterError> {
5463 if let Some(system_parameter_sync_config) = system_parameter_sync_config {
5464 tracing::info!("parameter sync on boot: start sync");
5465
5466 let mut params = SynchronizedParameters::new(SystemVars::default());
5506 let frontend_sync = async {
5507 let frontend = SystemParameterFrontend::from(&system_parameter_sync_config).await?;
5508 frontend.pull(&mut params);
5509 let ops = params
5510 .modified()
5511 .into_iter()
5512 .map(|param| {
5513 let name = param.name;
5514 let value = param.value;
5515 tracing::info!(name, value, initial = true, "sync parameter");
5516 (name, value)
5517 })
5518 .collect();
5519 tracing::info!("parameter sync on boot: end sync");
5520 Ok(Some(ops))
5521 };
5522 if !storage.has_system_config_synced_once().await? {
5523 frontend_sync.await
5524 } else {
5525 match mz_ore::future::timeout(system_parameter_sync_timeout, frontend_sync).await {
5526 Ok(ops) => Ok(ops),
5527 Err(TimeoutError::Inner(e)) => Err(e),
5528 Err(TimeoutError::DeadlineElapsed) => {
5529 tracing::info!("parameter sync on boot: sync has timed out");
5530 Ok(None)
5531 }
5532 }
5533 }
5534 } else {
5535 Ok(None)
5536 }
5537}
5538
5539#[derive(Debug)]
5540pub enum WatchSetResponse {
5541 StatementDependenciesReady(StatementLoggingId, StatementLifecycleEvent),
5542 AlterSinkReady(AlterSinkReadyContext),
5543 AlterMaterializedViewReady(AlterMaterializedViewReadyContext),
5544}
5545
5546#[derive(Debug)]
5547pub struct AlterSinkReadyContext {
5548 ctx: Option<ExecuteContext>,
5549 otel_ctx: OpenTelemetryContext,
5550 plan: AlterSinkPlan,
5551 plan_validity: PlanValidity,
5552 read_hold: ReadHolds,
5553}
5554
5555impl AlterSinkReadyContext {
5556 fn ctx(&mut self) -> &mut ExecuteContext {
5557 self.ctx.as_mut().expect("only cleared on drop")
5558 }
5559
5560 fn retire(mut self, result: Result<ExecuteResponse, AdapterError>) {
5561 self.ctx
5562 .take()
5563 .expect("only cleared on drop")
5564 .retire(result);
5565 }
5566}
5567
5568impl Drop for AlterSinkReadyContext {
5569 fn drop(&mut self) {
5570 if let Some(ctx) = self.ctx.take() {
5571 ctx.retire(Err(AdapterError::Canceled));
5572 }
5573 }
5574}
5575
5576#[derive(Debug)]
5577pub struct AlterMaterializedViewReadyContext {
5578 ctx: Option<ExecuteContext>,
5579 otel_ctx: OpenTelemetryContext,
5580 plan: plan::AlterMaterializedViewApplyReplacementPlan,
5581 plan_validity: PlanValidity,
5582}
5583
5584impl AlterMaterializedViewReadyContext {
5585 fn ctx(&mut self) -> &mut ExecuteContext {
5586 self.ctx.as_mut().expect("only cleared on drop")
5587 }
5588
5589 fn retire(mut self, result: Result<ExecuteResponse, AdapterError>) {
5590 self.ctx
5591 .take()
5592 .expect("only cleared on drop")
5593 .retire(result);
5594 }
5595}
5596
5597impl Drop for AlterMaterializedViewReadyContext {
5598 fn drop(&mut self) {
5599 if let Some(ctx) = self.ctx.take() {
5600 ctx.retire(Err(AdapterError::Canceled));
5601 }
5602 }
5603}
5604
5605#[derive(Debug)]
5608struct LockedVecDeque<T> {
5609 items: VecDeque<T>,
5610 lock: Arc<tokio::sync::Mutex<()>>,
5611}
5612
5613impl<T> LockedVecDeque<T> {
5614 pub fn new() -> Self {
5615 Self {
5616 items: VecDeque::new(),
5617 lock: Arc::new(tokio::sync::Mutex::new(())),
5618 }
5619 }
5620
5621 pub fn try_lock_owned(&self) -> Result<OwnedMutexGuard<()>, tokio::sync::TryLockError> {
5622 Arc::clone(&self.lock).try_lock_owned()
5623 }
5624
5625 pub fn is_empty(&self) -> bool {
5626 self.items.is_empty()
5627 }
5628
5629 pub fn push_back(&mut self, value: T) {
5630 self.items.push_back(value)
5631 }
5632
5633 pub fn pop_front(&mut self) -> Option<T> {
5634 self.items.pop_front()
5635 }
5636
5637 pub fn remove(&mut self, index: usize) -> Option<T> {
5638 self.items.remove(index)
5639 }
5640
5641 pub fn iter(&self) -> std::collections::vec_deque::Iter<'_, T> {
5642 self.items.iter()
5643 }
5644}
5645
5646#[derive(Debug)]
5647struct DeferredPlanStatement {
5648 ctx: ExecuteContext,
5649 ps: PlanStatement,
5650}
5651
5652#[derive(Debug)]
5653enum PlanStatement {
5654 Statement {
5655 stmt: Arc<Statement<Raw>>,
5656 params: Params,
5657 },
5658 Plan {
5659 plan: mz_sql::plan::Plan,
5660 resolved_ids: ResolvedIds,
5661 sql_impl_resolved_ids: ResolvedIds,
5662 },
5663}
5664
5665#[derive(Debug, Error)]
5666pub enum NetworkPolicyError {
5667 #[error("Access denied for address {0}")]
5668 AddressDenied(IpAddr),
5669 #[error("Access denied missing IP address")]
5670 MissingIp,
5671}
5672
5673pub(crate) fn validate_ip_with_policy_rules(
5674 ip: &IpAddr,
5675 rules: &Vec<NetworkPolicyRule>,
5676) -> Result<(), NetworkPolicyError> {
5677 if rules.iter().any(|r| r.address.0.contains(ip)) {
5680 Ok(())
5681 } else {
5682 Err(NetworkPolicyError::AddressDenied(ip.clone()))
5683 }
5684}
5685
5686pub(crate) fn infer_sql_type_for_catalog(
5687 hir_expr: &HirRelationExpr,
5688 mir_expr: &MirRelationExpr,
5689) -> SqlRelationType {
5690 let mut typ = hir_expr.top_level_typ();
5691 typ.backport_nullability_and_keys(&mir_expr.typ());
5692 typ
5693}
5694
5695#[cfg(test)]
5696mod execute_context_tests {
5697 use tokio::sync::{mpsc, oneshot};
5698
5699 use super::*;
5700 use crate::session::Session;
5701 use crate::util::ClientTransmitter;
5702
5703 #[mz_ore::test]
5706 fn test_retire_answers_client_when_runtime_shuts_down() {
5707 let runtime = tokio::runtime::Runtime::new().expect("can build runtime");
5708
5709 let (client_tx, mut client_rx) = oneshot::channel();
5710 let (internal_cmd_tx, _internal_cmd_rx) = mpsc::unbounded_channel();
5711
5712 runtime.block_on(async {
5713 let ctx = ExecuteContext::from_parts_with_response_barriers(
5714 ClientTransmitter::new(client_tx, internal_cmd_tx.clone()),
5715 internal_cmd_tx,
5716 Session::dummy(),
5717 ExecuteContextGuard::default(),
5718 vec![Box::pin(std::future::pending())],
5720 );
5721 ctx.retire(Ok(ExecuteResponse::StartedTransaction));
5722 });
5723
5724 drop(runtime);
5725
5726 let response = client_rx.try_recv().expect("client must be answered");
5727 assert!(
5728 matches!(response.result, Err(AdapterError::Internal(_))),
5729 "expected an internal error, got {:?}",
5730 response.result
5731 );
5732 }
5733}
5734
5735#[cfg(test)]
5736mod id_pool_tests {
5737 use super::IdPool;
5738
5739 #[mz_ore::test]
5740 fn test_empty_pool() {
5741 let mut pool = IdPool::empty();
5742 assert_eq!(pool.remaining(), 0);
5743 assert_eq!(pool.allocate(), None);
5744 assert_eq!(pool.allocate_many(1), None);
5745 }
5746
5747 #[mz_ore::test]
5748 fn test_allocate_single() {
5749 let mut pool = IdPool::empty();
5750 pool.refill(10, 13);
5751 assert_eq!(pool.remaining(), 3);
5752 assert_eq!(pool.allocate(), Some(10));
5753 assert_eq!(pool.allocate(), Some(11));
5754 assert_eq!(pool.allocate(), Some(12));
5755 assert_eq!(pool.remaining(), 0);
5756 assert_eq!(pool.allocate(), None);
5757 }
5758
5759 #[mz_ore::test]
5760 fn test_allocate_many() {
5761 let mut pool = IdPool::empty();
5762 pool.refill(100, 105);
5763 assert_eq!(pool.allocate_many(3), Some(vec![100, 101, 102]));
5764 assert_eq!(pool.remaining(), 2);
5765 assert_eq!(pool.allocate_many(3), None);
5767 assert_eq!(pool.allocate_many(2), Some(vec![103, 104]));
5769 assert_eq!(pool.remaining(), 0);
5770 }
5771
5772 #[mz_ore::test]
5773 fn test_allocate_many_zero() {
5774 let mut pool = IdPool::empty();
5775 pool.refill(1, 5);
5776 assert_eq!(pool.allocate_many(0), Some(vec![]));
5777 assert_eq!(pool.remaining(), 4);
5778 }
5779
5780 #[mz_ore::test]
5781 fn test_refill_resets_pool() {
5782 let mut pool = IdPool::empty();
5783 pool.refill(0, 2);
5784 assert_eq!(pool.allocate(), Some(0));
5785 pool.refill(50, 52);
5787 assert_eq!(pool.allocate(), Some(50));
5788 assert_eq!(pool.allocate(), Some(51));
5789 assert_eq!(pool.allocate(), None);
5790 }
5791
5792 #[mz_ore::test]
5793 fn test_mixed_allocate_and_allocate_many() {
5794 let mut pool = IdPool::empty();
5795 pool.refill(0, 10);
5796 assert_eq!(pool.allocate(), Some(0));
5797 assert_eq!(pool.allocate_many(3), Some(vec![1, 2, 3]));
5798 assert_eq!(pool.allocate(), Some(4));
5799 assert_eq!(pool.remaining(), 5);
5800 }
5801
5802 #[mz_ore::test]
5803 #[should_panic(expected = "invalid pool range")]
5804 fn test_refill_invalid_range_panics() {
5805 let mut pool = IdPool::empty();
5806 pool.refill(10, 5);
5807 }
5808}
5809
5810#[cfg(test)]
5811mod arrangement_sizes_pruner_tests {
5812 use mz_repr::catalog_item_id::CatalogItemId;
5813 use mz_repr::{Datum, Row};
5814
5815 use super::arrangement_sizes_expired_retractions;
5816
5817 fn history_row(ts_ms: i64) -> Row {
5821 let dt = mz_ore::now::to_datetime(ts_ms.try_into().expect("non-negative"));
5822 Row::pack_slice(&[
5823 Datum::String("r1"),
5824 Datum::String("u1"),
5825 Datum::Int64(123),
5826 Datum::TimestampTz(dt.try_into().expect("fits in TimestampTz")),
5827 ])
5828 }
5829
5830 fn item_id() -> CatalogItemId {
5831 CatalogItemId::User(42)
5833 }
5834
5835 #[mz_ore::test]
5836 fn empty_input_produces_no_retractions() {
5837 let out = arrangement_sizes_expired_retractions(Vec::new(), 1_000, item_id());
5838 assert!(out.is_empty());
5839 }
5840
5841 #[mz_ore::test]
5842 fn retracts_only_rows_strictly_before_cutoff() {
5843 let rows = vec![
5846 (history_row(100), 1),
5847 (history_row(500), 1),
5848 (history_row(1_000), 1), (history_row(5_000), 1),
5850 ];
5851 let out = arrangement_sizes_expired_retractions(rows, 1_000, item_id());
5852 assert_eq!(out.len(), 2);
5853 }
5854
5855 #[mz_ore::test]
5856 #[should_panic(expected = "consolidated contents should not contain retractions")]
5857 fn retraction_in_input_panics() {
5858 let rows = vec![(history_row(100), -1)];
5859 let _ = arrangement_sizes_expired_retractions(rows, 1_000, item_id());
5860 }
5861}