Skip to main content

mz_adapter/
coord.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Translation of SQL commands into timestamped `Controller` commands.
11//!
12//! The various SQL commands instruct the system to take actions that are not
13//! yet explicitly timestamped. On the other hand, the underlying data continually
14//! change as time moves forward. On the third hand, we greatly benefit from the
15//! information that some times are no longer of interest, so that we may
16//! compact the representation of the continually changing collections.
17//!
18//! The [`Coordinator`] curates these interactions by observing the progress
19//! collections make through time, choosing timestamps for its own commands,
20//! and eventually communicating that certain times have irretrievably "passed".
21//!
22//! ## Frontiers another way
23//!
24//! If the above description of frontiers left you with questions, this
25//! repackaged explanation might help.
26//!
27//! - `since` is the least recent time (i.e. oldest time) that you can read
28//!   from sources and be guaranteed that the returned data is accurate as of
29//!   that time.
30//!
31//!   Reads at times less than `since` may return values that were not actually
32//!   seen at the specified time, but arrived later (i.e. the results are
33//!   compacted).
34//!
35//!   For correctness' sake, the coordinator never chooses to read at a time
36//!   less than an arrangement's `since`.
37//!
38//! - `upper` is the first time after the most recent time that you can read
39//!   from sources and receive an immediate response. Alternately, it is the
40//!   least time at which the data may still change (that is the reason we may
41//!   not be able to respond immediately).
42//!
43//!   Reads at times >= `upper` may not immediately return because the answer
44//!   isn't known yet. However, once the `upper` is > the specified read time,
45//!   the read can return.
46//!
47//!   For the sake of returned values' freshness, the coordinator prefers
48//!   performing reads at an arrangement's `upper`. However, because we more
49//!   strongly prefer correctness, the coordinator will choose timestamps
50//!   greater than an object's `upper` if it is also being accessed alongside
51//!   objects whose `since` times are >= its `upper`.
52//!
53//! This illustration attempts to show, with time moving left to right, the
54//! relationship between `since` and `upper`.
55//!
56//! - `#`: possibly inaccurate results
57//! - `-`: immediate, correct response
58//! - `?`: not yet known
59//! - `s`: since
60//! - `u`: upper
61//! - `|`: eligible for coordinator to select
62//!
63//! ```nofmt
64//! ####s----u?????
65//!     |||||||||||
66//! ```
67//!
68
69use 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::{
95    ENABLE_SCOPED_SYSTEM_PARAMETERS, USER_ID_POOL_BATCH_SIZE,
96    WITH_0DT_DEPLOYMENT_CAUGHT_UP_CHECK_INTERVAL,
97};
98use mz_auth::password::Password;
99use mz_build_info::BuildInfo;
100use mz_catalog::builtin::{BUILTINS, BUILTINS_STATIC, MZ_STORAGE_USAGE_BY_SHARD};
101use mz_catalog::config::{AwsPrincipalContext, BuiltinItemMigrationConfig, ClusterReplicaSizeMap};
102use mz_catalog::durable::OpenableDurableCatalogState;
103use mz_catalog::expr_cache::{GlobalExpressions, LocalExpressions};
104use mz_catalog::memory::objects::{
105    CatalogEntry, CatalogItem, ClusterReplicaProcessStatus, ClusterVariantManaged, Connection,
106    DataSourceDesc, Table, TableDataSource,
107};
108use mz_cloud_resources::{CloudResourceController, VpcEndpointConfig, VpcEndpointEvent};
109use mz_compute_client::as_of_selection;
110use mz_compute_client::controller::error::{
111    CollectionLookupError, CollectionMissing, DataflowCreationError, InstanceMissing,
112};
113use mz_compute_types::ComputeInstanceId;
114use mz_compute_types::dataflows::DataflowDescription;
115use mz_compute_types::plan::LirRelationExpr;
116use mz_controller::clusters::{
117    ClusterConfig, ClusterEvent, ClusterStatus, ProcessId, ReplicaLocation,
118};
119use mz_controller::{ControllerConfig, Readiness};
120use mz_controller_types::{ClusterId, ReplicaId, WatchSetId};
121use mz_dyncfg::{ConfigUpdates, ParameterScope};
122use mz_expr::{MapFilterProject, MirRelationExpr, OptimizedMirRelationExpr, RowSetFinishing};
123use mz_license_keys::{ExpirationBehavior, ValidatedLicenseKey};
124use mz_orchestrator::OfflineReason;
125use mz_ore::cast::{CastFrom, CastInto, CastLossy};
126use mz_ore::channel::trigger::Trigger;
127use mz_ore::future::TimeoutError;
128use mz_ore::metrics::MetricsRegistry;
129use mz_ore::now::{EpochMillis, NowFn};
130use mz_ore::task::{JoinHandle, spawn};
131use mz_ore::thread::JoinHandleExt;
132use mz_ore::tracing::{OpenTelemetryContext, TracingHandle};
133use mz_ore::url::SensitiveUrl;
134use mz_ore::{
135    assert_none, instrument, soft_assert_eq_or_log, soft_assert_or_log, soft_panic_or_log, stack,
136};
137use mz_persist_client::PersistClient;
138use mz_persist_client::batch::ProtoBatch;
139use mz_persist_client::usage::{ShardsUsageReferenced, StorageUsageClient};
140use mz_repr::adt::numeric::Numeric;
141use mz_repr::explain::{ExplainConfig, ExplainFormat};
142use mz_repr::global_id::TransientIdGen;
143use mz_repr::optimize::{OptimizerFeatureOverrides, OptimizerFeatures, OverrideFrom};
144use mz_repr::role_id::RoleId;
145use mz_repr::{CatalogItemId, Diff, GlobalId, RelationDesc, SqlRelationType, Timestamp};
146use mz_secrets::cache::CachingSecretsReader;
147use mz_secrets::{SecretsController, SecretsReader};
148use mz_sql::ast::{Raw, Statement};
149use mz_sql::catalog::{CatalogCluster, EnvironmentId};
150use mz_sql::names::{QualifiedItemName, ResolvedIds, SchemaSpecifier};
151use mz_sql::optimizer_metrics::OptimizerMetrics;
152use mz_sql::plan::{
153    self, AlterSinkPlan, ConnectionDetails, CreateConnectionPlan, HirRelationExpr,
154    NetworkPolicyRule, OnTimeoutAction, Params, QueryWhen,
155};
156use mz_sql::session::user::User;
157use mz_sql::session::vars::{MAX_CREDIT_CONSUMPTION_RATE, SystemVars, Var};
158use mz_sql_parser::ast::ExplainStage;
159use mz_sql_parser::ast::display::AstDisplay;
160use mz_storage_client::client::TableData;
161use mz_storage_client::controller::{CollectionDescription, DataSource, ExportDescription};
162use mz_storage_types::connections::Connection as StorageConnection;
163use mz_storage_types::connections::ConnectionContext;
164use mz_storage_types::connections::inline::{IntoInlineConnection, ReferencedConnection};
165use mz_storage_types::read_holds::ReadHold;
166use mz_storage_types::sinks::{S3SinkFormat, StorageSinkDesc};
167use mz_storage_types::sources::kafka::KAFKA_PROGRESS_DESC;
168use mz_storage_types::sources::{IngestionDescription, SourceExport, Timeline};
169use mz_timestamp_oracle::{TimestampOracleConfig, WriteTimestamp};
170use mz_transform::dataflow::DataflowMetainfo;
171use opentelemetry::trace::TraceContextExt;
172use serde::Serialize;
173use thiserror::Error;
174use timely::progress::{Antichain, Timestamp as _};
175use tokio::runtime::Handle as TokioHandle;
176use tokio::select;
177use tokio::sync::{Notify, OwnedMutexGuard, mpsc, oneshot, watch};
178use tokio::time::{Interval, MissedTickBehavior};
179use tracing::{Instrument, Level, Span, debug, info, info_span, span, warn};
180use tracing_opentelemetry::OpenTelemetrySpanExt;
181use uuid::Uuid;
182
183use crate::active_compute_sink::{ActiveComputeSink, ActiveCopyFrom};
184use crate::catalog::{BuiltinTableUpdate, Catalog, OpenCatalogResult};
185use crate::client::{Client, Handle};
186use crate::command::{Command, ExecuteResponse};
187use crate::config::{
188    ClusterEvalContext, ReplicaEvalContext, ScopedParameters, ScopedParametersScope,
189    SynchronizedParameters, SystemParameterFrontend, SystemParameterSyncConfig,
190};
191use crate::coord::appends::{
192    BuiltinTableAppendCompletion, BuiltinTableAppendNotify, DeferredOp, GroupCommitPermit,
193    PendingWriteTxn,
194};
195use crate::coord::caught_up::CaughtUpCheckContext;
196use crate::coord::cluster_scheduling::SchedulingDecision;
197use crate::coord::id_bundle::CollectionIdBundle;
198use crate::coord::introspection::IntrospectionSubscribe;
199use crate::coord::peek::PendingPeek;
200use crate::coord::statement_logging::StatementLogging;
201use crate::coord::timeline::{TimelineContext, TimelineState};
202use crate::coord::timestamp_selection::{TimestampContext, TimestampDetermination};
203use crate::coord::validity::PlanValidity;
204use crate::error::AdapterError;
205use crate::explain::insights::PlanInsightsContext;
206use crate::explain::optimizer_trace::{DispatchGuard, OptimizerTrace};
207use crate::metrics::Metrics;
208use crate::optimize::dataflows::{ComputeInstanceSnapshot, DataflowBuilder};
209use crate::optimize::{self, Optimize, OptimizerConfig};
210use crate::session::{EndTransactionAction, Session};
211use crate::statement_logging::{
212    StatementEndedExecutionReason, StatementLifecycleEvent, StatementLoggingId,
213};
214use crate::util::{ClientTransmitter, ResultExt, sort_topological};
215use crate::webhook::{WebhookAppenderInvalidator, WebhookConcurrencyLimiter};
216use crate::{AdapterNotice, ReadHolds, flags};
217
218pub(crate) mod appends;
219pub(crate) mod catalog_serving;
220pub(crate) mod cluster_controller;
221pub(crate) mod cluster_scheduling;
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/// A pool of pre-allocated user IDs to avoid per-DDL persist writes.
247///
248/// IDs in the range `[next, upper)` are available for allocation.
249/// When exhausted, the pool must be refilled via the catalog.
250///
251/// # Correctness
252///
253/// The pool is owned by [`Coordinator`], which processes all requests
254/// on a single-threaded event loop. Because every access requires
255/// `&mut self` on the coordinator, there is no concurrent access to the
256/// pool — no additional synchronization is needed.
257///
258/// Global ID uniqueness is guaranteed because each refill calls
259/// [`Catalog::allocate_user_ids`], which performs a durable persist
260/// write that atomically reserves the entire batch before any IDs from
261/// it are handed out. If the process crashes after a refill but before
262/// all pre-allocated IDs are consumed, the unused IDs form harmless
263/// gaps in the sequence — user IDs are not required to be contiguous.
264///
265/// This guarantee holds even if multiple `environmentd` processes run
266/// concurrently. Each process has its own independent pool,
267/// but every refill goes through the shared persist-backed catalog,
268/// which serializes allocations across all callers. Two processes
269/// will therefore never receive overlapping ID ranges,
270/// for the same reason they could not before this pool existed.
271#[derive(Debug)]
272pub(crate) struct IdPool {
273    next: u64,
274    upper: u64,
275}
276
277impl IdPool {
278    /// Creates an empty pool.
279    pub fn empty() -> Self {
280        IdPool { next: 0, upper: 0 }
281    }
282
283    /// Allocates a single ID from the pool, returning `None` if exhausted.
284    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    /// Allocates `n` consecutive IDs from the pool, returning `None` if
295    /// insufficient IDs remain.
296    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    /// Returns the number of IDs remaining in the pool.
307    pub fn remaining(&self) -> u64 {
308        self.upper - self.next
309    }
310
311    /// Refills the pool with the given range `[next, upper)`.
312    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)]
320pub enum Message {
321    Command(OpenTelemetryContext, Command),
322    ControllerReady {
323        controller: ControllerReadiness,
324    },
325    PurifiedStatementReady(PurifiedStatementReady),
326    CreateConnectionValidationReady(CreateConnectionValidationReady),
327    AlterConnectionValidationReady(AlterConnectionValidationReady),
328    TryDeferred {
329        /// The connection that created this op.
330        conn_id: ConnectionId,
331        /// The write lock that notified us our deferred op might be able to run.
332        ///
333        /// Note: While we never want to hold a partial set of locks, it can be important to hold
334        /// onto the _one_ that notified us our op might be ready. If there are multiple operations
335        /// waiting on a single collection, and we don't hold this lock through retyring the op,
336        /// then everything waiting on this collection will get retried causing traffic in the
337        /// Coordinator's message queue.
338        ///
339        /// See [`DeferredOp::can_be_optimistically_retried`] for more detail.
340        acquired_lock: Option<(CatalogItemId, tokio::sync::OwnedMutexGuard<()>)>,
341    },
342    /// Initiates a group commit.
343    GroupCommitInitiate(Span, Option<GroupCommitPermit>),
344    DeferredStatementReady,
345    AdvanceTimelines,
346    ClusterEvent(ClusterEvent),
347    CancelPendingPeeks {
348        conn_id: ConnectionId,
349    },
350    LinearizeReads,
351    StagedBatches {
352        conn_id: ConnectionId,
353        table_id: CatalogItemId,
354        batches: Vec<Result<ProtoBatch, String>>,
355    },
356    StorageUsageSchedule,
357    StorageUsageFetch,
358    StorageUsageUpdate(ShardsUsageReferenced),
359    StorageUsagePrune(Vec<BuiltinTableUpdate>),
360    ArrangementSizesSchedule,
361    ArrangementSizesSnapshot,
362    ArrangementSizesPrune(Vec<BuiltinTableUpdate>),
363    /// Performs any cleanup and logging actions necessary for
364    /// finalizing a statement execution.
365    RetireExecute {
366        data: ExecuteContextExtra,
367        otel_ctx: OpenTelemetryContext,
368        reason: StatementEndedExecutionReason,
369    },
370    ExecuteSingleStatementTransaction {
371        ctx: ExecuteContext,
372        otel_ctx: OpenTelemetryContext,
373        stmt: Arc<Statement<Raw>>,
374        params: mz_sql::plan::Params,
375    },
376    PeekStageReady {
377        ctx: ExecuteContext,
378        span: Span,
379        stage: PeekStage,
380    },
381    CreateIndexStageReady {
382        ctx: ExecuteContext,
383        span: Span,
384        stage: CreateIndexStage,
385    },
386    CreateViewStageReady {
387        ctx: ExecuteContext,
388        span: Span,
389        stage: CreateViewStage,
390    },
391    CreateMaterializedViewStageReady {
392        ctx: ExecuteContext,
393        span: Span,
394        stage: CreateMaterializedViewStage,
395    },
396    SubscribeStageReady {
397        ctx: ExecuteContext,
398        span: Span,
399        stage: SubscribeStage,
400    },
401    IntrospectionSubscribeStageReady {
402        span: Span,
403        stage: IntrospectionSubscribeStage,
404    },
405    SecretStageReady {
406        ctx: ExecuteContext,
407        span: Span,
408        stage: SecretStage,
409    },
410    ClusterStageReady {
411        ctx: ExecuteContext,
412        span: Span,
413        stage: ClusterStage,
414    },
415    ExplainTimestampStageReady {
416        ctx: ExecuteContext,
417        span: Span,
418        stage: ExplainTimestampStage,
419    },
420    DrainStatementLog,
421    PrivateLinkVpcEndpointEvents(Vec<VpcEndpointEvent>),
422    CheckSchedulingPolicies,
423
424    /// Scheduling policy decisions about turning clusters On/Off.
425    /// `Vec<(policy name, Vec of decisions by the policy)>`
426    /// A cluster will be On if and only if there is at least one On decision for it.
427    /// Scheduling decisions for clusters that have `SCHEDULE = MANUAL` are ignored.
428    SchedulingDecisions(Vec<(&'static str, Vec<(ClusterId, SchedulingDecision)>)>),
429
430    /// One pull/apply call from the cluster controller task, answered on the main
431    /// coordinator message loop from the catalog and live controller signals.
432    /// See [`cluster_controller`].
433    ClusterControllerRequest(cluster_controller::ClusterControllerRequest),
434}
435
436impl Message {
437    /// Returns a string to identify the kind of [`Message`], useful for logging.
438    pub const fn kind(&self) -> &'static str {
439        match self {
440            Message::Command(_, msg) => match msg {
441                Command::CatalogSnapshot { .. } => "command-catalog_snapshot",
442                Command::Startup { .. } => "command-startup",
443                Command::Execute { .. } => "command-execute",
444                Command::Commit { .. } => "command-commit",
445                Command::CancelRequest { .. } => "command-cancel_request",
446                Command::PrivilegedCancelRequest { .. } => "command-privileged_cancel_request",
447                Command::GetWebhook { .. } => "command-get_webhook",
448                Command::GetSystemVars { .. } => "command-get_system_vars",
449                Command::SetSystemVars { .. } => "command-set_system_vars",
450                Command::UpdateScopedSystemParameters { .. } => {
451                    "command-update_scoped_system_parameters"
452                }
453                Command::InstallScopedSystemParameterFrontend { .. } => {
454                    "command-install_scoped_system_parameter_frontend"
455                }
456                Command::Terminate { .. } => "command-terminate",
457                Command::RetireExecute { .. } => "command-retire_execute",
458                Command::CheckConsistency { .. } => "command-check_consistency",
459                Command::Dump { .. } => "command-dump",
460                Command::AuthenticatePassword { .. } => "command-auth_check",
461                Command::AuthenticateGetSASLChallenge { .. } => "command-auth_get_sasl_challenge",
462                Command::AuthenticateVerifySASLProof { .. } => "command-auth_verify_sasl_proof",
463                Command::CheckRoleCanLogin { .. } => "command-check_role_can_login",
464                Command::GetComputeInstanceClient { .. } => "get-compute-instance-client",
465                Command::GetOracle { .. } => "get-oracle",
466                Command::DetermineRealTimeRecentTimestamp { .. } => {
467                    "determine-real-time-recent-timestamp"
468                }
469                Command::GetTransactionReadHoldsBundle { .. } => {
470                    "get-transaction-read-holds-bundle"
471                }
472                Command::StoreTransactionReadHolds { .. } => "store-transaction-read-holds",
473                Command::ExecuteSlowPathPeek { .. } => "execute-slow-path-peek",
474                Command::ExecuteSubscribe { .. } => "execute-subscribe",
475                Command::CopyToPreflight { .. } => "copy-to-preflight",
476                Command::ExecuteCopyTo { .. } => "execute-copy-to",
477                Command::ExecuteSideEffectingFunc { .. } => "execute-side-effecting-func",
478                Command::LookupConnection { .. } => "lookup-connection",
479                Command::RegisterFrontendPeek { .. } => "register-frontend-peek",
480                Command::UnregisterFrontendPeek { .. } => "unregister-frontend-peek",
481                Command::ExplainTimestamp { .. } => "explain-timestamp",
482                Command::FrontendStatementLogging(..) => "frontend-statement-logging",
483                Command::StartCopyFromStdin { .. } => "start-copy-from-stdin",
484                Command::InjectAuditEvents { .. } => "inject-audit-events",
485            },
486            Message::ControllerReady {
487                controller: ControllerReadiness::Compute,
488            } => "controller_ready(compute)",
489            Message::ControllerReady {
490                controller: ControllerReadiness::Storage,
491            } => "controller_ready(storage)",
492            Message::ControllerReady {
493                controller: ControllerReadiness::Metrics,
494            } => "controller_ready(metrics)",
495            Message::ControllerReady {
496                controller: ControllerReadiness::Internal,
497            } => "controller_ready(internal)",
498            Message::PurifiedStatementReady(_) => "purified_statement_ready",
499            Message::CreateConnectionValidationReady(_) => "create_connection_validation_ready",
500            Message::TryDeferred { .. } => "try_deferred",
501            Message::GroupCommitInitiate(..) => "group_commit_initiate",
502            Message::AdvanceTimelines => "advance_timelines",
503            Message::ClusterEvent(_) => "cluster_event",
504            Message::CancelPendingPeeks { .. } => "cancel_pending_peeks",
505            Message::LinearizeReads => "linearize_reads",
506            Message::StagedBatches { .. } => "staged_batches",
507            Message::StorageUsageSchedule => "storage_usage_schedule",
508            Message::StorageUsageFetch => "storage_usage_fetch",
509            Message::StorageUsageUpdate(_) => "storage_usage_update",
510            Message::StorageUsagePrune(_) => "storage_usage_prune",
511            Message::ArrangementSizesSchedule => "arrangement_sizes_schedule",
512            Message::ArrangementSizesSnapshot => "arrangement_sizes_snapshot",
513            Message::ArrangementSizesPrune(_) => "arrangement_sizes_prune",
514            Message::RetireExecute { .. } => "retire_execute",
515            Message::ExecuteSingleStatementTransaction { .. } => {
516                "execute_single_statement_transaction"
517            }
518            Message::PeekStageReady { .. } => "peek_stage_ready",
519            Message::ExplainTimestampStageReady { .. } => "explain_timestamp_stage_ready",
520            Message::CreateIndexStageReady { .. } => "create_index_stage_ready",
521            Message::CreateViewStageReady { .. } => "create_view_stage_ready",
522            Message::CreateMaterializedViewStageReady { .. } => {
523                "create_materialized_view_stage_ready"
524            }
525            Message::SubscribeStageReady { .. } => "subscribe_stage_ready",
526            Message::IntrospectionSubscribeStageReady { .. } => {
527                "introspection_subscribe_stage_ready"
528            }
529            Message::SecretStageReady { .. } => "secret_stage_ready",
530            Message::ClusterStageReady { .. } => "cluster_stage_ready",
531            Message::DrainStatementLog => "drain_statement_log",
532            Message::AlterConnectionValidationReady(..) => "alter_connection_validation_ready",
533            Message::PrivateLinkVpcEndpointEvents(_) => "private_link_vpc_endpoint_events",
534            Message::CheckSchedulingPolicies => "check_scheduling_policies",
535            Message::SchedulingDecisions { .. } => "scheduling_decision",
536            Message::ClusterControllerRequest(_) => "cluster_controller_request",
537            Message::DeferredStatementReady => "deferred_statement_ready",
538        }
539    }
540}
541
542/// The reason for why a controller needs processing on the main loop.
543#[derive(Debug)]
544pub enum ControllerReadiness {
545    /// The storage controller is ready.
546    Storage,
547    /// The compute controller is ready.
548    Compute,
549    /// A batch of metric data is ready.
550    Metrics,
551    /// An internally-generated message is ready to be returned.
552    Internal,
553}
554
555#[derive(Derivative)]
556#[derivative(Debug)]
557pub struct BackgroundWorkResult<T> {
558    #[derivative(Debug = "ignore")]
559    pub ctx: ExecuteContext,
560    pub result: Result<T, AdapterError>,
561    pub params: Params,
562    pub plan_validity: PlanValidity,
563    pub original_stmt: Arc<Statement<Raw>>,
564    pub otel_ctx: OpenTelemetryContext,
565}
566
567pub type PurifiedStatementReady = BackgroundWorkResult<mz_sql::pure::PurifiedStatement>;
568
569#[derive(Derivative)]
570#[derivative(Debug)]
571pub struct ValidationReady<T> {
572    #[derivative(Debug = "ignore")]
573    pub ctx: ExecuteContext,
574    pub result: Result<T, AdapterError>,
575    pub resolved_ids: ResolvedIds,
576    pub connection_id: CatalogItemId,
577    pub connection_gid: GlobalId,
578    pub plan_validity: PlanValidity,
579    pub otel_ctx: OpenTelemetryContext,
580}
581
582pub type CreateConnectionValidationReady = ValidationReady<CreateConnectionPlan>;
583pub type AlterConnectionValidationReady = ValidationReady<Connection>;
584
585#[derive(Debug)]
586pub enum PeekStage {
587    /// Common stages across SELECT, EXPLAIN and COPY TO queries.
588    LinearizeTimestamp(PeekStageLinearizeTimestamp),
589    RealTimeRecency(PeekStageRealTimeRecency),
590    TimestampReadHold(PeekStageTimestampReadHold),
591    Optimize(PeekStageOptimize),
592    /// Final stage for a peek.
593    Finish(PeekStageFinish),
594    /// Final stage for an explain.
595    ExplainPlan(PeekStageExplainPlan),
596    ExplainPushdown(PeekStageExplainPushdown),
597    /// Preflight checks for a copy to operation.
598    CopyToPreflight(PeekStageCopyTo),
599    /// Final stage for a copy to which involves shipping the dataflow.
600    CopyToDataflow(PeekStageCopyTo),
601}
602
603#[derive(Debug)]
604pub struct CopyToContext {
605    /// The `RelationDesc` of the data to be copied.
606    pub desc: RelationDesc,
607    /// The destination uri of the external service where the data will be copied.
608    pub uri: Uri,
609    /// Connection information required to connect to the external service to copy the data.
610    pub connection: StorageConnection<ReferencedConnection>,
611    /// The ID of the CONNECTION object to be used for copying the data.
612    pub connection_id: CatalogItemId,
613    /// Format params to format the data.
614    pub format: S3SinkFormat,
615    /// Approximate max file size of each uploaded file.
616    pub max_file_size: u64,
617    /// Number of batches the output of the COPY TO will be partitioned into
618    /// to distribute the load across workers deterministically.
619    /// This is only an option since it's not set when CopyToContext is instantiated
620    /// but immediately after in the PeekStageValidate stage.
621    pub output_batch_count: Option<u64>,
622}
623
624#[derive(Debug)]
625pub struct PeekStageLinearizeTimestamp {
626    validity: PlanValidity,
627    plan: mz_sql::plan::SelectPlan,
628    max_query_result_size: Option<u64>,
629    source_ids: BTreeSet<GlobalId>,
630    target_replica: Option<ReplicaId>,
631    timeline_context: TimelineContext,
632    optimizer: optimize::PeekOptimizer,
633    /// An optional context set iff the state machine is initiated from
634    /// sequencing an EXPLAIN for this statement.
635    explain_ctx: ExplainContext,
636}
637
638#[derive(Debug)]
639pub struct PeekStageRealTimeRecency {
640    validity: PlanValidity,
641    plan: mz_sql::plan::SelectPlan,
642    max_query_result_size: Option<u64>,
643    source_ids: BTreeSet<GlobalId>,
644    target_replica: Option<ReplicaId>,
645    timeline_context: TimelineContext,
646    oracle_read_ts: Option<Timestamp>,
647    optimizer: optimize::PeekOptimizer,
648    /// An optional context set iff the state machine is initiated from
649    /// sequencing an EXPLAIN for this statement.
650    explain_ctx: ExplainContext,
651}
652
653#[derive(Debug)]
654pub struct PeekStageTimestampReadHold {
655    validity: PlanValidity,
656    plan: mz_sql::plan::SelectPlan,
657    max_query_result_size: Option<u64>,
658    source_ids: BTreeSet<GlobalId>,
659    target_replica: Option<ReplicaId>,
660    timeline_context: TimelineContext,
661    oracle_read_ts: Option<Timestamp>,
662    real_time_recency_ts: Option<mz_repr::Timestamp>,
663    optimizer: optimize::PeekOptimizer,
664    /// An optional context set iff the state machine is initiated from
665    /// sequencing an EXPLAIN for this statement.
666    explain_ctx: ExplainContext,
667}
668
669#[derive(Debug)]
670pub struct PeekStageOptimize {
671    validity: PlanValidity,
672    plan: mz_sql::plan::SelectPlan,
673    max_query_result_size: Option<u64>,
674    source_ids: BTreeSet<GlobalId>,
675    id_bundle: CollectionIdBundle,
676    target_replica: Option<ReplicaId>,
677    determination: TimestampDetermination,
678    optimizer: optimize::PeekOptimizer,
679    /// An optional context set iff the state machine is initiated from
680    /// sequencing an EXPLAIN for this statement.
681    explain_ctx: ExplainContext,
682}
683
684#[derive(Debug)]
685pub struct PeekStageFinish {
686    validity: PlanValidity,
687    plan: mz_sql::plan::SelectPlan,
688    max_query_result_size: Option<u64>,
689    id_bundle: CollectionIdBundle,
690    target_replica: Option<ReplicaId>,
691    source_ids: BTreeSet<GlobalId>,
692    determination: TimestampDetermination,
693    cluster_id: ComputeInstanceId,
694    finishing: RowSetFinishing,
695    /// When present, an optimizer trace to be used for emitting a plan insights
696    /// notice.
697    plan_insights_optimizer_trace: Option<OptimizerTrace>,
698    insights_ctx: Option<Box<PlanInsightsContext>>,
699    global_lir_plan: optimize::peek::GlobalLirPlan,
700    optimization_finished_at: EpochMillis,
701}
702
703#[derive(Debug)]
704pub struct PeekStageCopyTo {
705    validity: PlanValidity,
706    optimizer: optimize::copy_to::Optimizer,
707    global_lir_plan: optimize::copy_to::GlobalLirPlan,
708    optimization_finished_at: EpochMillis,
709    target_replica: Option<ReplicaId>,
710    source_ids: BTreeSet<GlobalId>,
711}
712
713#[derive(Debug)]
714pub struct PeekStageExplainPlan {
715    validity: PlanValidity,
716    optimizer: optimize::peek::Optimizer,
717    df_meta: DataflowMetainfo,
718    explain_ctx: ExplainPlanContext,
719    insights_ctx: Option<Box<PlanInsightsContext>>,
720}
721
722#[derive(Debug)]
723pub struct PeekStageExplainPushdown {
724    validity: PlanValidity,
725    determination: TimestampDetermination,
726    imports: BTreeMap<GlobalId, MapFilterProject>,
727}
728
729#[derive(Debug)]
730pub enum CreateIndexStage {
731    Optimize(CreateIndexOptimize),
732    Finish(CreateIndexFinish),
733    Explain(CreateIndexExplain),
734}
735
736#[derive(Debug)]
737pub struct CreateIndexOptimize {
738    validity: PlanValidity,
739    plan: plan::CreateIndexPlan,
740    resolved_ids: ResolvedIds,
741    /// An optional context set iff the state machine is initiated from
742    /// sequencing an EXPLAIN for this statement.
743    explain_ctx: ExplainContext,
744}
745
746#[derive(Debug)]
747pub struct CreateIndexFinish {
748    validity: PlanValidity,
749    item_id: CatalogItemId,
750    global_id: GlobalId,
751    plan: plan::CreateIndexPlan,
752    resolved_ids: ResolvedIds,
753    global_mir_plan: optimize::index::GlobalMirPlan,
754    global_lir_plan: optimize::index::GlobalLirPlan,
755    optimizer_features: OptimizerFeatures,
756}
757
758#[derive(Debug)]
759pub struct CreateIndexExplain {
760    validity: PlanValidity,
761    exported_index_id: GlobalId,
762    plan: plan::CreateIndexPlan,
763    df_meta: DataflowMetainfo,
764    explain_ctx: ExplainPlanContext,
765}
766
767#[derive(Debug)]
768pub enum CreateViewStage {
769    Optimize(CreateViewOptimize),
770    Finish(CreateViewFinish),
771    Explain(CreateViewExplain),
772}
773
774#[derive(Debug)]
775pub struct CreateViewOptimize {
776    validity: PlanValidity,
777    plan: plan::CreateViewPlan,
778    resolved_ids: ResolvedIds,
779    /// An optional context set iff the state machine is initiated from
780    /// sequencing an EXPLAIN for this statement.
781    explain_ctx: ExplainContext,
782}
783
784#[derive(Debug)]
785pub struct CreateViewFinish {
786    validity: PlanValidity,
787    /// ID of this item in the Catalog.
788    item_id: CatalogItemId,
789    /// ID by with Compute will reference this View.
790    global_id: GlobalId,
791    plan: plan::CreateViewPlan,
792    /// IDs of objects resolved during name resolution.
793    resolved_ids: ResolvedIds,
794    optimized_expr: OptimizedMirRelationExpr,
795}
796
797#[derive(Debug)]
798pub struct CreateViewExplain {
799    validity: PlanValidity,
800    id: GlobalId,
801    plan: plan::CreateViewPlan,
802    explain_ctx: ExplainPlanContext,
803}
804
805#[derive(Debug)]
806pub enum ExplainTimestampStage {
807    Optimize(ExplainTimestampOptimize),
808    RealTimeRecency(ExplainTimestampRealTimeRecency),
809    LinearizeTimestamp(ExplainTimestampLinearizeTimestamp),
810    Finish(ExplainTimestampFinish),
811}
812
813#[derive(Debug)]
814pub struct ExplainTimestampOptimize {
815    validity: PlanValidity,
816    plan: plan::ExplainTimestampPlan,
817    cluster_id: ClusterId,
818}
819
820#[derive(Debug)]
821pub struct ExplainTimestampRealTimeRecency {
822    validity: PlanValidity,
823    format: ExplainFormat,
824    optimized_plan: OptimizedMirRelationExpr,
825    cluster_id: ClusterId,
826    when: QueryWhen,
827}
828
829#[derive(Debug)]
830pub struct ExplainTimestampLinearizeTimestamp {
831    validity: PlanValidity,
832    format: ExplainFormat,
833    optimized_plan: OptimizedMirRelationExpr,
834    cluster_id: ClusterId,
835    source_ids: BTreeSet<GlobalId>,
836    when: QueryWhen,
837    real_time_recency_ts: Option<Timestamp>,
838}
839
840#[derive(Debug)]
841pub struct ExplainTimestampFinish {
842    validity: PlanValidity,
843    format: ExplainFormat,
844    cluster_id: ClusterId,
845    source_ids: BTreeSet<GlobalId>,
846    when: QueryWhen,
847    real_time_recency_ts: Option<Timestamp>,
848    /// The timeline context derived in the preceding `LinearizeTimestamp`
849    /// stage, carried forward so it stays consistent with `oracle_read_ts`.
850    timeline_context: TimelineContext,
851    /// The linearized read timestamp, read off the coordinator loop in the
852    /// preceding `LinearizeTimestamp` stage. `None` when no linearized read is
853    /// needed.
854    oracle_read_ts: Option<Timestamp>,
855}
856
857#[derive(Debug)]
858pub enum ClusterStage {
859    Alter(AlterCluster),
860    WaitForHydrated(AlterClusterWaitForHydrated),
861    Finalize(AlterClusterFinalize),
862}
863
864#[derive(Debug)]
865pub struct AlterCluster {
866    validity: PlanValidity,
867    plan: plan::AlterClusterPlan,
868}
869
870#[derive(Debug)]
871pub struct AlterClusterWaitForHydrated {
872    validity: PlanValidity,
873    plan: plan::AlterClusterPlan,
874    new_config: ClusterVariantManaged,
875    workload_class: Option<String>,
876    timeout_time: Instant,
877    on_timeout: OnTimeoutAction,
878}
879
880#[derive(Debug)]
881pub struct AlterClusterFinalize {
882    validity: PlanValidity,
883    plan: plan::AlterClusterPlan,
884    new_config: ClusterVariantManaged,
885    workload_class: Option<String>,
886}
887
888#[derive(Debug)]
889pub enum ExplainContext {
890    /// The ordinary, non-explain variant of the statement.
891    None,
892    /// The `EXPLAIN <level> PLAN FOR <explainee>` version of the statement.
893    Plan(ExplainPlanContext),
894    /// Generate a notice containing the `EXPLAIN PLAN INSIGHTS` output
895    /// alongside the query's normal output.
896    PlanInsightsNotice(OptimizerTrace),
897    /// `EXPLAIN FILTER PUSHDOWN`
898    Pushdown,
899}
900
901impl ExplainContext {
902    /// If available for this context, wrap the [`OptimizerTrace`] into a
903    /// [`tracing::Dispatch`] and set it as default, returning the resulting
904    /// guard in a `Some(guard)` option.
905    pub(crate) fn dispatch_guard(&self) -> Option<DispatchGuard<'_>> {
906        let optimizer_trace = match self {
907            ExplainContext::Plan(explain_ctx) => Some(&explain_ctx.optimizer_trace),
908            ExplainContext::PlanInsightsNotice(optimizer_trace) => Some(optimizer_trace),
909            _ => None,
910        };
911        optimizer_trace.map(|optimizer_trace| optimizer_trace.as_guard())
912    }
913
914    pub(crate) fn needs_cluster(&self) -> bool {
915        match self {
916            ExplainContext::None => true,
917            ExplainContext::Plan(..) => false,
918            ExplainContext::PlanInsightsNotice(..) => true,
919            ExplainContext::Pushdown => false,
920        }
921    }
922
923    pub(crate) fn needs_plan_insights(&self) -> bool {
924        matches!(
925            self,
926            ExplainContext::Plan(ExplainPlanContext {
927                stage: ExplainStage::PlanInsights,
928                ..
929            }) | ExplainContext::PlanInsightsNotice(_)
930        )
931    }
932}
933
934#[derive(Debug)]
935pub struct ExplainPlanContext {
936    /// EXPLAIN BROKEN is internal syntax for showing EXPLAIN output despite an internal error in
937    /// the optimizer: we don't immediately bail out from peek sequencing when an internal optimizer
938    /// error happens, but go on with trying to show the requested EXPLAIN stage. This can still
939    /// succeed if the requested EXPLAIN stage is before the point where the error happened.
940    pub broken: bool,
941    pub config: ExplainConfig,
942    pub format: ExplainFormat,
943    pub stage: ExplainStage,
944    pub replan: Option<GlobalId>,
945    pub desc: Option<RelationDesc>,
946    pub optimizer_trace: OptimizerTrace,
947}
948
949#[derive(Debug)]
950pub enum CreateMaterializedViewStage {
951    Optimize(CreateMaterializedViewOptimize),
952    Finish(CreateMaterializedViewFinish),
953    Explain(CreateMaterializedViewExplain),
954}
955
956#[derive(Debug)]
957pub struct CreateMaterializedViewOptimize {
958    validity: PlanValidity,
959    plan: plan::CreateMaterializedViewPlan,
960    resolved_ids: ResolvedIds,
961    /// An optional context set iff the state machine is initiated from
962    /// sequencing an EXPLAIN for this statement.
963    explain_ctx: ExplainContext,
964}
965
966#[derive(Debug)]
967pub struct CreateMaterializedViewFinish {
968    /// The ID of this Materialized View in the Catalog.
969    item_id: CatalogItemId,
970    /// The ID of the durable pTVC backing this Materialized View.
971    global_id: GlobalId,
972    validity: PlanValidity,
973    plan: plan::CreateMaterializedViewPlan,
974    resolved_ids: ResolvedIds,
975    local_mir_plan: optimize::materialized_view::LocalMirPlan,
976    global_mir_plan: optimize::materialized_view::GlobalMirPlan,
977    global_lir_plan: optimize::materialized_view::GlobalLirPlan,
978    optimizer_features: OptimizerFeatures,
979}
980
981#[derive(Debug)]
982pub struct CreateMaterializedViewExplain {
983    global_id: GlobalId,
984    validity: PlanValidity,
985    plan: plan::CreateMaterializedViewPlan,
986    df_meta: DataflowMetainfo,
987    explain_ctx: ExplainPlanContext,
988}
989
990#[derive(Debug)]
991pub enum SubscribeStage {
992    OptimizeMir(SubscribeOptimizeMir),
993    LinearizeTimestamp(SubscribeLinearizeTimestamp),
994    TimestampOptimizeLir(SubscribeTimestampOptimizeLir),
995    Finish(SubscribeFinish),
996    Explain(SubscribeExplain),
997}
998
999#[derive(Debug)]
1000pub struct SubscribeOptimizeMir {
1001    validity: PlanValidity,
1002    plan: plan::SubscribePlan,
1003    timeline: TimelineContext,
1004    dependency_ids: BTreeSet<GlobalId>,
1005    cluster_id: ComputeInstanceId,
1006    replica_id: Option<ReplicaId>,
1007    /// An optional context set iff the state machine is initiated from
1008    /// sequencing an EXPLAIN for this statement.
1009    explain_ctx: ExplainContext,
1010}
1011
1012#[derive(Debug)]
1013pub struct SubscribeLinearizeTimestamp {
1014    validity: PlanValidity,
1015    plan: plan::SubscribePlan,
1016    timeline: TimelineContext,
1017    optimizer: optimize::subscribe::Optimizer,
1018    global_mir_plan: optimize::subscribe::GlobalMirPlan<optimize::subscribe::Unresolved>,
1019    dependency_ids: BTreeSet<GlobalId>,
1020    replica_id: Option<ReplicaId>,
1021    /// An optional context set iff the state machine is initiated from
1022    /// sequencing an EXPLAIN for this statement.
1023    explain_ctx: ExplainContext,
1024}
1025
1026#[derive(Debug)]
1027pub struct SubscribeTimestampOptimizeLir {
1028    validity: PlanValidity,
1029    plan: plan::SubscribePlan,
1030    timeline: TimelineContext,
1031    optimizer: optimize::subscribe::Optimizer,
1032    global_mir_plan: optimize::subscribe::GlobalMirPlan<optimize::subscribe::Unresolved>,
1033    dependency_ids: BTreeSet<GlobalId>,
1034    replica_id: Option<ReplicaId>,
1035    /// The linearized read timestamp, read off the coordinator loop in the
1036    /// preceding `LinearizeTimestamp` stage. `None` when no linearized read is
1037    /// needed.
1038    oracle_read_ts: Option<Timestamp>,
1039    /// An optional context set iff the state machine is initiated from
1040    /// sequencing an EXPLAIN for this statement.
1041    explain_ctx: ExplainContext,
1042}
1043
1044#[derive(Debug)]
1045pub struct SubscribeFinish {
1046    validity: PlanValidity,
1047    cluster_id: ComputeInstanceId,
1048    replica_id: Option<ReplicaId>,
1049    plan: plan::SubscribePlan,
1050    global_lir_plan: optimize::subscribe::GlobalLirPlan,
1051    dependency_ids: BTreeSet<GlobalId>,
1052}
1053
1054#[derive(Debug)]
1055pub struct SubscribeExplain {
1056    validity: PlanValidity,
1057    optimizer: optimize::subscribe::Optimizer,
1058    df_meta: DataflowMetainfo,
1059    cluster_id: ComputeInstanceId,
1060    explain_ctx: ExplainPlanContext,
1061}
1062
1063#[derive(Debug)]
1064pub enum IntrospectionSubscribeStage {
1065    OptimizeMir(IntrospectionSubscribeOptimizeMir),
1066    TimestampOptimizeLir(IntrospectionSubscribeTimestampOptimizeLir),
1067    Finish(IntrospectionSubscribeFinish),
1068}
1069
1070#[derive(Debug)]
1071pub struct IntrospectionSubscribeOptimizeMir {
1072    validity: PlanValidity,
1073    plan: plan::SubscribePlan,
1074    subscribe_id: GlobalId,
1075    cluster_id: ComputeInstanceId,
1076    replica_id: ReplicaId,
1077}
1078
1079#[derive(Debug)]
1080pub struct IntrospectionSubscribeTimestampOptimizeLir {
1081    validity: PlanValidity,
1082    optimizer: optimize::subscribe::Optimizer,
1083    global_mir_plan: optimize::subscribe::GlobalMirPlan<optimize::subscribe::Unresolved>,
1084    cluster_id: ComputeInstanceId,
1085    replica_id: ReplicaId,
1086}
1087
1088#[derive(Debug)]
1089pub struct IntrospectionSubscribeFinish {
1090    validity: PlanValidity,
1091    global_lir_plan: optimize::subscribe::GlobalLirPlan,
1092    read_holds: ReadHolds,
1093    cluster_id: ComputeInstanceId,
1094    replica_id: ReplicaId,
1095}
1096
1097#[derive(Debug)]
1098pub enum SecretStage {
1099    CreateEnsure(CreateSecretEnsure),
1100    CreateFinish(CreateSecretFinish),
1101    RotateKeysEnsure(RotateKeysSecretEnsure),
1102    RotateKeysFinish(RotateKeysSecretFinish),
1103    Alter(AlterSecret),
1104}
1105
1106#[derive(Debug)]
1107pub struct CreateSecretEnsure {
1108    validity: PlanValidity,
1109    plan: plan::CreateSecretPlan,
1110}
1111
1112#[derive(Debug)]
1113pub struct CreateSecretFinish {
1114    validity: PlanValidity,
1115    item_id: CatalogItemId,
1116    global_id: GlobalId,
1117    plan: plan::CreateSecretPlan,
1118}
1119
1120#[derive(Debug)]
1121pub struct RotateKeysSecretEnsure {
1122    validity: PlanValidity,
1123    id: CatalogItemId,
1124}
1125
1126#[derive(Debug)]
1127pub struct RotateKeysSecretFinish {
1128    validity: PlanValidity,
1129    ops: Vec<crate::catalog::Op>,
1130}
1131
1132#[derive(Debug)]
1133pub struct AlterSecret {
1134    validity: PlanValidity,
1135    plan: plan::AlterSecretPlan,
1136}
1137
1138/// An enum describing which cluster to run a statement on.
1139///
1140/// One example usage would be that if a query depends only on system tables, we might
1141/// automatically run it on the catalog server cluster to benefit from indexes that exist there.
1142#[derive(Debug, Copy, Clone, PartialEq, Eq)]
1143pub enum TargetCluster {
1144    /// The catalog server cluster.
1145    CatalogServer,
1146    /// The current user's active cluster.
1147    Active,
1148    /// The cluster selected at the start of a transaction.
1149    Transaction(ClusterId),
1150}
1151
1152/// Result types for each stage of a sequence.
1153pub(crate) enum StageResult<T> {
1154    /// A task was spawned that will return the next stage.
1155    Handle(JoinHandle<Result<T, AdapterError>>),
1156    /// A task was spawned that will return a response for the client.
1157    HandleRetire(JoinHandle<Result<ExecuteResponse, AdapterError>>),
1158    /// The next stage is immediately ready and will execute.
1159    Immediate(T),
1160    /// The final stage was executed and is ready to respond to the client.
1161    Response(ExecuteResponse),
1162}
1163
1164/// Common functionality for [Coordinator::sequence_staged].
1165pub(crate) trait Staged: Send {
1166    type Ctx: StagedContext;
1167
1168    fn validity(&mut self) -> &mut PlanValidity;
1169
1170    /// Returns the next stage or final result.
1171    async fn stage(
1172        self,
1173        coord: &mut Coordinator,
1174        ctx: &mut Self::Ctx,
1175    ) -> Result<StageResult<Box<Self>>, AdapterError>;
1176
1177    /// Prepares a message for the Coordinator.
1178    fn message(self, ctx: Self::Ctx, span: Span) -> Message;
1179
1180    /// Whether it is safe to SQL cancel this stage.
1181    fn cancel_enabled(&self) -> bool;
1182}
1183
1184pub trait StagedContext {
1185    fn retire(self, result: Result<ExecuteResponse, AdapterError>);
1186    fn session(&self) -> Option<&Session>;
1187}
1188
1189impl StagedContext for ExecuteContext {
1190    fn retire(self, result: Result<ExecuteResponse, AdapterError>) {
1191        self.retire(result);
1192    }
1193
1194    fn session(&self) -> Option<&Session> {
1195        Some(self.session())
1196    }
1197}
1198
1199impl StagedContext for () {
1200    fn retire(self, _result: Result<ExecuteResponse, AdapterError>) {}
1201
1202    fn session(&self) -> Option<&Session> {
1203        None
1204    }
1205}
1206
1207/// Configures a coordinator.
1208pub struct Config {
1209    pub controller_config: ControllerConfig,
1210    pub controller_envd_epoch: NonZeroI64,
1211    pub storage: Box<dyn mz_catalog::durable::DurableCatalogState>,
1212    pub timestamp_oracle_url: Option<SensitiveUrl>,
1213    pub unsafe_mode: bool,
1214    pub all_features: bool,
1215    pub build_info: &'static BuildInfo,
1216    pub environment_id: EnvironmentId,
1217    pub metrics_registry: MetricsRegistry,
1218    pub now: NowFn,
1219    pub secrets_controller: Arc<dyn SecretsController>,
1220    pub cloud_resource_controller: Option<Arc<dyn CloudResourceController>>,
1221    pub availability_zones: Vec<String>,
1222    pub cluster_replica_sizes: ClusterReplicaSizeMap,
1223    pub builtin_system_cluster_config: BootstrapBuiltinClusterConfig,
1224    pub builtin_catalog_server_cluster_config: BootstrapBuiltinClusterConfig,
1225    pub builtin_probe_cluster_config: BootstrapBuiltinClusterConfig,
1226    pub builtin_support_cluster_config: BootstrapBuiltinClusterConfig,
1227    pub builtin_analytics_cluster_config: BootstrapBuiltinClusterConfig,
1228    pub system_parameter_defaults: BTreeMap<String, String>,
1229    pub storage_usage_client: StorageUsageClient,
1230    pub storage_usage_collection_interval: Duration,
1231    pub storage_usage_retention_period: Option<Duration>,
1232    pub segment_client: Option<mz_segment::Client>,
1233    pub egress_addresses: Vec<IpNet>,
1234    pub remote_system_parameters: Option<BTreeMap<String, String>>,
1235    pub aws_account_id: Option<String>,
1236    pub aws_privatelink_availability_zones: Option<Vec<String>>,
1237    pub connection_context: ConnectionContext,
1238    pub connection_limit_callback: Box<dyn Fn(u64, u64) -> () + Send + Sync + 'static>,
1239    pub webhook_concurrency_limit: WebhookConcurrencyLimiter,
1240    pub http_host_name: Option<String>,
1241    pub tracing_handle: TracingHandle,
1242    /// Whether or not to start controllers in read-only mode. This is only
1243    /// meant for use during development of read-only clusters and 0dt upgrades
1244    /// and should go away once we have proper orchestration during upgrades.
1245    pub read_only_controllers: bool,
1246
1247    /// A trigger that signals that the current deployment has caught up with a
1248    /// previous deployment. Only used during 0dt deployment, while in read-only
1249    /// mode.
1250    pub caught_up_trigger: Option<Trigger>,
1251
1252    pub helm_chart_version: Option<String>,
1253    pub license_key: ValidatedLicenseKey,
1254    pub external_login_password_mz_system: Option<Password>,
1255    pub force_builtin_schema_migration: Option<String>,
1256}
1257
1258/// Metadata about an active connection.
1259#[derive(Debug, Serialize)]
1260pub struct ConnMeta {
1261    /// Pgwire specifies that every connection have a 32-bit secret associated
1262    /// with it, that is known to both the client and the server. Cancellation
1263    /// requests are required to authenticate with the secret of the connection
1264    /// that they are targeting.
1265    secret_key: u32,
1266    /// The time when the session's connection was initiated.
1267    connected_at: EpochMillis,
1268    user: User,
1269    application_name: String,
1270    uuid: Uuid,
1271    conn_id: ConnectionId,
1272    client_ip: Option<IpAddr>,
1273
1274    /// Sinks that will need to be dropped when the current transaction, if
1275    /// any, is cleared.
1276    drop_sinks: BTreeSet<GlobalId>,
1277
1278    /// Lock for the Coordinator's deferred statements that is dropped on transaction clear.
1279    #[serde(skip)]
1280    deferred_lock: Option<OwnedMutexGuard<()>>,
1281
1282    /// Cluster reconfigurations that will need to be
1283    /// cleaned up when the current transaction is cleared
1284    pending_cluster_alters: BTreeSet<ClusterId>,
1285
1286    /// Channel on which to send notices to a session.
1287    #[serde(skip)]
1288    notice_tx: mpsc::UnboundedSender<AdapterNotice>,
1289
1290    /// The role that initiated the database context. Fixed for the duration of the connection.
1291    /// WARNING: This role reference is not updated when the role is dropped.
1292    /// Consumers should not assume that this role exist.
1293    authenticated_role: RoleId,
1294}
1295
1296impl ConnMeta {
1297    pub fn conn_id(&self) -> &ConnectionId {
1298        &self.conn_id
1299    }
1300
1301    pub fn user(&self) -> &User {
1302        &self.user
1303    }
1304
1305    pub fn application_name(&self) -> &str {
1306        &self.application_name
1307    }
1308
1309    pub fn authenticated_role_id(&self) -> &RoleId {
1310        &self.authenticated_role
1311    }
1312
1313    pub fn uuid(&self) -> Uuid {
1314        self.uuid
1315    }
1316
1317    pub fn client_ip(&self) -> Option<IpAddr> {
1318        self.client_ip
1319    }
1320
1321    pub fn connected_at(&self) -> EpochMillis {
1322        self.connected_at
1323    }
1324}
1325
1326#[derive(Debug)]
1327/// A pending transaction waiting to be committed.
1328pub struct PendingTxn {
1329    /// Context used to send a response back to the client.
1330    ctx: ExecuteContext,
1331    /// Client response for transaction.
1332    response: Result<PendingTxnResponse, AdapterError>,
1333    /// The action to take at the end of the transaction.
1334    action: EndTransactionAction,
1335}
1336
1337#[derive(Debug)]
1338/// The response we'll send for a [`PendingTxn`].
1339pub enum PendingTxnResponse {
1340    /// The transaction will be committed.
1341    Committed {
1342        /// Parameters that will change, and their values, once this transaction is complete.
1343        params: BTreeMap<&'static str, String>,
1344    },
1345    /// The transaction will be rolled back.
1346    Rolledback {
1347        /// Parameters that will change, and their values, once this transaction is complete.
1348        params: BTreeMap<&'static str, String>,
1349    },
1350}
1351
1352impl PendingTxnResponse {
1353    pub fn extend_params(&mut self, p: impl IntoIterator<Item = (&'static str, String)>) {
1354        match self {
1355            PendingTxnResponse::Committed { params }
1356            | PendingTxnResponse::Rolledback { params } => params.extend(p),
1357        }
1358    }
1359}
1360
1361impl From<PendingTxnResponse> for ExecuteResponse {
1362    fn from(value: PendingTxnResponse) -> Self {
1363        match value {
1364            PendingTxnResponse::Committed { params } => {
1365                ExecuteResponse::TransactionCommitted { params }
1366            }
1367            PendingTxnResponse::Rolledback { params } => {
1368                ExecuteResponse::TransactionRolledBack { params }
1369            }
1370        }
1371    }
1372}
1373
1374#[derive(Debug)]
1375/// A pending read transaction waiting to be linearized along with metadata about it's state
1376pub struct PendingReadTxn {
1377    /// The transaction type
1378    txn: PendingRead,
1379    /// The timestamp context of the transaction.
1380    timestamp_context: TimestampContext,
1381    /// When we created this pending txn, when the transaction ends. Only used for metrics.
1382    created: Instant,
1383    /// Number of times we requeued the processing of this pending read txn.
1384    /// Requeueing is necessary if the time we executed the query is after the current oracle time;
1385    /// see [`Coordinator::message_linearize_reads`] for more details.
1386    num_requeues: u64,
1387    /// Telemetry context.
1388    otel_ctx: OpenTelemetryContext,
1389}
1390
1391impl PendingReadTxn {
1392    /// Return the timestamp context of the pending read transaction.
1393    pub fn timestamp_context(&self) -> &TimestampContext {
1394        &self.timestamp_context
1395    }
1396
1397    pub(crate) fn take_context(self) -> ExecuteContext {
1398        self.txn.take_context()
1399    }
1400}
1401
1402#[derive(Debug)]
1403/// A pending read transaction waiting to be linearized.
1404enum PendingRead {
1405    Read {
1406        /// The inner transaction.
1407        txn: PendingTxn,
1408    },
1409    ReadThenWrite {
1410        /// Context used to send a response back to the client.
1411        ctx: ExecuteContext,
1412        /// Channel used to alert the transaction that the read has been linearized and send back
1413        /// `ctx`.
1414        tx: oneshot::Sender<Option<ExecuteContext>>,
1415    },
1416}
1417
1418impl PendingRead {
1419    /// Alert the client that the read has been linearized.
1420    ///
1421    /// If it is necessary to finalize an execute, return the state necessary to do so
1422    /// (execution context and result)
1423    #[instrument(level = "debug")]
1424    pub fn finish(self) -> Option<(ExecuteContext, Result<ExecuteResponse, AdapterError>)> {
1425        match self {
1426            PendingRead::Read {
1427                txn:
1428                    PendingTxn {
1429                        mut ctx,
1430                        response,
1431                        action,
1432                    },
1433                ..
1434            } => {
1435                let changed = ctx.session_mut().vars_mut().end_transaction(action);
1436                // Append any parameters that changed to the response.
1437                let response = response.map(|mut r| {
1438                    r.extend_params(changed);
1439                    ExecuteResponse::from(r)
1440                });
1441
1442                Some((ctx, response))
1443            }
1444            PendingRead::ReadThenWrite { ctx, tx, .. } => {
1445                // Ignore errors if the caller has hung up.
1446                let _ = tx.send(Some(ctx));
1447                None
1448            }
1449        }
1450    }
1451
1452    fn label(&self) -> &'static str {
1453        match self {
1454            PendingRead::Read { .. } => "read",
1455            PendingRead::ReadThenWrite { .. } => "read_then_write",
1456        }
1457    }
1458
1459    pub(crate) fn take_context(self) -> ExecuteContext {
1460        match self {
1461            PendingRead::Read { txn, .. } => txn.ctx,
1462            PendingRead::ReadThenWrite { ctx, tx, .. } => {
1463                // Inform the transaction that we've taken their context.
1464                // Ignore errors if the caller has hung up.
1465                let _ = tx.send(None);
1466                ctx
1467            }
1468        }
1469    }
1470}
1471
1472/// State that the coordinator must process as part of retiring
1473/// command execution.  `ExecuteContextExtra::Default` is guaranteed
1474/// to produce a value that will cause the coordinator to do nothing, and
1475/// is intended for use by code that invokes the execution processing flow
1476/// (i.e., `sequence_plan`) without actually being a statement execution.
1477///
1478/// This is a pure data struct containing only the statement logging ID.
1479/// For auto-retire-on-drop behavior, use `ExecuteContextGuard` which wraps
1480/// this struct and owns the channel for sending retirement messages.
1481#[derive(Debug, Default)]
1482#[must_use]
1483pub struct ExecuteContextExtra {
1484    statement_uuid: Option<StatementLoggingId>,
1485}
1486
1487impl ExecuteContextExtra {
1488    pub(crate) fn new(statement_uuid: Option<StatementLoggingId>) -> Self {
1489        Self { statement_uuid }
1490    }
1491    pub fn is_trivial(&self) -> bool {
1492        self.statement_uuid.is_none()
1493    }
1494    pub fn contents(&self) -> Option<StatementLoggingId> {
1495        self.statement_uuid
1496    }
1497    /// Consume this extra and return the statement UUID for retirement.
1498    /// This should only be called from code that knows what to do to finish
1499    /// up logging based on the inner value.
1500    #[must_use]
1501    pub(crate) fn retire(self) -> Option<StatementLoggingId> {
1502        self.statement_uuid
1503    }
1504}
1505
1506/// A guard that wraps `ExecuteContextExtra` and owns a channel for sending
1507/// retirement messages to the coordinator.
1508///
1509/// If this guard is dropped with a `Some` `statement_uuid` in its inner
1510/// `ExecuteContextExtra`, the `Drop` implementation will automatically send a
1511/// `Message::RetireExecute` to log the statement ending.
1512/// This handles cases like connection drops where the context cannot be
1513/// explicitly retired.
1514/// See <https://github.com/MaterializeInc/database-issues/issues/7304>
1515#[derive(Debug)]
1516#[must_use]
1517pub struct ExecuteContextGuard {
1518    extra: ExecuteContextExtra,
1519    /// Channel for sending messages to the coordinator. Used for auto-retiring on drop.
1520    /// For `Default` instances, this is a dummy sender (receiver already dropped), so
1521    /// sends will fail silently - which is the desired behavior since Default instances
1522    /// should only be used for non-logged statements.
1523    coordinator_tx: mpsc::UnboundedSender<Message>,
1524}
1525
1526impl Default for ExecuteContextGuard {
1527    fn default() -> Self {
1528        // Create a dummy sender by immediately dropping the receiver.
1529        // Any send on this channel will fail silently, which is the desired
1530        // behavior for Default instances (non-logged statements).
1531        let (tx, _rx) = mpsc::unbounded_channel();
1532        Self {
1533            extra: ExecuteContextExtra::default(),
1534            coordinator_tx: tx,
1535        }
1536    }
1537}
1538
1539impl ExecuteContextGuard {
1540    pub(crate) fn new(
1541        statement_uuid: Option<StatementLoggingId>,
1542        coordinator_tx: mpsc::UnboundedSender<Message>,
1543    ) -> Self {
1544        Self {
1545            extra: ExecuteContextExtra::new(statement_uuid),
1546            coordinator_tx,
1547        }
1548    }
1549    pub fn is_trivial(&self) -> bool {
1550        self.extra.is_trivial()
1551    }
1552    pub fn contents(&self) -> Option<StatementLoggingId> {
1553        self.extra.contents()
1554    }
1555    /// Take responsibility for the contents.  This should only be
1556    /// called from code that knows what to do to finish up logging
1557    /// based on the inner value.
1558    ///
1559    /// Returns the inner `ExecuteContextExtra`, consuming the guard without
1560    /// triggering the auto-retire behavior.
1561    pub(crate) fn defuse(mut self) -> ExecuteContextExtra {
1562        // Taking statement_uuid prevents the Drop impl from sending a retire message
1563        std::mem::take(&mut self.extra)
1564    }
1565}
1566
1567impl Drop for ExecuteContextGuard {
1568    fn drop(&mut self) {
1569        if let Some(statement_uuid) = self.extra.statement_uuid.take() {
1570            // Auto-retire since the guard was dropped without explicit retirement (likely due
1571            // to connection drop).
1572            let msg = Message::RetireExecute {
1573                data: ExecuteContextExtra {
1574                    statement_uuid: Some(statement_uuid),
1575                },
1576                otel_ctx: OpenTelemetryContext::obtain(),
1577                reason: StatementEndedExecutionReason::Aborted,
1578            };
1579            // Send may fail for Default instances (dummy sender), which is fine since
1580            // Default instances should only be used for non-logged statements.
1581            let _ = self.coordinator_tx.send(msg);
1582        }
1583    }
1584}
1585
1586/// Bundle of state related to statement execution.
1587///
1588/// This struct collects a bundle of state that needs to be threaded
1589/// through various functions as part of statement execution.
1590/// It is used to finalize execution, by calling `retire`. Finalizing execution
1591/// involves sending the session back to the pgwire layer so that it
1592/// may be used to process further commands. It also involves
1593/// performing some work on the main coordinator thread
1594/// (e.g., recording the time at which the statement finished
1595/// executing). The state necessary to perform this work is bundled in
1596/// the `ExecuteContextGuard` object.
1597#[derive(Debug)]
1598pub struct ExecuteContext {
1599    inner: Box<ExecuteContextInner>,
1600}
1601
1602impl std::ops::Deref for ExecuteContext {
1603    type Target = ExecuteContextInner;
1604    fn deref(&self) -> &Self::Target {
1605        &*self.inner
1606    }
1607}
1608
1609impl std::ops::DerefMut for ExecuteContext {
1610    fn deref_mut(&mut self) -> &mut Self::Target {
1611        &mut *self.inner
1612    }
1613}
1614
1615#[derive(Derivative)]
1616#[derivative(Debug)]
1617pub struct ExecuteContextInner {
1618    tx: ClientTransmitter<ExecuteResponse>,
1619    internal_cmd_tx: mpsc::UnboundedSender<Message>,
1620    session: Session,
1621    extra: ExecuteContextGuard,
1622    #[derivative(Debug = "ignore")]
1623    response_barriers: Vec<BuiltinTableAppendNotify>,
1624}
1625
1626impl ExecuteContext {
1627    pub fn session(&self) -> &Session {
1628        &self.session
1629    }
1630
1631    pub fn session_mut(&mut self) -> &mut Session {
1632        &mut self.session
1633    }
1634
1635    pub fn tx(&self) -> &ClientTransmitter<ExecuteResponse> {
1636        &self.tx
1637    }
1638
1639    pub fn tx_mut(&mut self) -> &mut ClientTransmitter<ExecuteResponse> {
1640        &mut self.tx
1641    }
1642
1643    pub fn from_parts(
1644        tx: ClientTransmitter<ExecuteResponse>,
1645        internal_cmd_tx: mpsc::UnboundedSender<Message>,
1646        session: Session,
1647        extra: ExecuteContextGuard,
1648    ) -> Self {
1649        Self::from_parts_with_response_barriers(tx, internal_cmd_tx, session, extra, Vec::new())
1650    }
1651
1652    pub fn from_parts_with_response_barriers(
1653        tx: ClientTransmitter<ExecuteResponse>,
1654        internal_cmd_tx: mpsc::UnboundedSender<Message>,
1655        session: Session,
1656        extra: ExecuteContextGuard,
1657        response_barriers: Vec<BuiltinTableAppendNotify>,
1658    ) -> Self {
1659        Self {
1660            inner: ExecuteContextInner {
1661                tx,
1662                session,
1663                extra,
1664                response_barriers,
1665                internal_cmd_tx,
1666            }
1667            .into(),
1668        }
1669    }
1670
1671    /// By calling this function, the caller takes responsibility for
1672    /// dealing with the instance of `ExecuteContextGuard`. This is
1673    /// intended to support protocols (like `COPY FROM`) that involve
1674    /// multiple passes of sending the session back and forth between
1675    /// the coordinator and the pgwire layer. As part of any such
1676    /// protocol, we must ensure that the `ExecuteContextGuard`
1677    /// (possibly wrapped in a new `ExecuteContext`) is passed back to the coordinator for
1678    /// eventual retirement. The returned response barriers must stay attached
1679    /// to the user-visible response path.
1680    pub fn into_parts(
1681        self,
1682    ) -> (
1683        ClientTransmitter<ExecuteResponse>,
1684        mpsc::UnboundedSender<Message>,
1685        Session,
1686        ExecuteContextGuard,
1687        Vec<BuiltinTableAppendNotify>,
1688    ) {
1689        let ExecuteContextInner {
1690            tx,
1691            internal_cmd_tx,
1692            session,
1693            extra,
1694            response_barriers,
1695        } = *self.inner;
1696        (tx, internal_cmd_tx, session, extra, response_barriers)
1697    }
1698
1699    /// Retire the execution, by sending a message to the coordinator.
1700    #[instrument(level = "debug")]
1701    pub fn retire(self, result: Result<ExecuteResponse, AdapterError>) {
1702        let ExecuteContextInner {
1703            tx,
1704            internal_cmd_tx,
1705            session,
1706            extra,
1707            response_barriers,
1708        } = *self.inner;
1709        if response_barriers.is_empty() {
1710            retire_execution_context(tx, internal_cmd_tx, session, extra, result);
1711        } else {
1712            spawn(
1713                || "execute_context::retire_after_response_barriers",
1714                async move {
1715                    for barrier in response_barriers {
1716                        barrier.await;
1717                    }
1718                    retire_execution_context(tx, internal_cmd_tx, session, extra, result);
1719                },
1720            );
1721        }
1722    }
1723
1724    /// Delays sending this statement's response until `barrier` resolves.
1725    pub(crate) fn delay_response_until(&mut self, barrier: BuiltinTableAppendCompletion) {
1726        self.response_barriers.push(barrier.into_notify());
1727    }
1728
1729    pub fn extra(&self) -> &ExecuteContextGuard {
1730        &self.extra
1731    }
1732
1733    pub fn extra_mut(&mut self) -> &mut ExecuteContextGuard {
1734        &mut self.extra
1735    }
1736}
1737
1738fn retire_execution_context(
1739    tx: ClientTransmitter<ExecuteResponse>,
1740    internal_cmd_tx: mpsc::UnboundedSender<Message>,
1741    session: Session,
1742    extra: ExecuteContextGuard,
1743    result: Result<ExecuteResponse, AdapterError>,
1744) {
1745    let reason = if extra.is_trivial() {
1746        None
1747    } else {
1748        Some((&result).into())
1749    };
1750    tx.send(result, session);
1751    if let Some(reason) = reason {
1752        let extra = extra.defuse();
1753        if let Err(e) = internal_cmd_tx.send(Message::RetireExecute {
1754            otel_ctx: OpenTelemetryContext::obtain(),
1755            data: extra,
1756            reason,
1757        }) {
1758            warn!("internal_cmd_rx dropped before we could send: {:?}", e);
1759        }
1760    }
1761}
1762
1763#[derive(Debug)]
1764struct ClusterReplicaStatuses(
1765    BTreeMap<ClusterId, BTreeMap<ReplicaId, BTreeMap<ProcessId, ClusterReplicaProcessStatus>>>,
1766);
1767
1768impl ClusterReplicaStatuses {
1769    pub(crate) fn new() -> ClusterReplicaStatuses {
1770        ClusterReplicaStatuses(BTreeMap::new())
1771    }
1772
1773    /// Initializes the statuses of the specified cluster.
1774    ///
1775    /// Panics if the cluster statuses are already initialized.
1776    pub(crate) fn initialize_cluster_statuses(&mut self, cluster_id: ClusterId) {
1777        let prev = self.0.insert(cluster_id, BTreeMap::new());
1778        assert_eq!(
1779            prev, None,
1780            "cluster {cluster_id} statuses already initialized"
1781        );
1782    }
1783
1784    /// Initializes the statuses of the specified cluster replica.
1785    ///
1786    /// Panics if the cluster replica statuses are already initialized.
1787    pub(crate) fn initialize_cluster_replica_statuses(
1788        &mut self,
1789        cluster_id: ClusterId,
1790        replica_id: ReplicaId,
1791        num_processes: usize,
1792        time: DateTime<Utc>,
1793    ) {
1794        tracing::info!(
1795            ?cluster_id,
1796            ?replica_id,
1797            ?time,
1798            "initializing cluster replica status"
1799        );
1800        let replica_statuses = self.0.entry(cluster_id).or_default();
1801        let process_statuses = (0..num_processes)
1802            .map(|process_id| {
1803                let status = ClusterReplicaProcessStatus {
1804                    status: ClusterStatus::Offline(Some(OfflineReason::Initializing)),
1805                    restart_count: 0,
1806                    time: time.clone(),
1807                };
1808                (u64::cast_from(process_id), status)
1809            })
1810            .collect();
1811        let prev = replica_statuses.insert(replica_id, process_statuses);
1812        assert_none!(
1813            prev,
1814            "cluster replica {cluster_id}.{replica_id} statuses already initialized"
1815        );
1816    }
1817
1818    /// Removes the statuses of the specified cluster.
1819    ///
1820    /// Panics if the cluster does not exist.
1821    pub(crate) fn remove_cluster_statuses(
1822        &mut self,
1823        cluster_id: &ClusterId,
1824    ) -> BTreeMap<ReplicaId, BTreeMap<ProcessId, ClusterReplicaProcessStatus>> {
1825        let prev = self.0.remove(cluster_id);
1826        prev.unwrap_or_else(|| panic!("unknown cluster: {cluster_id}"))
1827    }
1828
1829    /// Removes the statuses of the specified cluster replica.
1830    ///
1831    /// Panics if the cluster or replica does not exist.
1832    pub(crate) fn remove_cluster_replica_statuses(
1833        &mut self,
1834        cluster_id: &ClusterId,
1835        replica_id: &ReplicaId,
1836    ) -> BTreeMap<ProcessId, ClusterReplicaProcessStatus> {
1837        let replica_statuses = self
1838            .0
1839            .get_mut(cluster_id)
1840            .unwrap_or_else(|| panic!("unknown cluster: {cluster_id}"));
1841        let prev = replica_statuses.remove(replica_id);
1842        prev.unwrap_or_else(|| panic!("unknown cluster replica: {cluster_id}.{replica_id}"))
1843    }
1844
1845    /// Inserts or updates the status of the specified cluster replica process.
1846    ///
1847    /// Panics if the cluster or replica does not exist.
1848    pub(crate) fn ensure_cluster_status(
1849        &mut self,
1850        cluster_id: ClusterId,
1851        replica_id: ReplicaId,
1852        process_id: ProcessId,
1853        status: ClusterReplicaProcessStatus,
1854    ) {
1855        let replica_statuses = self
1856            .0
1857            .get_mut(&cluster_id)
1858            .unwrap_or_else(|| panic!("unknown cluster: {cluster_id}"))
1859            .get_mut(&replica_id)
1860            .unwrap_or_else(|| panic!("unknown cluster replica: {cluster_id}.{replica_id}"));
1861        replica_statuses.insert(process_id, status);
1862    }
1863
1864    /// Computes the status of the cluster replica as a whole.
1865    ///
1866    /// Panics if `cluster_id` or `replica_id` don't exist.
1867    pub fn get_cluster_replica_status(
1868        &self,
1869        cluster_id: ClusterId,
1870        replica_id: ReplicaId,
1871    ) -> ClusterStatus {
1872        let process_status = self.get_cluster_replica_statuses(cluster_id, replica_id);
1873        Self::cluster_replica_status(process_status)
1874    }
1875
1876    /// Computes the status of the cluster replica as a whole.
1877    pub fn cluster_replica_status(
1878        process_status: &BTreeMap<ProcessId, ClusterReplicaProcessStatus>,
1879    ) -> ClusterStatus {
1880        process_status
1881            .values()
1882            .fold(ClusterStatus::Online, |s, p| match (s, p.status) {
1883                (ClusterStatus::Online, ClusterStatus::Online) => ClusterStatus::Online,
1884                (x, y) => {
1885                    let reason_x = match x {
1886                        ClusterStatus::Offline(reason) => reason,
1887                        ClusterStatus::Online => None,
1888                    };
1889                    let reason_y = match y {
1890                        ClusterStatus::Offline(reason) => reason,
1891                        ClusterStatus::Online => None,
1892                    };
1893                    // Arbitrarily pick the first known not-ready reason.
1894                    ClusterStatus::Offline(reason_x.or(reason_y))
1895                }
1896            })
1897    }
1898
1899    /// Gets the statuses of the given cluster replica.
1900    ///
1901    /// Panics if the cluster or replica does not exist
1902    pub(crate) fn get_cluster_replica_statuses(
1903        &self,
1904        cluster_id: ClusterId,
1905        replica_id: ReplicaId,
1906    ) -> &BTreeMap<ProcessId, ClusterReplicaProcessStatus> {
1907        self.try_get_cluster_replica_statuses(cluster_id, replica_id)
1908            .unwrap_or_else(|| panic!("unknown cluster replica: {cluster_id}.{replica_id}"))
1909    }
1910
1911    /// Gets the statuses of the given cluster replica.
1912    pub(crate) fn try_get_cluster_replica_statuses(
1913        &self,
1914        cluster_id: ClusterId,
1915        replica_id: ReplicaId,
1916    ) -> Option<&BTreeMap<ProcessId, ClusterReplicaProcessStatus>> {
1917        self.try_get_cluster_statuses(cluster_id)
1918            .and_then(|statuses| statuses.get(&replica_id))
1919    }
1920
1921    /// Gets the statuses of the given cluster.
1922    pub(crate) fn try_get_cluster_statuses(
1923        &self,
1924        cluster_id: ClusterId,
1925    ) -> Option<&BTreeMap<ReplicaId, BTreeMap<ProcessId, ClusterReplicaProcessStatus>>> {
1926        self.0.get(&cluster_id)
1927    }
1928}
1929
1930/// Glues the external world to the Timely workers.
1931#[derive(Derivative)]
1932#[derivative(Debug)]
1933pub struct Coordinator {
1934    /// The controller for the storage and compute layers.
1935    #[derivative(Debug = "ignore")]
1936    controller: mz_controller::Controller,
1937    /// The catalog in an Arc suitable for readonly references. The Arc allows
1938    /// us to hand out cheap copies of the catalog to functions that can use it
1939    /// off of the main coordinator thread. If the coordinator needs to mutate
1940    /// the catalog, call [`Self::catalog_mut`], which will clone this struct member,
1941    /// allowing it to be mutated here while the other off-thread references can
1942    /// read their catalog as long as needed. In the future we would like this
1943    /// to be a pTVC, but for now this is sufficient.
1944    catalog: Arc<Catalog>,
1945
1946    /// A client for persist. Initially, this is only used for reading stashed
1947    /// peek responses out of batches.
1948    persist_client: PersistClient,
1949
1950    /// Channel to manage internal commands from the coordinator to itself.
1951    internal_cmd_tx: mpsc::UnboundedSender<Message>,
1952    /// Notification that triggers a group commit.
1953    group_commit_tx: appends::GroupCommitNotifier,
1954    /// Wakes the cluster controller task to reconcile immediately instead of
1955    /// waiting out its tick interval. Notified after catalog transactions that
1956    /// change durable cluster state.
1957    reconcile_now: Arc<Notify>,
1958
1959    /// Channel for strict serializable reads ready to commit.
1960    strict_serializable_reads_tx: mpsc::UnboundedSender<(ConnectionId, PendingReadTxn)>,
1961
1962    /// Signals that pending strict serializable reads should be re-checked
1963    /// because the timestamp oracle may have advanced. Awaited below group commit
1964    /// in [`Coordinator::serve`]; see that branch for the ordering rationale.
1965    linearize_reads_notify: Arc<Notify>,
1966
1967    /// Mechanism for totally ordering write and read timestamps, so that all reads
1968    /// reflect exactly the set of writes that precede them, and no writes that follow.
1969    global_timelines: BTreeMap<Timeline, TimelineState>,
1970
1971    /// A generator for transient [`GlobalId`]s, shareable with other threads.
1972    transient_id_gen: Arc<TransientIdGen>,
1973    /// A map from connection ID to metadata about that connection for all
1974    /// active connections.
1975    active_conns: BTreeMap<ConnectionId, ConnMeta>,
1976
1977    /// For each transaction, the read holds taken to support any performed reads.
1978    ///
1979    /// Upon completing a transaction, these read holds should be dropped.
1980    txn_read_holds: BTreeMap<ConnectionId, read_policy::ReadHolds>,
1981
1982    /// Access to the peek fields should be restricted to methods in the [`peek`] API.
1983    /// A map from pending peek ids to the queue into which responses are sent, and
1984    /// the connection id of the client that initiated the peek.
1985    pending_peeks: BTreeMap<Uuid, PendingPeek>,
1986    /// A map from client connection ids to a set of all pending peeks for that client.
1987    client_pending_peeks: BTreeMap<ConnectionId, BTreeMap<Uuid, ClusterId>>,
1988
1989    /// A map from client connection ids to pending linearize read transaction.
1990    pending_linearize_read_txns: BTreeMap<ConnectionId, PendingReadTxn>,
1991
1992    /// A map from the compute sink ID to it's state description.
1993    active_compute_sinks: BTreeMap<GlobalId, ActiveComputeSink>,
1994    /// A map from active webhooks to their invalidation handle.
1995    active_webhooks: BTreeMap<CatalogItemId, WebhookAppenderInvalidator>,
1996    /// A map of active `COPY FROM` statements. The Coordinator waits for `clusterd`
1997    /// to stage Batches in Persist that we will then link into the shard.
1998    active_copies: BTreeMap<ConnectionId, ActiveCopyFrom>,
1999
2000    /// Connection-scoped cancellation watches.
2001    ///
2002    /// Each entry is a watch channel whose value is `false` until cancellation
2003    /// is requested for that connection, at which point it is set to `true`.
2004    ///
2005    /// Consumers install/remove these watches while they have cancellable work
2006    /// in flight.
2007    connection_cancel_watches: BTreeMap<ConnectionId, (watch::Sender<bool>, watch::Receiver<bool>)>,
2008    /// Active introspection subscribes.
2009    introspection_subscribes: BTreeMap<GlobalId, IntrospectionSubscribe>,
2010
2011    /// Locks that grant access to a specific object, populated lazily as objects are written to.
2012    write_locks: BTreeMap<CatalogItemId, Arc<tokio::sync::Mutex<()>>>,
2013    /// Plans that are currently deferred and waiting on a write lock.
2014    deferred_write_ops: BTreeMap<ConnectionId, DeferredOp>,
2015
2016    /// Pending writes waiting for a group commit.
2017    pending_writes: Vec<PendingWriteTxn>,
2018
2019    /// For the realtime timeline, an explicit SELECT or INSERT on a table will bump the
2020    /// table's timestamps, but there are cases where timestamps are not bumped but
2021    /// we expect the closed timestamps to advance (`AS OF X`, SUBSCRIBing views over
2022    /// RT sources and tables). To address these, spawn a task that forces table
2023    /// timestamps to close on a regular interval. This roughly tracks the behavior
2024    /// of realtime sources that close off timestamps on an interval.
2025    ///
2026    /// For non-realtime timelines, nothing pushes the timestamps forward, so we must do
2027    /// it manually.
2028    advance_timelines_interval: Interval,
2029
2030    /// Serialized DDL. DDL must be serialized because:
2031    /// - Many of them do off-thread work and need to verify the catalog is in a valid state, but
2032    ///   [`PlanValidity`] does not currently support tracking all changes. Doing that correctly
2033    ///   seems to be more difficult than it's worth, so we would instead re-plan and re-sequence
2034    ///   the statements.
2035    /// - Re-planning a statement is hard because Coordinator and Session state is mutated at
2036    ///   various points, and we would need to correctly reset those changes before re-planning and
2037    ///   re-sequencing.
2038    serialized_ddl: LockedVecDeque<DeferredPlanStatement>,
2039
2040    /// Handle to secret manager that can create and delete secrets from
2041    /// an arbitrary secret storage engine.
2042    secrets_controller: Arc<dyn SecretsController>,
2043    /// A secrets reader than maintains an in-memory cache, where values have a set TTL.
2044    caching_secrets_reader: CachingSecretsReader,
2045
2046    /// Handle to a manager that can create and delete kubernetes resources
2047    /// (ie: VpcEndpoint objects)
2048    cloud_resource_controller: Option<Arc<dyn CloudResourceController>>,
2049
2050    /// Persist client for fetching storage metadata such as size metrics.
2051    storage_usage_client: StorageUsageClient,
2052    /// The interval at which to collect storage usage information.
2053    storage_usage_collection_interval: Duration,
2054
2055    /// Segment analytics client.
2056    #[derivative(Debug = "ignore")]
2057    segment_client: Option<mz_segment::Client>,
2058
2059    /// Coordinator metrics.
2060    metrics: Metrics,
2061    /// Optimizer metrics.
2062    optimizer_metrics: OptimizerMetrics,
2063
2064    /// Tracing handle.
2065    tracing_handle: TracingHandle,
2066
2067    /// Data used by the statement logging feature.
2068    statement_logging: StatementLogging,
2069
2070    /// Limit for how many concurrent webhook requests we allow.
2071    webhook_concurrency_limit: WebhookConcurrencyLimiter,
2072
2073    /// Optional config for the timestamp oracle. This is _required_ when
2074    /// a timestamp oracle backend is configured.
2075    timestamp_oracle_config: Option<TimestampOracleConfig>,
2076
2077    /// Periodically asks cluster scheduling policies to make their decisions.
2078    check_cluster_scheduling_policies_interval: Interval,
2079
2080    /// This keeps the last On/Off decision for each cluster and each scheduling policy.
2081    /// (Clusters that have been dropped or are otherwise out of scope for automatic scheduling are
2082    /// periodically cleaned up from this Map.)
2083    cluster_scheduling_decisions: BTreeMap<ClusterId, BTreeMap<&'static str, SchedulingDecision>>,
2084
2085    /// When doing 0dt upgrades/in read-only mode, periodically ask all known
2086    /// clusters/collections whether they are caught up.
2087    caught_up_check_interval: Interval,
2088
2089    /// Context needed to check whether all clusters/collections have caught up.
2090    /// Only used during 0dt deployment, while in read-only mode.
2091    caught_up_check: Option<CaughtUpCheckContext>,
2092
2093    /// The metrics registry, handed to the catalog info-metrics background task
2094    /// so it can register and own its `*_info` series.
2095    catalog_info_metrics_registry: MetricsRegistry,
2096
2097    /// The shared system-parameter frontend, installed by the sync loop once it
2098    /// initializes (and re-installed on reconnect). `None` until then, for
2099    /// example before LaunchDarkly connects, where a newly-created object
2100    /// resolves to the environment-wide value (the cold-cache fallback). Used to
2101    /// resolve a new cluster's or replica's scoped overrides synchronously at
2102    /// create time, so its first plan or first controller configuration is
2103    /// correct rather than waiting for the next sync tick. See the scoped
2104    /// feature flags design.
2105    scoped_frontend: Option<Arc<SystemParameterFrontend>>,
2106
2107    /// Tracks the state associated with the currently installed watchsets.
2108    installed_watch_sets: BTreeMap<WatchSetId, (ConnectionId, WatchSetResponse)>,
2109
2110    /// Tracks the currently installed watchsets for each connection.
2111    connection_watch_sets: BTreeMap<ConnectionId, BTreeSet<WatchSetId>>,
2112
2113    /// Tracks the statuses of all cluster replicas.
2114    cluster_replica_statuses: ClusterReplicaStatuses,
2115
2116    /// Whether or not to start controllers in read-only mode. This is only
2117    /// meant for use during development of read-only clusters and 0dt upgrades
2118    /// and should go away once we have proper orchestration during upgrades.
2119    read_only_controllers: bool,
2120
2121    /// Updates to builtin tables that are being buffered while we are in
2122    /// read-only mode. We apply these all at once when coming out of read-only
2123    /// mode.
2124    ///
2125    /// This is a `Some` while in read-only mode and will be replaced by a
2126    /// `None` when we transition out of read-only mode and write out any
2127    /// buffered updates.
2128    buffered_builtin_table_updates: Option<Vec<BuiltinTableUpdate>>,
2129
2130    license_key: ValidatedLicenseKey,
2131
2132    /// Pre-allocated pool of user IDs to amortize persist writes across DDL operations.
2133    user_id_pool: IdPool,
2134}
2135
2136impl Coordinator {
2137    /// Persists the scoped system-parameter working copy and reconciles it into
2138    /// the per-scope resolution boundaries.
2139    ///
2140    /// The system-parameter sync loop and the create-time fold
2141    /// (`scoped_overrides_create_op`, folded into the create transaction) are the
2142    /// only writers, both serialized on the coordinator loop. The diff is
2143    /// persisted to the
2144    /// durable cache (so values survive an `environmentd` restart and an LD
2145    /// outage) via `Op::UpdateScopedSystemParameters`, which also updates the
2146    /// in-memory working copy in [`CatalogState`] and the
2147    /// `mz_cluster_system_parameters` / `mz_replica_system_parameters`
2148    /// introspection relations. The `replica`-scoped overrides reach the compute
2149    /// controller's per-replica dyncfg layer through the catalog implication for
2150    /// the persisted change. The `cluster`-scoped layer is resolved at plan time
2151    /// via [`CatalogState::cluster_scoped_optimizer_overrides`].
2152    ///
2153    /// [`CatalogState`]: crate::catalog::CatalogState
2154    /// [`CatalogState::cluster_scoped_optimizer_overrides`]: crate::catalog::CatalogState::cluster_scoped_optimizer_overrides
2155    pub(crate) async fn reconcile_scoped_system_parameters(
2156        &mut self,
2157        scoped: ScopedParameters,
2158        prune_scope: Option<ScopedParametersScope>,
2159    ) {
2160        // Nothing changed: skip the durable write. This is the common case on
2161        // most sync ticks.
2162        if self.catalog().state().scoped_system_parameters() == &scoped {
2163            return;
2164        }
2165
2166        // Persist the diff and update the in-memory working copy + introspection
2167        // through the catalog transaction, serialized on the coordinator loop
2168        // with the create-time fold. The replica-scoped
2169        // controller push is derived from this transaction's diff by the catalog
2170        // implication. `prune_scope` bounds removals to the evaluated objects, so
2171        // a concurrently-created object's override is not wiped. Best-effort: a
2172        // failure here is logged and retried on the next sync tick.
2173        if let Err(e) = self
2174            .catalog_transact(
2175                None,
2176                vec![crate::catalog::Op::UpdateScopedSystemParameters {
2177                    scoped,
2178                    prune_scope,
2179                }],
2180            )
2181            .await
2182        {
2183            tracing::warn!("failed to persist scoped system parameters: {e}");
2184        }
2185    }
2186
2187    /// Evaluates the scoped overrides for freshly-created objects from explicit
2188    /// eval contexts and returns an [`Op::UpdateScopedSystemParameters`] to fold
2189    /// into the same transaction that creates them.
2190    ///
2191    /// The objects are not yet in the catalog, so the contexts are built from
2192    /// plan data and pre-allocated ids. Folding the op into the create
2193    /// transaction makes its committed diff drive the replica-scoped controller
2194    /// push, as a catalog implication, before `create_replica`. A new replica's
2195    /// first configuration then carries its overrides rather than the env-wide
2196    /// values. Render-frozen flags (e.g. the column-paged batcher, chosen at
2197    /// arrangement-build time) make a later push too late, which is why this
2198    /// happens in the create transaction rather than the next sync tick.
2199    ///
2200    /// Returns `None` when the feature is gated off, the shared frontend is not
2201    /// yet installed (e.g. before LaunchDarkly connects), or no override
2202    /// applies. The new objects then resolve to the environment-wide value, and
2203    /// the periodic sync loop remains the authoritative full-state reconciler.
2204    ///
2205    /// [`Op::UpdateScopedSystemParameters`]: crate::catalog::Op::UpdateScopedSystemParameters
2206    fn scoped_overrides_create_op(
2207        &self,
2208        clusters: &[ClusterEvalContext],
2209        replicas: &[ReplicaEvalContext],
2210    ) -> Option<crate::catalog::Op> {
2211        let frontend = self.scoped_frontend.clone()?;
2212        let catalog = self.catalog();
2213        let system_config = catalog.system_config();
2214        if !ENABLE_SCOPED_SYSTEM_PARAMETERS.get(system_config.dyncfgs()) {
2215            return None;
2216        }
2217
2218        // Partition the synced parameters by scope class, as the sync loop does,
2219        // so we evaluate exactly the flags in use at each scope.
2220        let replica_param_names: Vec<&'static str> = system_config
2221            .iter_synced()
2222            .filter(|var| var.scope() == ParameterScope::Replica)
2223            .map(|var| var.name())
2224            .collect();
2225        let cluster_param_names: Vec<&'static str> = system_config
2226            .iter_synced()
2227            .filter(|var| var.scope() == ParameterScope::Cluster)
2228            .map(|var| var.name())
2229            .collect();
2230
2231        let params = SynchronizedParameters::new(system_config.clone());
2232        let mut evaluated = ScopedParameters::default();
2233        if !cluster_param_names.is_empty() && !clusters.is_empty() {
2234            evaluated.cluster =
2235                frontend.pull_cluster_overrides(&params, &cluster_param_names, clusters);
2236        }
2237        if !replica_param_names.is_empty() && !replicas.is_empty() {
2238            evaluated.replica =
2239                frontend.pull_replica_overrides(&params, &replica_param_names, replicas);
2240        }
2241        if evaluated.is_empty() {
2242            return None;
2243        }
2244
2245        // Prune only within the objects being created. They have no prior rows,
2246        // so nothing is removed, and this op never touches another object whose
2247        // override a concurrent reconcile may be writing.
2248        let prune_scope = ScopedParametersScope {
2249            clusters: clusters.iter().map(|cluster| cluster.cluster_id).collect(),
2250            replicas: replicas.iter().map(|replica| replica.replica_id).collect(),
2251        };
2252        Some(crate::catalog::Op::UpdateScopedSystemParameters {
2253            scoped: evaluated,
2254            prune_scope: Some(prune_scope),
2255        })
2256    }
2257
2258    /// Resolves the replica-local scoped overrides from the catalog working copy
2259    /// into the compute controller's per-replica dyncfg layer, then re-pushes
2260    /// the environment-wide compute configuration so replicas observe the new
2261    /// values. Driven by the catalog implication for replica-scoped
2262    /// configuration changes, and called once on bootstrap.
2263    pub(crate) fn push_replica_dyncfg_overrides(&mut self) {
2264        // Clone the (sparse) replica overrides so we don't hold a catalog borrow
2265        // across the mutable controller calls below.
2266        let replica_overrides = self
2267            .catalog()
2268            .state()
2269            .scoped_system_parameters()
2270            .replica
2271            .clone();
2272
2273        let dyncfgs = self.catalog().system_config().dyncfgs();
2274        let mut instance_overrides: BTreeMap<
2275            ComputeInstanceId,
2276            BTreeMap<ReplicaId, ConfigUpdates>,
2277        > = BTreeMap::new();
2278        for cluster in self.catalog().clusters() {
2279            for replica in cluster.replicas() {
2280                let Some(values) = replica_overrides.get(&replica.replica_id) else {
2281                    continue;
2282                };
2283                let mut updates = ConfigUpdates::default();
2284                for (name, value) in values {
2285                    let Some(entry) = dyncfgs.entry(name) else {
2286                        // A replica-local parameter that is not a dyncfg has no
2287                        // per-replica realization, so skip it.
2288                        continue;
2289                    };
2290                    match entry.parse_val(value) {
2291                        Ok(val) => updates.add_dynamic(name, val),
2292                        Err(e) => {
2293                            tracing::warn!(%name, %value, "cannot parse scoped override: {e}")
2294                        }
2295                    }
2296                }
2297                if !updates.updates.is_empty() {
2298                    instance_overrides
2299                        .entry(cluster.id)
2300                        .or_default()
2301                        .insert(replica.replica_id, updates);
2302                }
2303            }
2304        }
2305
2306        // Only the compute controller's per-replica dyncfg layer is pushed, but on
2307        // `clusterd` that also reaches storage. Compute and storage share one
2308        // process, and the compute worker's `handle_update_configuration` applies
2309        // the pushed dyncfg updates both to compute's own worker `ConfigSet` and to
2310        // the shared persist client `ConfigSet` (`persist_clients.cfg()`) that the
2311        // co-located storage server reads from the same `Arc`. So persist-backed and
2312        // process-global replica-local configs such as the persist pager, LZ4,
2313        // persist client tuning, and `lgalloc` take effect on storage too. The only
2314        // gap would be a future `Replica`-scoped config realized solely in the
2315        // storage worker's own `ConfigSet`, of which none exists today.
2316        self.controller
2317            .compute
2318            .update_replica_dyncfg_overrides(instance_overrides);
2319        // Re-push the env-wide compute config so existing replicas pick up their
2320        // (possibly changed) overrides. This also reverts a removed override: the
2321        // per-replica layer no longer carries the key, so the replica falls back
2322        // to the env-wide value, which `compute_config` always includes because it
2323        // renders the full dyncfg set.
2324        let compute_config = crate::flags::compute_config(self.catalog().system_config());
2325        self.controller.compute.update_configuration(compute_config);
2326    }
2327
2328    /// Returns the cluster-coherent scoped optimizer-feature overrides for
2329    /// `cluster_id`. See
2330    /// [`CatalogState::cluster_scoped_optimizer_overrides`](crate::catalog::CatalogState::cluster_scoped_optimizer_overrides).
2331    pub(crate) fn cluster_scoped_optimizer_overrides(
2332        &self,
2333        cluster_id: ClusterId,
2334    ) -> OptimizerFeatureOverrides {
2335        self.catalog()
2336            .state()
2337            .cluster_scoped_optimizer_overrides(cluster_id)
2338    }
2339
2340    /// Initializes coordinator state based on the contained catalog. Must be
2341    /// called after creating the coordinator and before calling the
2342    /// `Coordinator::serve` method.
2343    #[instrument(name = "coord::bootstrap")]
2344    pub(crate) async fn bootstrap(
2345        &mut self,
2346        boot_ts: Timestamp,
2347        migrated_storage_collections_0dt: BTreeSet<CatalogItemId>,
2348        mut builtin_table_updates: Vec<BuiltinTableUpdate>,
2349        cached_global_exprs: BTreeMap<GlobalId, GlobalExpressions>,
2350        uncached_local_exprs: BTreeMap<GlobalId, LocalExpressions>,
2351    ) -> Result<(), AdapterError> {
2352        let bootstrap_start = Instant::now();
2353        info!("startup: coordinator init: bootstrap beginning");
2354        info!("startup: coordinator init: bootstrap: preamble beginning");
2355
2356        // Initialize cluster replica statuses.
2357        // Gross iterator is to avoid partial borrow issues.
2358        let cluster_statuses: Vec<(_, Vec<_>)> = self
2359            .catalog()
2360            .clusters()
2361            .map(|cluster| {
2362                (
2363                    cluster.id(),
2364                    cluster
2365                        .replicas()
2366                        .map(|replica| {
2367                            (replica.replica_id, replica.config.location.num_processes())
2368                        })
2369                        .collect(),
2370                )
2371            })
2372            .collect();
2373        let now = self.now_datetime();
2374        for (cluster_id, replica_statuses) in cluster_statuses {
2375            self.cluster_replica_statuses
2376                .initialize_cluster_statuses(cluster_id);
2377            for (replica_id, num_processes) in replica_statuses {
2378                self.cluster_replica_statuses
2379                    .initialize_cluster_replica_statuses(
2380                        cluster_id,
2381                        replica_id,
2382                        num_processes,
2383                        now,
2384                    );
2385            }
2386        }
2387
2388        let system_config = self.catalog().system_config();
2389
2390        // Inform metrics about the initial system configuration.
2391        mz_metrics::update_dyncfg(&system_config.dyncfg_updates());
2392
2393        // Inform the controllers about their initial configuration.
2394        let compute_config = flags::compute_config(system_config);
2395        let storage_config = flags::storage_config(system_config);
2396        let scheduling_config = flags::orchestrator_scheduling_config(system_config);
2397        let dyncfg_updates = system_config.dyncfg_updates();
2398        self.controller.compute.update_configuration(compute_config);
2399        self.controller.storage.update_parameters(storage_config);
2400        self.controller
2401            .update_orchestrator_scheduling_config(scheduling_config);
2402        self.controller.update_configuration(dyncfg_updates);
2403
2404        // Skip the credit consumption check at bootstrap under DisableClusterCreation behavior:
2405        // this codepath validates existing replicas at startup, not cluster creation, so it
2406        // must not block startup. New cluster creation is still gated by the DDL-time check.
2407        // The Disable case is already handled by a bail! in main.rs before we reach here.
2408        let enforce_credit_limit_at_bootstrap = !matches!(
2409            self.license_key.expiration_behavior,
2410            ExpirationBehavior::DisableClusterCreation,
2411        );
2412        if enforce_credit_limit_at_bootstrap {
2413            self.validate_resource_limit_numeric(
2414                Numeric::zero(),
2415                self.current_credit_consumption_rate(),
2416                |system_vars| {
2417                    self.license_key
2418                        .max_credit_consumption_rate()
2419                        .map_or_else(|| system_vars.max_credit_consumption_rate(), Numeric::from)
2420                },
2421                "cluster replica",
2422                MAX_CREDIT_CONSUMPTION_RATE.name(),
2423            )?;
2424        }
2425
2426        let mut policies_to_set: BTreeMap<CompactionWindow, CollectionIdBundle> =
2427            Default::default();
2428
2429        let enable_worker_core_affinity =
2430            self.catalog().system_config().enable_worker_core_affinity();
2431        let enable_storage_introspection_logs = self
2432            .catalog()
2433            .system_config()
2434            .enable_storage_introspection_logs();
2435        for instance in self.catalog.clusters() {
2436            self.controller.create_cluster(
2437                instance.id,
2438                ClusterConfig {
2439                    arranged_logs: instance.log_indexes.clone(),
2440                    workload_class: instance.config.workload_class.clone(),
2441                },
2442            )?;
2443            for replica in instance.replicas() {
2444                let role = instance.role();
2445                self.controller.create_replica(
2446                    instance.id,
2447                    replica.replica_id,
2448                    instance.name.clone(),
2449                    replica.name.clone(),
2450                    role,
2451                    replica.config.clone(),
2452                    enable_worker_core_affinity,
2453                    enable_storage_introspection_logs,
2454                )?;
2455            }
2456        }
2457
2458        // Now that the compute instances and their replicas exist, push the
2459        // replica-local scoped overrides into the compute controller so existing
2460        // replicas observe them at startup. The scoped (per-cluster and
2461        // per-replica) working copy was restored from the durable cache into
2462        // `CatalogState` while opening the catalog, so the last-known values are
2463        // in effect before the first parameter sync and through a sync outage.
2464        // This must run after the creation loop above: the push iterates the
2465        // controller's instances, so before they exist it is a no-op. It also
2466        // runs before dataflows are rendered later in bootstrap, so render-frozen
2467        // replica flags take effect. The cluster-coherent layer is read at plan
2468        // time.
2469        self.push_replica_dyncfg_overrides();
2470
2471        info!(
2472            "startup: coordinator init: bootstrap: preamble complete in {:?}",
2473            bootstrap_start.elapsed()
2474        );
2475
2476        let init_storage_collections_start = Instant::now();
2477        info!("startup: coordinator init: bootstrap: storage collections init beginning");
2478        self.bootstrap_storage_collections(&migrated_storage_collections_0dt)
2479            .await;
2480        info!(
2481            "startup: coordinator init: bootstrap: storage collections init complete in {:?}",
2482            init_storage_collections_start.elapsed()
2483        );
2484
2485        // The storage controller knows about the introspection collections now, so we can start
2486        // sinking introspection updates in the compute controller. It makes sense to do that as
2487        // soon as possible, to avoid updates piling up in the compute controller's internal
2488        // buffers.
2489        self.controller.start_compute_introspection_sink();
2490
2491        let sorting_start = Instant::now();
2492        info!("startup: coordinator init: bootstrap: sorting catalog entries");
2493        let entries = self.bootstrap_sort_catalog_entries();
2494        info!(
2495            "startup: coordinator init: bootstrap: sorting catalog entries complete in {:?}",
2496            sorting_start.elapsed()
2497        );
2498
2499        let optimize_dataflows_start = Instant::now();
2500        info!("startup: coordinator init: bootstrap: optimize dataflow plans beginning");
2501        let uncached_global_exps = self.bootstrap_dataflow_plans(&entries, cached_global_exprs)?;
2502        info!(
2503            "startup: coordinator init: bootstrap: optimize dataflow plans complete in {:?}",
2504            optimize_dataflows_start.elapsed()
2505        );
2506
2507        // We don't need to wait for the cache to update.
2508        let _fut = self.catalog().update_expression_cache(
2509            uncached_local_exprs.into_iter().collect(),
2510            uncached_global_exps.into_iter().collect(),
2511            Default::default(),
2512        );
2513
2514        // Select dataflow as-ofs. This step relies on the storage collections created by
2515        // `bootstrap_storage_collections` and the dataflow plans created by
2516        // `bootstrap_dataflow_plans`.
2517        let bootstrap_as_ofs_start = Instant::now();
2518        info!("startup: coordinator init: bootstrap: dataflow as-of bootstrapping beginning");
2519        let dataflow_read_holds = self.bootstrap_dataflow_as_ofs().await;
2520        info!(
2521            "startup: coordinator init: bootstrap: dataflow as-of bootstrapping complete in {:?}",
2522            bootstrap_as_ofs_start.elapsed()
2523        );
2524
2525        let postamble_start = Instant::now();
2526        info!("startup: coordinator init: bootstrap: postamble beginning");
2527
2528        let logs: BTreeSet<_> = BUILTINS::logs()
2529            .map(|log| self.catalog().resolve_builtin_log(log))
2530            .flat_map(|item_id| self.catalog().get_global_ids(&item_id))
2531            .collect();
2532
2533        let mut privatelink_connections = BTreeMap::new();
2534
2535        for entry in &entries {
2536            debug!(
2537                "coordinator init: installing {} {}",
2538                entry.item().typ(),
2539                entry.id()
2540            );
2541            let mut policy = entry.item().initial_logical_compaction_window();
2542            match entry.item() {
2543                // Currently catalog item rebuild assumes that sinks and
2544                // indexes are always built individually and does not store information
2545                // about how it was built. If we start building multiple sinks and/or indexes
2546                // using a single dataflow, we have to make sure the rebuild process re-runs
2547                // the same multiple-build dataflow.
2548                CatalogItem::Source(source) => {
2549                    // Propagate source compaction windows to subsources if needed.
2550                    if source.custom_logical_compaction_window.is_none() {
2551                        if let DataSourceDesc::IngestionExport { ingestion_id, .. } =
2552                            source.data_source
2553                        {
2554                            policy = Some(
2555                                self.catalog()
2556                                    .get_entry(&ingestion_id)
2557                                    .source()
2558                                    .expect("must be source")
2559                                    .custom_logical_compaction_window
2560                                    .unwrap_or_default(),
2561                            );
2562                        }
2563                    }
2564                    policies_to_set
2565                        .entry(policy.expect("sources have a compaction window"))
2566                        .or_insert_with(Default::default)
2567                        .storage_ids
2568                        .insert(source.global_id());
2569                }
2570                CatalogItem::Table(table) => {
2571                    policies_to_set
2572                        .entry(policy.expect("tables have a compaction window"))
2573                        .or_insert_with(Default::default)
2574                        .storage_ids
2575                        .extend(table.global_ids());
2576                }
2577                CatalogItem::Index(idx) => {
2578                    let policy_entry = policies_to_set
2579                        .entry(policy.expect("indexes have a compaction window"))
2580                        .or_insert_with(Default::default);
2581
2582                    if logs.contains(&idx.on) {
2583                        policy_entry
2584                            .compute_ids
2585                            .entry(idx.cluster_id)
2586                            .or_insert_with(BTreeSet::new)
2587                            .insert(idx.global_id());
2588                    } else {
2589                        let df_desc = self
2590                            .catalog()
2591                            .try_get_physical_plan(&idx.global_id())
2592                            .expect("added in `bootstrap_dataflow_plans`")
2593                            .clone();
2594
2595                        let df_meta = self
2596                            .catalog()
2597                            .try_get_dataflow_metainfo(&idx.global_id())
2598                            .expect("added in `bootstrap_dataflow_plans`");
2599
2600                        if self.catalog().state().system_config().enable_mz_notices() {
2601                            // Collect optimization hint updates.
2602                            self.catalog().state().pack_optimizer_notices(
2603                                &mut builtin_table_updates,
2604                                df_meta.optimizer_notices.iter(),
2605                                Diff::ONE,
2606                            );
2607                        }
2608
2609                        // What follows is morally equivalent to `self.ship_dataflow(df, idx.cluster_id)`,
2610                        // but we cannot call that as it will also downgrade the read hold on the index.
2611                        policy_entry
2612                            .compute_ids
2613                            .entry(idx.cluster_id)
2614                            .or_insert_with(Default::default)
2615                            .extend(df_desc.export_ids());
2616
2617                        self.controller
2618                            .compute
2619                            .create_dataflow(idx.cluster_id, df_desc, None)
2620                            .unwrap_or_terminate("cannot fail to create dataflows");
2621                    }
2622                }
2623                CatalogItem::View(_) => (),
2624                CatalogItem::MaterializedView(mview) => {
2625                    policies_to_set
2626                        .entry(policy.expect("materialized views have a compaction window"))
2627                        .or_insert_with(Default::default)
2628                        .storage_ids
2629                        .insert(mview.global_id_writes());
2630
2631                    let mut df_desc = self
2632                        .catalog()
2633                        .try_get_physical_plan(&mview.global_id_writes())
2634                        .expect("added in `bootstrap_dataflow_plans`")
2635                        .clone();
2636
2637                    if let Some(initial_as_of) = mview.initial_as_of.clone() {
2638                        df_desc.set_initial_as_of(initial_as_of);
2639                    }
2640
2641                    // If we have a refresh schedule that has a last refresh, then set the `until` to the last refresh.
2642                    let until = mview
2643                        .refresh_schedule
2644                        .as_ref()
2645                        .and_then(|s| s.last_refresh())
2646                        .and_then(|r| r.try_step_forward());
2647                    if let Some(until) = until {
2648                        df_desc.until.meet_assign(&Antichain::from_elem(until));
2649                    }
2650
2651                    let df_meta = self
2652                        .catalog()
2653                        .try_get_dataflow_metainfo(&mview.global_id_writes())
2654                        .expect("added in `bootstrap_dataflow_plans`");
2655
2656                    if self.catalog().state().system_config().enable_mz_notices() {
2657                        // Collect optimization hint updates.
2658                        self.catalog().state().pack_optimizer_notices(
2659                            &mut builtin_table_updates,
2660                            df_meta.optimizer_notices.iter(),
2661                            Diff::ONE,
2662                        );
2663                    }
2664
2665                    self.ship_dataflow(df_desc, mview.cluster_id, mview.target_replica)
2666                        .await;
2667
2668                    // If this is a replacement MV, it must remain read-only until the replacement
2669                    // gets applied.
2670                    if mview.replacement_target.is_none() {
2671                        self.allow_writes(mview.cluster_id, mview.global_id_writes());
2672                    }
2673                }
2674                CatalogItem::Sink(sink) => {
2675                    policies_to_set
2676                        .entry(CompactionWindow::Default)
2677                        .or_insert_with(Default::default)
2678                        .storage_ids
2679                        .insert(sink.global_id());
2680                }
2681                CatalogItem::Connection(catalog_connection) => {
2682                    if let ConnectionDetails::AwsPrivatelink(conn) = &catalog_connection.details {
2683                        privatelink_connections.insert(
2684                            entry.id(),
2685                            VpcEndpointConfig {
2686                                aws_service_name: conn.service_name.clone(),
2687                                availability_zone_ids: conn.availability_zones.clone(),
2688                            },
2689                        );
2690                    }
2691                }
2692                // Nothing to do for these cases
2693                CatalogItem::Log(_)
2694                | CatalogItem::Type(_)
2695                | CatalogItem::Func(_)
2696                | CatalogItem::Secret(_) => {}
2697            }
2698        }
2699
2700        if let Some(cloud_resource_controller) = &self.cloud_resource_controller {
2701            // Clean up any extraneous VpcEndpoints that shouldn't exist.
2702            let existing_vpc_endpoints = cloud_resource_controller
2703                .list_vpc_endpoints()
2704                .await
2705                .context("list vpc endpoints")?;
2706            let existing_vpc_endpoints = BTreeSet::from_iter(existing_vpc_endpoints.into_keys());
2707            let desired_vpc_endpoints = privatelink_connections.keys().cloned().collect();
2708            let vpc_endpoints_to_remove = existing_vpc_endpoints.difference(&desired_vpc_endpoints);
2709            for id in vpc_endpoints_to_remove {
2710                cloud_resource_controller
2711                    .delete_vpc_endpoint(*id)
2712                    .await
2713                    .context("deleting extraneous vpc endpoint")?;
2714            }
2715
2716            // Ensure desired VpcEndpoints are up to date.
2717            for (id, spec) in privatelink_connections {
2718                cloud_resource_controller
2719                    .ensure_vpc_endpoint(id, spec)
2720                    .await
2721                    .context("ensuring vpc endpoint")?;
2722            }
2723        }
2724
2725        // Having installed all entries, creating all constraints, we can now drop read holds and
2726        // relax read policies.
2727        drop(dataflow_read_holds);
2728        // TODO -- Improve `initialize_read_policies` API so we can avoid calling this in a loop.
2729        for (cw, policies) in policies_to_set {
2730            self.initialize_read_policies(&policies, cw).await;
2731        }
2732
2733        // Expose mapping from T-shirt sizes to actual sizes
2734        builtin_table_updates.extend(
2735            self.catalog().state().resolve_builtin_table_updates(
2736                self.catalog().state().pack_all_replica_size_updates(),
2737            ),
2738        );
2739
2740        debug!("startup: coordinator init: bootstrap: initializing migrated builtin tables");
2741        // When 0dt is enabled, we create new shards for any migrated builtin storage collections.
2742        // In read-only mode, the migrated builtin tables (which are a subset of migrated builtin
2743        // storage collections) need to be back-filled so that any dependent dataflow can be
2744        // hydrated. Additionally, these shards are not registered with the txn-shard, and cannot
2745        // be registered while in read-only, so they are written to directly.
2746        let migrated_updates_fut = if self.controller.read_only() {
2747            let min_timestamp = Timestamp::minimum();
2748            let migrated_builtin_table_updates: Vec<_> = builtin_table_updates
2749                .extract_if(.., |update| {
2750                    let gid = self.catalog().get_entry(&update.id).latest_global_id();
2751                    migrated_storage_collections_0dt.contains(&update.id)
2752                        && self
2753                            .controller
2754                            .storage_collections
2755                            .collection_frontiers(gid)
2756                            .expect("all tables are registered")
2757                            .write_frontier
2758                            .elements()
2759                            == &[min_timestamp]
2760                })
2761                .collect();
2762            if migrated_builtin_table_updates.is_empty() {
2763                futures::future::ready(()).boxed()
2764            } else {
2765                // Group all updates per-table.
2766                let mut grouped_appends: BTreeMap<GlobalId, Vec<TableData>> = BTreeMap::new();
2767                for update in migrated_builtin_table_updates {
2768                    let gid = self.catalog().get_entry(&update.id).latest_global_id();
2769                    grouped_appends.entry(gid).or_default().push(update.data);
2770                }
2771                info!(
2772                    "coordinator init: rehydrating migrated builtin tables in read-only mode: {:?}",
2773                    grouped_appends.keys().collect::<Vec<_>>()
2774                );
2775
2776                // Consolidate Row data, staged batches must already be consolidated.
2777                let mut all_appends = Vec::with_capacity(grouped_appends.len());
2778                for (item_id, table_data) in grouped_appends.into_iter() {
2779                    let mut all_rows = Vec::new();
2780                    let mut all_data = Vec::new();
2781                    for data in table_data {
2782                        match data {
2783                            TableData::Rows(rows) => all_rows.extend(rows),
2784                            TableData::Batches(_) => all_data.push(data),
2785                        }
2786                    }
2787                    differential_dataflow::consolidation::consolidate(&mut all_rows);
2788                    all_data.push(TableData::Rows(all_rows));
2789
2790                    // TODO(parkmycar): Use SmallVec throughout.
2791                    all_appends.push((item_id, all_data));
2792                }
2793
2794                let fut = self
2795                    .controller
2796                    .storage
2797                    .append_table(min_timestamp, boot_ts.step_forward(), all_appends)
2798                    .expect("cannot fail to append");
2799                async {
2800                    fut.await
2801                        .expect("One-shot shouldn't be dropped during bootstrap")
2802                        .unwrap_or_terminate("cannot fail to append")
2803                }
2804                .boxed()
2805            }
2806        } else {
2807            futures::future::ready(()).boxed()
2808        };
2809
2810        info!(
2811            "startup: coordinator init: bootstrap: postamble complete in {:?}",
2812            postamble_start.elapsed()
2813        );
2814
2815        let builtin_update_start = Instant::now();
2816        info!("startup: coordinator init: bootstrap: generate builtin updates beginning");
2817
2818        if self.controller.read_only() {
2819            info!(
2820                "coordinator init: bootstrap: stashing builtin table updates while in read-only mode"
2821            );
2822
2823            self.buffered_builtin_table_updates
2824                .as_mut()
2825                .expect("in read-only mode")
2826                .append(&mut builtin_table_updates);
2827        } else {
2828            self.bootstrap_tables(&entries, builtin_table_updates).await;
2829        };
2830        info!(
2831            "startup: coordinator init: bootstrap: generate builtin updates complete in {:?}",
2832            builtin_update_start.elapsed()
2833        );
2834
2835        let cleanup_secrets_start = Instant::now();
2836        info!("startup: coordinator init: bootstrap: generate secret cleanup beginning");
2837        // Cleanup orphaned secrets. Errors during list() or delete() do not
2838        // need to prevent bootstrap from succeeding; we will retry next
2839        // startup.
2840        {
2841            // Destructure Self so we can selectively move fields into the async
2842            // task.
2843            let Self {
2844                secrets_controller,
2845                catalog,
2846                ..
2847            } = self;
2848
2849            let next_user_item_id = catalog.get_next_user_item_id().await?;
2850            let next_system_item_id = catalog.get_next_system_item_id().await?;
2851            let read_only = self.controller.read_only();
2852            // Fetch all IDs from the catalog to future-proof against other
2853            // things using secrets. Today, SECRET and CONNECTION objects use
2854            // secrets_controller.ensure, but more things could in the future
2855            // that would be easy to miss adding here.
2856            let catalog_ids: BTreeSet<CatalogItemId> =
2857                catalog.entries().map(|entry| entry.id()).collect();
2858            let secrets_controller = Arc::clone(secrets_controller);
2859
2860            spawn(|| "cleanup-orphaned-secrets", async move {
2861                if read_only {
2862                    info!(
2863                        "coordinator init: not cleaning up orphaned secrets while in read-only mode"
2864                    );
2865                    return;
2866                }
2867                info!("coordinator init: cleaning up orphaned secrets");
2868
2869                match secrets_controller.list().await {
2870                    Ok(controller_secrets) => {
2871                        let controller_secrets: BTreeSet<CatalogItemId> =
2872                            controller_secrets.into_iter().collect();
2873                        let orphaned = controller_secrets.difference(&catalog_ids);
2874                        for id in orphaned {
2875                            let id_too_large = match id {
2876                                CatalogItemId::System(id) => *id >= next_system_item_id,
2877                                CatalogItemId::User(id) => *id >= next_user_item_id,
2878                                CatalogItemId::IntrospectionSourceIndex(_)
2879                                | CatalogItemId::Transient(_) => false,
2880                            };
2881                            if id_too_large {
2882                                info!(
2883                                    %next_user_item_id, %next_system_item_id,
2884                                    "coordinator init: not deleting orphaned secret {id} that was likely created by a newer deploy generation"
2885                                );
2886                            } else {
2887                                info!("coordinator init: deleting orphaned secret {id}");
2888                                fail_point!("orphan_secrets");
2889                                if let Err(e) = secrets_controller.delete(*id).await {
2890                                    warn!(
2891                                        "Dropping orphaned secret has encountered an error: {}",
2892                                        e
2893                                    );
2894                                }
2895                            }
2896                        }
2897                    }
2898                    Err(e) => warn!("Failed to list secrets during orphan cleanup: {:?}", e),
2899                }
2900            });
2901        }
2902        info!(
2903            "startup: coordinator init: bootstrap: generate secret cleanup complete in {:?}",
2904            cleanup_secrets_start.elapsed()
2905        );
2906
2907        // Run all of our final steps concurrently.
2908        let final_steps_start = Instant::now();
2909        info!(
2910            "startup: coordinator init: bootstrap: migrate builtin tables in read-only mode beginning"
2911        );
2912        migrated_updates_fut
2913            .instrument(info_span!("coord::bootstrap::final"))
2914            .await;
2915
2916        debug!(
2917            "startup: coordinator init: bootstrap: announcing completion of initialization to controller"
2918        );
2919        // Announce the completion of initialization.
2920        self.controller.initialization_complete();
2921
2922        // Initialize unified introspection.
2923        self.bootstrap_introspection_subscribes().await;
2924
2925        info!(
2926            "startup: coordinator init: bootstrap: migrate builtin tables in read-only mode complete in {:?}",
2927            final_steps_start.elapsed()
2928        );
2929
2930        info!(
2931            "startup: coordinator init: bootstrap complete in {:?}",
2932            bootstrap_start.elapsed()
2933        );
2934        Ok(())
2935    }
2936
2937    /// Prepares tables for writing by resetting them to a known state and
2938    /// appending the given builtin table updates. The timestamp oracle
2939    /// will be advanced to the write timestamp of the append when this
2940    /// method returns.
2941    #[allow(clippy::async_yields_async)]
2942    #[instrument]
2943    async fn bootstrap_tables(
2944        &mut self,
2945        entries: &[CatalogEntry],
2946        mut builtin_table_updates: Vec<BuiltinTableUpdate>,
2947    ) {
2948        /// Smaller helper struct of metadata for bootstrapping tables.
2949        struct TableMetadata<'a> {
2950            id: CatalogItemId,
2951            name: &'a QualifiedItemName,
2952            table: &'a Table,
2953        }
2954
2955        // Filter our entries down to just tables.
2956        let table_metas: Vec<_> = entries
2957            .into_iter()
2958            .filter_map(|entry| {
2959                entry.table().map(|table| TableMetadata {
2960                    id: entry.id(),
2961                    name: entry.name(),
2962                    table,
2963                })
2964            })
2965            .collect();
2966
2967        // Append empty batches to advance the timestamp of all tables.
2968        debug!("coordinator init: advancing all tables to current timestamp");
2969        let WriteTimestamp {
2970            timestamp: write_ts,
2971            advance_to,
2972        } = self.get_local_write_ts().await;
2973        let appends = table_metas
2974            .iter()
2975            .map(|meta| (meta.table.global_id_writes(), Vec::new()))
2976            .collect();
2977        // Append the tables in the background. We apply the write timestamp before getting a read
2978        // timestamp and reading a snapshot of each table, so the snapshots will block on their own
2979        // until the appends are complete.
2980        let table_fence_rx = self
2981            .controller
2982            .storage
2983            .append_table(write_ts.clone(), advance_to, appends)
2984            .expect("invalid updates");
2985
2986        self.apply_local_write(write_ts).await;
2987
2988        // Add builtin table updates the clear the contents of all system tables
2989        debug!("coordinator init: resetting system tables");
2990        let read_ts = self.get_local_read_ts().await;
2991
2992        // Filter out the 'mz_storage_usage_by_shard' table since we need to retain that info for
2993        // billing purposes.
2994        let mz_storage_usage_by_shard_schema: SchemaSpecifier = self
2995            .catalog()
2996            .resolve_system_schema(MZ_STORAGE_USAGE_BY_SHARD.schema)
2997            .into();
2998        let is_storage_usage_by_shard = |meta: &TableMetadata| -> bool {
2999            meta.name.item == MZ_STORAGE_USAGE_BY_SHARD.name
3000                && meta.name.qualifiers.schema_spec == mz_storage_usage_by_shard_schema
3001        };
3002
3003        let mut retraction_tasks = Vec::new();
3004        let system_tables: Vec<_> = table_metas
3005            .iter()
3006            .filter(|meta| meta.id.is_system() && !is_storage_usage_by_shard(meta))
3007            .collect();
3008
3009        for system_table in system_tables {
3010            let table_id = system_table.id;
3011            let full_name = self.catalog().resolve_full_name(system_table.name, None);
3012            debug!("coordinator init: resetting system table {full_name} ({table_id})");
3013
3014            // Fetch the current contents of the table for retraction.
3015            let snapshot_fut = self
3016                .controller
3017                .storage_collections
3018                .snapshot_cursor(system_table.table.global_id_writes(), read_ts);
3019            let batch_fut = self
3020                .controller
3021                .storage_collections
3022                .create_update_builder(system_table.table.global_id_writes());
3023
3024            let task = spawn(|| format!("snapshot-{table_id}"), async move {
3025                // Create a TimestamplessUpdateBuilder.
3026                let mut batch = batch_fut
3027                    .await
3028                    .unwrap_or_terminate("cannot fail to create a batch for a BuiltinTable");
3029                tracing::info!(?table_id, "starting snapshot");
3030                // Get a cursor which will emit a consolidated snapshot.
3031                let mut snapshot_cursor = snapshot_fut
3032                    .await
3033                    .unwrap_or_terminate("cannot fail to snapshot");
3034
3035                // Retract the current contents, spilling into our builder.
3036                while let Some(values) = snapshot_cursor.next().await {
3037                    for (key, _t, d) in values {
3038                        let d_invert = d.neg();
3039                        batch.add(&key, &(), &d_invert).await;
3040                    }
3041                }
3042                tracing::info!(?table_id, "finished snapshot");
3043
3044                let batch = batch.finish().await;
3045                BuiltinTableUpdate::batch(table_id, batch)
3046            });
3047            retraction_tasks.push(task);
3048        }
3049
3050        let retractions_res = futures::future::join_all(retraction_tasks).await;
3051        for retractions in retractions_res {
3052            builtin_table_updates.push(retractions);
3053        }
3054
3055        // Now that the snapshots are complete, the appends must also be complete.
3056        table_fence_rx
3057            .await
3058            .expect("One-shot shouldn't be dropped during bootstrap")
3059            .unwrap_or_terminate("cannot fail to append");
3060
3061        info!("coordinator init: sending builtin table updates");
3062        let (_builtin_updates_fut, write_ts) = self
3063            .builtin_table_update()
3064            .execute(builtin_table_updates)
3065            .await;
3066        info!(?write_ts, "our write ts");
3067        if let Some(write_ts) = write_ts {
3068            self.apply_local_write(write_ts).await;
3069        }
3070    }
3071
3072    /// Initializes all storage collections required by catalog objects in the storage controller.
3073    ///
3074    /// This method takes care of collection creation, as well as migration of existing
3075    /// collections.
3076    ///
3077    /// Creating all storage collections in a single `create_collections` call, rather than on
3078    /// demand, is more efficient as it reduces the number of writes to durable storage. It also
3079    /// allows subsequent bootstrap logic to fetch metadata (such as frontiers) of arbitrary
3080    /// storage collections, without needing to worry about dependency order.
3081    ///
3082    /// `migrated_storage_collections` is a set of builtin storage collections that have been
3083    /// migrated and should be handled specially.
3084    #[instrument]
3085    async fn bootstrap_storage_collections(
3086        &mut self,
3087        migrated_storage_collections: &BTreeSet<CatalogItemId>,
3088    ) {
3089        let catalog = self.catalog();
3090
3091        let source_desc = |object_id: GlobalId,
3092                           data_source: &DataSourceDesc,
3093                           desc: &RelationDesc,
3094                           timeline: &Timeline| {
3095            let data_source = match data_source.clone() {
3096                // Re-announce the source description.
3097                DataSourceDesc::Ingestion { desc, cluster_id } => {
3098                    let desc = desc.into_inline_connection(catalog.state());
3099                    let ingestion = IngestionDescription::new(desc, cluster_id, object_id);
3100                    DataSource::Ingestion(ingestion)
3101                }
3102                DataSourceDesc::OldSyntaxIngestion {
3103                    desc,
3104                    progress_subsource,
3105                    data_config,
3106                    details,
3107                    cluster_id,
3108                } => {
3109                    let desc = desc.into_inline_connection(catalog.state());
3110                    let data_config = data_config.into_inline_connection(catalog.state());
3111                    // TODO(parkmycar): We should probably check the type here, but I'm not sure if
3112                    // this will always be a Source or a Table.
3113                    let progress_subsource =
3114                        catalog.get_entry(&progress_subsource).latest_global_id();
3115                    let mut ingestion =
3116                        IngestionDescription::new(desc, cluster_id, progress_subsource);
3117                    let legacy_export = SourceExport {
3118                        storage_metadata: (),
3119                        data_config,
3120                        details,
3121                    };
3122                    ingestion.source_exports.insert(object_id, legacy_export);
3123
3124                    DataSource::Ingestion(ingestion)
3125                }
3126                DataSourceDesc::IngestionExport {
3127                    ingestion_id,
3128                    external_reference: _,
3129                    details,
3130                    data_config,
3131                } => {
3132                    // TODO(parkmycar): We should probably check the type here, but I'm not sure if
3133                    // this will always be a Source or a Table.
3134                    let ingestion_id = catalog.get_entry(&ingestion_id).latest_global_id();
3135
3136                    DataSource::IngestionExport {
3137                        ingestion_id,
3138                        details,
3139                        data_config: data_config.into_inline_connection(catalog.state()),
3140                    }
3141                }
3142                DataSourceDesc::Webhook { .. } => DataSource::Webhook,
3143                DataSourceDesc::Progress => DataSource::Progress,
3144                DataSourceDesc::Introspection(introspection) => {
3145                    DataSource::Introspection(introspection)
3146                }
3147                DataSourceDesc::Catalog => DataSource::Other,
3148            };
3149            CollectionDescription {
3150                desc: desc.clone(),
3151                data_source,
3152                since: None,
3153                timeline: Some(timeline.clone()),
3154                primary: None,
3155            }
3156        };
3157
3158        let mut compute_collections = vec![];
3159        let mut collections = vec![];
3160        for entry in catalog.entries() {
3161            match entry.item() {
3162                CatalogItem::Source(source) => {
3163                    collections.push((
3164                        source.global_id(),
3165                        source_desc(
3166                            source.global_id(),
3167                            &source.data_source,
3168                            &source.desc,
3169                            &source.timeline,
3170                        ),
3171                    ));
3172                }
3173                CatalogItem::Table(table) => {
3174                    match &table.data_source {
3175                        TableDataSource::TableWrites { defaults: _ } => {
3176                            let versions: BTreeMap<_, _> = table
3177                                .collection_descs()
3178                                .map(|(gid, version, desc)| (version, (gid, desc)))
3179                                .collect();
3180                            let collection_descs = versions.iter().map(|(version, (gid, desc))| {
3181                                let next_version = version.bump();
3182                                let primary_collection =
3183                                    versions.get(&next_version).map(|(gid, _desc)| gid).copied();
3184                                let mut collection_desc =
3185                                    CollectionDescription::for_table(desc.clone());
3186                                collection_desc.primary = primary_collection;
3187
3188                                (*gid, collection_desc)
3189                            });
3190                            collections.extend(collection_descs);
3191                        }
3192                        TableDataSource::DataSource {
3193                            desc: data_source_desc,
3194                            timeline,
3195                        } => {
3196                            // TODO(alter_table): Support versioning tables that read from sources.
3197                            soft_assert_eq_or_log!(table.collections.len(), 1);
3198                            let collection_descs =
3199                                table.collection_descs().map(|(gid, _version, desc)| {
3200                                    (
3201                                        gid,
3202                                        source_desc(
3203                                            entry.latest_global_id(),
3204                                            data_source_desc,
3205                                            &desc,
3206                                            timeline,
3207                                        ),
3208                                    )
3209                                });
3210                            collections.extend(collection_descs);
3211                        }
3212                    };
3213                }
3214                CatalogItem::MaterializedView(mv) => {
3215                    let collection_descs = mv.collection_descs().map(|(gid, _version, desc)| {
3216                        let collection_desc =
3217                            CollectionDescription::for_other(desc, mv.initial_as_of.clone());
3218                        (gid, collection_desc)
3219                    });
3220
3221                    collections.extend(collection_descs);
3222                    compute_collections.push((mv.global_id_writes(), mv.desc.latest()));
3223                }
3224                CatalogItem::Sink(sink) => {
3225                    let storage_sink_from_entry = self.catalog().get_entry_by_global_id(&sink.from);
3226                    let from_desc = storage_sink_from_entry
3227                        .relation_desc()
3228                        .expect("sinks can only be built on items with descs")
3229                        .into_owned();
3230                    let collection_desc = CollectionDescription {
3231                        // TODO(sinks): make generic once we have more than one sink type.
3232                        desc: KAFKA_PROGRESS_DESC.clone(),
3233                        data_source: DataSource::Sink {
3234                            desc: ExportDescription {
3235                                sink: StorageSinkDesc {
3236                                    from: sink.from,
3237                                    from_desc,
3238                                    connection: sink
3239                                        .connection
3240                                        .clone()
3241                                        .into_inline_connection(self.catalog().state()),
3242                                    envelope: sink.envelope,
3243                                    as_of: Antichain::from_elem(Timestamp::minimum()),
3244                                    with_snapshot: sink.with_snapshot,
3245                                    version: sink.version,
3246                                    from_storage_metadata: (),
3247                                    to_storage_metadata: (),
3248                                    commit_interval: sink.commit_interval,
3249                                },
3250                                instance_id: sink.cluster_id,
3251                            },
3252                        },
3253                        since: None,
3254                        timeline: None,
3255                        primary: None,
3256                    };
3257                    collections.push((sink.global_id, collection_desc));
3258                }
3259                CatalogItem::Log(_)
3260                | CatalogItem::View(_)
3261                | CatalogItem::Index(_)
3262                | CatalogItem::Type(_)
3263                | CatalogItem::Func(_)
3264                | CatalogItem::Secret(_)
3265                | CatalogItem::Connection(_) => (),
3266            }
3267        }
3268
3269        let register_ts = if self.controller.read_only() {
3270            self.get_local_read_ts().await
3271        } else {
3272            // Getting a write timestamp bumps the write timestamp in the
3273            // oracle, which we're not allowed in read-only mode.
3274            self.get_local_write_ts().await.timestamp
3275        };
3276
3277        let storage_metadata = self.catalog.state().storage_metadata();
3278        let migrated_storage_collections = migrated_storage_collections
3279            .into_iter()
3280            .flat_map(|item_id| self.catalog.get_entry(item_id).global_ids())
3281            .collect();
3282
3283        // Before possibly creating collections, make sure their schemas are correct.
3284        //
3285        // Across different versions of Materialize the nullability of columns can change based on
3286        // updates to our optimizer.
3287        self.controller
3288            .storage
3289            .evolve_nullability_for_bootstrap(storage_metadata, compute_collections)
3290            .await
3291            .unwrap_or_terminate("cannot fail to evolve collections");
3292
3293        // New builtin storage collections are by default created with [0] since/upper frontiers.
3294        // For collections that have dependencies on other collections (MVs, CTs), this can violate
3295        // the frontier invariants assumed by as-of selection. For example, as-of selection expects
3296        // to be able to pick up computing a materialized view from its most recent upper, but if
3297        // that upper is [0] it's likely that the required times are not available anymore in the
3298        // MV inputs.
3299        //
3300        // To avoid violating frontier invariants, we need to bump their sinces to times greater
3301        // than all of their upstream storage inputs. To know the since of a storage input, it has
3302        // to be registered with the storage controller first. Thus we register collections in
3303        // layers: Each iteration registers the collections whose dependencies are all already
3304        // registered.
3305        let mut pending: BTreeMap<_, _> = collections.into_iter().collect();
3306
3307        // Precompute storage-collection dependencies for each collection.
3308        let transitive_dep_gids: BTreeMap<_, _> = pending
3309            .keys()
3310            .map(|gid| {
3311                let entry = self.catalog.get_entry_by_global_id(gid);
3312                let item_id = entry.id();
3313                let deps = self.catalog.state().transitive_uses(item_id);
3314                let dep_gids: BTreeSet<_> = deps
3315                    // Ignore self-dependencies. For example, `transitive_uses` includes the input ID,
3316                    // and CTs can depend on themselves.
3317                    .filter(|dep_id| *dep_id != item_id)
3318                    .map(|dep_id| self.catalog.get_entry(&dep_id).latest_global_id())
3319                    // Ignore dependencies on objects that are not storage collections.
3320                    .filter(|dep_gid| pending.contains_key(dep_gid))
3321                    .collect();
3322                (*gid, dep_gids)
3323            })
3324            .collect();
3325
3326        while !pending.is_empty() {
3327            // Drain collections whose dependencies have all been registered already
3328            // (i.e., are not in `pending`).
3329            let ready_gids: BTreeSet<_> = pending
3330                .keys()
3331                .filter(|gid| {
3332                    let mut deps = transitive_dep_gids[gid].iter();
3333                    !deps.any(|dep_gid| pending.contains_key(dep_gid))
3334                })
3335                .copied()
3336                .collect();
3337            let mut ready: Vec<_> = pending
3338                .extract_if(.., |gid, _| ready_gids.contains(gid))
3339                .collect();
3340
3341            // Bump sinces of builtin collections.
3342            for (gid, collection) in &mut ready {
3343                // Don't silently overwrite an explicitly specified `since`.
3344                if !gid.is_system() || collection.since.is_some() {
3345                    continue;
3346                }
3347
3348                let mut derived_since = Antichain::from_elem(Timestamp::MIN);
3349                for dep_gid in &transitive_dep_gids[gid] {
3350                    let (since, _) = self
3351                        .controller
3352                        .storage
3353                        .collection_frontiers(*dep_gid)
3354                        .expect("previously registered");
3355                    derived_since.join_assign(&since);
3356                }
3357                collection.since = Some(derived_since);
3358            }
3359
3360            if ready.is_empty() {
3361                soft_panic_or_log!(
3362                    "cycle in storage collections: {:?}",
3363                    pending.keys().collect::<Vec<_>>(),
3364                );
3365                // We get here only due to a bug. Rather than crash-looping, we try our best to
3366                // reach a sane state by attempting to register all the remaining collections at
3367                // once.
3368                ready = mem::take(&mut pending).into_iter().collect();
3369            }
3370
3371            self.controller
3372                .storage
3373                .create_collections_for_bootstrap(
3374                    storage_metadata,
3375                    Some(register_ts),
3376                    ready,
3377                    &migrated_storage_collections,
3378                )
3379                .await
3380                .unwrap_or_terminate("cannot fail to create collections");
3381        }
3382
3383        if !self.controller.read_only() {
3384            self.apply_local_write(register_ts).await;
3385        }
3386    }
3387
3388    /// Returns the current list of catalog entries, sorted into an appropriate order for
3389    /// bootstrapping.
3390    ///
3391    /// The returned entries are in dependency order. Indexes are sorted immediately after the
3392    /// objects they index, to ensure that all dependants of these indexed objects can make use of
3393    /// the respective indexes.
3394    fn bootstrap_sort_catalog_entries(&self) -> Vec<CatalogEntry> {
3395        let mut indexes_on = BTreeMap::<_, Vec<_>>::new();
3396        let mut non_indexes = Vec::new();
3397        for entry in self.catalog().entries().cloned() {
3398            if let Some(index) = entry.index() {
3399                let on = self.catalog().get_entry_by_global_id(&index.on);
3400                indexes_on.entry(on.id()).or_default().push(entry);
3401            } else {
3402                non_indexes.push(entry);
3403            }
3404        }
3405
3406        let key_fn = |entry: &CatalogEntry| entry.id;
3407        let dependencies_fn = |entry: &CatalogEntry| entry.uses();
3408        sort_topological(&mut non_indexes, key_fn, dependencies_fn);
3409
3410        let mut result = Vec::new();
3411        for entry in non_indexes {
3412            let id = entry.id();
3413            result.push(entry);
3414            if let Some(mut indexes) = indexes_on.remove(&id) {
3415                result.append(&mut indexes);
3416            }
3417        }
3418
3419        soft_assert_or_log!(
3420            indexes_on.is_empty(),
3421            "indexes with missing dependencies: {indexes_on:?}",
3422        );
3423
3424        result
3425    }
3426
3427    /// Invokes the optimizer on all indexes and materialized views in the catalog and inserts the
3428    /// resulting dataflow plans into the catalog state.
3429    ///
3430    /// `ordered_catalog_entries` must be sorted in dependency order, with dependencies ordered
3431    /// before their dependants.
3432    ///
3433    /// This method does not perform timestamp selection for the dataflows, nor does it create them
3434    /// in the compute controller. Both of these steps happen later during bootstrapping.
3435    ///
3436    /// Returns a map of expressions that were not cached.
3437    #[instrument]
3438    fn bootstrap_dataflow_plans(
3439        &mut self,
3440        ordered_catalog_entries: &[CatalogEntry],
3441        mut cached_global_exprs: BTreeMap<GlobalId, GlobalExpressions>,
3442    ) -> Result<BTreeMap<GlobalId, GlobalExpressions>, AdapterError> {
3443        // The optimizer expects to be able to query its `ComputeInstanceSnapshot` for
3444        // collections the current dataflow can depend on. But since we don't yet install anything
3445        // on compute instances, the snapshot information is incomplete. We fix that by manually
3446        // updating `ComputeInstanceSnapshot` objects to ensure they contain collections previously
3447        // optimized.
3448        let mut instance_snapshots = BTreeMap::new();
3449        let mut uncached_expressions = BTreeMap::new();
3450
3451        let optimizer_config = |catalog: &Catalog, cluster_id| {
3452            let system_config = catalog.system_config();
3453            let overrides = catalog.get_cluster(cluster_id).config.features();
3454            OptimizerConfig::from(system_config)
3455                .override_from(&overrides)
3456                // A cluster-scoped LaunchDarkly rule beats a manual `FEATURES`
3457                // pin.
3458                .override_from(
3459                    &catalog
3460                        .state()
3461                        .cluster_scoped_optimizer_overrides(cluster_id),
3462                )
3463        };
3464
3465        for entry in ordered_catalog_entries {
3466            match entry.item() {
3467                CatalogItem::Index(idx) => {
3468                    // Collect optimizer parameters.
3469                    let compute_instance =
3470                        instance_snapshots.entry(idx.cluster_id).or_insert_with(|| {
3471                            self.instance_snapshot(idx.cluster_id)
3472                                .expect("compute instance exists")
3473                        });
3474                    let global_id = idx.global_id();
3475
3476                    // The index may already be installed on the compute instance. For example,
3477                    // this is the case for introspection indexes.
3478                    if compute_instance.contains_collection(&global_id) {
3479                        continue;
3480                    }
3481
3482                    let optimizer_config = optimizer_config(&self.catalog, idx.cluster_id);
3483
3484                    let (optimized_plan, physical_plan, metainfo) =
3485                        match cached_global_exprs.remove(&global_id) {
3486                            Some(global_expressions)
3487                                if global_expressions.optimizer_features
3488                                    == optimizer_config.features =>
3489                            {
3490                                debug!("global expression cache hit for {global_id:?}");
3491                                (
3492                                    global_expressions.global_mir,
3493                                    global_expressions.physical_plan,
3494                                    global_expressions.dataflow_metainfos,
3495                                )
3496                            }
3497                            Some(_) | None => {
3498                                let (optimized_plan, global_lir_plan) = {
3499                                    // Build an optimizer for this INDEX.
3500                                    let mut optimizer = optimize::index::Optimizer::new(
3501                                        self.owned_catalog(),
3502                                        compute_instance.clone(),
3503                                        global_id,
3504                                        optimizer_config.clone(),
3505                                        self.optimizer_metrics(),
3506                                    );
3507
3508                                    // MIR ⇒ MIR optimization (global)
3509                                    let index_plan = optimize::index::Index::new(
3510                                        entry.name().clone(),
3511                                        idx.on,
3512                                        idx.keys.to_vec(),
3513                                    );
3514                                    let global_mir_plan = optimizer.optimize(index_plan)?;
3515                                    let optimized_plan = global_mir_plan.df_desc().clone();
3516
3517                                    // MIR ⇒ LIR lowering and LIR ⇒ LIR optimization (global)
3518                                    let global_lir_plan = optimizer.optimize(global_mir_plan)?;
3519
3520                                    (optimized_plan, global_lir_plan)
3521                                };
3522
3523                                let (physical_plan, metainfo) = global_lir_plan.unapply();
3524                                let metainfo = {
3525                                    // Pre-allocate a vector of transient GlobalIds for each notice.
3526                                    let notice_ids =
3527                                        std::iter::repeat_with(|| self.allocate_transient_id())
3528                                            .map(|(_item_id, gid)| gid)
3529                                            .take(metainfo.optimizer_notices.len())
3530                                            .collect::<Vec<_>>();
3531                                    // Return a metainfo with rendered notices.
3532                                    self.catalog().render_notices(
3533                                        metainfo,
3534                                        notice_ids,
3535                                        Some(idx.global_id()),
3536                                    )
3537                                };
3538                                uncached_expressions.insert(
3539                                    global_id,
3540                                    GlobalExpressions {
3541                                        global_mir: optimized_plan.clone(),
3542                                        physical_plan: physical_plan.clone(),
3543                                        dataflow_metainfos: metainfo.clone(),
3544                                        optimizer_features: optimizer_config.features.clone(),
3545                                    },
3546                                );
3547                                (optimized_plan, physical_plan, metainfo)
3548                            }
3549                        };
3550
3551                    let catalog = self.catalog_mut();
3552                    catalog.set_optimized_plan(idx.global_id(), optimized_plan);
3553                    catalog.set_physical_plan(idx.global_id(), physical_plan);
3554                    catalog.set_dataflow_metainfo(idx.global_id(), metainfo);
3555
3556                    compute_instance.insert_collection(idx.global_id());
3557                }
3558                CatalogItem::MaterializedView(mv) => {
3559                    // Collect optimizer parameters.
3560                    let compute_instance =
3561                        instance_snapshots.entry(mv.cluster_id).or_insert_with(|| {
3562                            self.instance_snapshot(mv.cluster_id)
3563                                .expect("compute instance exists")
3564                        });
3565                    let global_id = mv.global_id_writes();
3566
3567                    let optimizer_config = optimizer_config(&self.catalog, mv.cluster_id);
3568
3569                    let (optimized_plan, physical_plan, metainfo) = match cached_global_exprs
3570                        .remove(&global_id)
3571                    {
3572                        Some(global_expressions)
3573                            if global_expressions.optimizer_features
3574                                == optimizer_config.features =>
3575                        {
3576                            debug!("global expression cache hit for {global_id:?}");
3577                            (
3578                                global_expressions.global_mir,
3579                                global_expressions.physical_plan,
3580                                global_expressions.dataflow_metainfos,
3581                            )
3582                        }
3583                        Some(_) | None => {
3584                            let (_, internal_view_id) = self.allocate_transient_id();
3585                            let debug_name = self
3586                                .catalog()
3587                                .resolve_full_name(entry.name(), None)
3588                                .to_string();
3589
3590                            let (optimized_plan, global_lir_plan) = {
3591                                // Build an optimizer for this MATERIALIZED VIEW.
3592                                let mut optimizer = optimize::materialized_view::Optimizer::new(
3593                                    self.owned_catalog().as_optimizer_catalog(),
3594                                    compute_instance.clone(),
3595                                    global_id,
3596                                    internal_view_id,
3597                                    mv.desc.latest().iter_names().cloned().collect(),
3598                                    mv.non_null_assertions.clone(),
3599                                    mv.refresh_schedule.clone(),
3600                                    debug_name,
3601                                    optimizer_config.clone(),
3602                                    self.optimizer_metrics(),
3603                                );
3604
3605                                // MIR ⇒ MIR optimization (global)
3606                                // We make sure to use the HIR SQL type (since MIR SQL types may not be coherent).
3607                                let typ = infer_sql_type_for_catalog(
3608                                    &mv.raw_expr,
3609                                    &mv.locally_optimized_expr.as_ref().clone(),
3610                                );
3611                                let global_mir_plan = optimizer
3612                                    .optimize((mv.locally_optimized_expr.as_ref().clone(), typ))?;
3613                                let optimized_plan = global_mir_plan.df_desc().clone();
3614
3615                                // MIR ⇒ LIR lowering and LIR ⇒ LIR optimization (global)
3616                                let global_lir_plan = optimizer.optimize(global_mir_plan)?;
3617
3618                                (optimized_plan, global_lir_plan)
3619                            };
3620
3621                            let (physical_plan, metainfo) = global_lir_plan.unapply();
3622                            let metainfo = {
3623                                // Pre-allocate a vector of transient GlobalIds for each notice.
3624                                let notice_ids =
3625                                    std::iter::repeat_with(|| self.allocate_transient_id())
3626                                        .map(|(_item_id, global_id)| global_id)
3627                                        .take(metainfo.optimizer_notices.len())
3628                                        .collect::<Vec<_>>();
3629                                // Return a metainfo with rendered notices.
3630                                self.catalog().render_notices(
3631                                    metainfo,
3632                                    notice_ids,
3633                                    Some(mv.global_id_writes()),
3634                                )
3635                            };
3636                            uncached_expressions.insert(
3637                                global_id,
3638                                GlobalExpressions {
3639                                    global_mir: optimized_plan.clone(),
3640                                    physical_plan: physical_plan.clone(),
3641                                    dataflow_metainfos: metainfo.clone(),
3642                                    optimizer_features: optimizer_config.features.clone(),
3643                                },
3644                            );
3645                            (optimized_plan, physical_plan, metainfo)
3646                        }
3647                    };
3648
3649                    let catalog = self.catalog_mut();
3650                    catalog.set_optimized_plan(mv.global_id_writes(), optimized_plan);
3651                    catalog.set_physical_plan(mv.global_id_writes(), physical_plan);
3652                    catalog.set_dataflow_metainfo(mv.global_id_writes(), metainfo);
3653
3654                    compute_instance.insert_collection(mv.global_id_writes());
3655                }
3656                CatalogItem::Table(_)
3657                | CatalogItem::Source(_)
3658                | CatalogItem::Log(_)
3659                | CatalogItem::View(_)
3660                | CatalogItem::Sink(_)
3661                | CatalogItem::Type(_)
3662                | CatalogItem::Func(_)
3663                | CatalogItem::Secret(_)
3664                | CatalogItem::Connection(_) => (),
3665            }
3666        }
3667
3668        Ok(uncached_expressions)
3669    }
3670
3671    /// Selects for each compute dataflow an as-of suitable for bootstrapping it.
3672    ///
3673    /// Returns a set of [`ReadHold`]s that ensures the read frontiers of involved collections stay
3674    /// in place and that must not be dropped before all compute dataflows have been created with
3675    /// the compute controller.
3676    ///
3677    /// This method expects all storage collections and dataflow plans to be available, so it must
3678    /// run after [`Coordinator::bootstrap_storage_collections`] and
3679    /// [`Coordinator::bootstrap_dataflow_plans`].
3680    async fn bootstrap_dataflow_as_ofs(&mut self) -> BTreeMap<GlobalId, ReadHold> {
3681        let mut catalog_ids = Vec::new();
3682        let mut dataflows = Vec::new();
3683        let mut read_policies = BTreeMap::new();
3684        for entry in self.catalog.entries() {
3685            let gid = match entry.item() {
3686                CatalogItem::Index(idx) => idx.global_id(),
3687                CatalogItem::MaterializedView(mv) => mv.global_id_writes(),
3688                CatalogItem::Table(_)
3689                | CatalogItem::Source(_)
3690                | CatalogItem::Log(_)
3691                | CatalogItem::View(_)
3692                | CatalogItem::Sink(_)
3693                | CatalogItem::Type(_)
3694                | CatalogItem::Func(_)
3695                | CatalogItem::Secret(_)
3696                | CatalogItem::Connection(_) => continue,
3697            };
3698            if let Some(plan) = self.catalog.try_get_physical_plan(&gid) {
3699                catalog_ids.push(gid);
3700                dataflows.push(plan.clone());
3701
3702                if let Some(compaction_window) = entry.item().initial_logical_compaction_window() {
3703                    read_policies.insert(gid, compaction_window.into());
3704                }
3705            }
3706        }
3707
3708        let read_ts = self.get_local_read_ts().await;
3709        let read_holds = as_of_selection::run(
3710            &mut dataflows,
3711            &read_policies,
3712            &*self.controller.storage_collections,
3713            read_ts,
3714            self.controller.read_only(),
3715        );
3716
3717        let catalog = self.catalog_mut();
3718        for (id, plan) in catalog_ids.into_iter().zip_eq(dataflows) {
3719            catalog.set_physical_plan(id, plan);
3720        }
3721
3722        read_holds
3723    }
3724
3725    /// Serves the coordinator, receiving commands from users over `cmd_rx`
3726    /// and feedback from dataflow workers over `feedback_rx`.
3727    ///
3728    /// You must call `bootstrap` before calling this method.
3729    ///
3730    /// BOXED FUTURE: As of Nov 2023 the returned Future from this function was 92KB. This would
3731    /// get stored on the stack which is bad for runtime performance, and blow up our stack usage.
3732    /// Because of that we purposefully move this Future onto the heap (i.e. Box it).
3733    fn serve(
3734        mut self,
3735        mut internal_cmd_rx: mpsc::UnboundedReceiver<Message>,
3736        mut strict_serializable_reads_rx: mpsc::UnboundedReceiver<(ConnectionId, PendingReadTxn)>,
3737        mut cmd_rx: mpsc::UnboundedReceiver<(OpenTelemetryContext, Command)>,
3738        group_commit_rx: appends::GroupCommitWaiter,
3739    ) -> LocalBoxFuture<'static, ()> {
3740        async move {
3741            // Watcher that listens for and reports cluster service status changes.
3742            let mut cluster_events = self.controller.events_stream();
3743            let last_message = Arc::new(Mutex::new(LastMessage {
3744                kind: "none",
3745                stmt: None,
3746            }));
3747
3748            let (idle_tx, mut idle_rx) = tokio::sync::mpsc::channel(1);
3749            let idle_metric = self.metrics.queue_busy_seconds.clone();
3750            let last_message_watchdog = Arc::clone(&last_message);
3751
3752            spawn(|| "coord watchdog", async move {
3753                // Every 5 seconds, attempt to measure how long it takes for the
3754                // coord select loop to be empty, because this message is the last
3755                // processed. If it is idle, this will result in some microseconds
3756                // of measurement.
3757                let mut interval = tokio::time::interval(Duration::from_secs(5));
3758                // If we end up having to wait more than 5 seconds for the coord to respond, then the
3759                // behavior of Delay results in the interval "restarting" from whenever we yield
3760                // instead of trying to catch up.
3761                interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
3762
3763                // Track if we become stuck to de-dupe error reporting.
3764                let mut coord_stuck = false;
3765
3766                loop {
3767                    interval.tick().await;
3768
3769                    // Wait for space in the channel, if we timeout then the coordinator is stuck!
3770                    let duration = tokio::time::Duration::from_secs(30);
3771                    let timeout = tokio::time::timeout(duration, idle_tx.reserve()).await;
3772                    let Ok(maybe_permit) = timeout else {
3773                        // Only log if we're newly stuck, to prevent logging repeatedly.
3774                        if !coord_stuck {
3775                            let last_message = last_message_watchdog.lock().expect("poisoned");
3776                            tracing::warn!(
3777                                last_message_kind = %last_message.kind,
3778                                last_message_sql = %last_message.stmt_to_string(),
3779                                "coordinator stuck for {duration:?}",
3780                            );
3781                        }
3782                        coord_stuck = true;
3783
3784                        continue;
3785                    };
3786
3787                    // We got a permit, we're not stuck!
3788                    if coord_stuck {
3789                        tracing::info!("Coordinator became unstuck");
3790                    }
3791                    coord_stuck = false;
3792
3793                    // If we failed to acquire a permit it's because we're shutting down.
3794                    let Ok(permit) = maybe_permit else {
3795                        break;
3796                    };
3797
3798                    permit.send(idle_metric.start_timer());
3799                }
3800            });
3801
3802            self.schedule_storage_usage_collection().await;
3803            self.schedule_arrangement_sizes_collection().await;
3804            self.spawn_privatelink_vpc_endpoints_watch_task();
3805            self.spawn_statement_logging_task();
3806            self.spawn_catalog_info_metrics_task();
3807            self.spawn_cluster_controller_task();
3808            flags::tracing_config(self.catalog.system_config()).apply(&self.tracing_handle);
3809
3810            // Report if the handling of a single message takes longer than this threshold.
3811            let warn_threshold = self
3812                .catalog()
3813                .system_config()
3814                .coord_slow_message_warn_threshold();
3815
3816            // How many messages we'd like to batch up before processing them. Must be > 0.
3817            const MESSAGE_BATCH: usize = 64;
3818            let mut messages = Vec::with_capacity(MESSAGE_BATCH);
3819            let mut cmd_messages = Vec::with_capacity(MESSAGE_BATCH);
3820
3821            let message_batch = self.metrics.message_batch.clone();
3822
3823            // A persisted `Notified` future for the linearize re-check signal.
3824            // It must outlive a single loop iteration and be re-`set` only after
3825            // it completes: a fresh `notified()` per iteration could drop a
3826            // wakeup that arrives while a higher-priority branch wins the same
3827            // poll, stranding pending reads. Keeping it pinned across iterations
3828            // leaves it registered, so no wakeup is lost.
3829            let linearize_reads_notify = Arc::clone(&self.linearize_reads_notify);
3830            let linearize_reads_notified = linearize_reads_notify.notified();
3831            tokio::pin!(linearize_reads_notified);
3832
3833            loop {
3834                // Before adding a branch to this select loop, please ensure that the branch is
3835                // cancellation safe and add a comment explaining why. You can refer here for more
3836                // info: https://docs.rs/tokio/latest/tokio/macro.select.html#cancellation-safety
3837                select! {
3838                    // We prioritize internal commands over other commands. However, we work through
3839                    // batches of commands in some branches of this select, which means that even if
3840                    // a command generates internal commands, we will work through the current batch
3841                    // before receiving a new batch of commands.
3842                    biased;
3843
3844                    // `recv_many()` on `UnboundedReceiver` is cancellation safe:
3845                    // https://docs.rs/tokio/1.38.0/tokio/sync/mpsc/struct.UnboundedReceiver.html#cancel-safety-1
3846                    // Receive a batch of commands.
3847                    _ = internal_cmd_rx.recv_many(&mut messages, MESSAGE_BATCH) => {},
3848                    // `next()` on any stream is cancel-safe:
3849                    // https://docs.rs/tokio-stream/0.1.9/tokio_stream/trait.StreamExt.html#cancel-safety
3850                    // Receive a single command.
3851                    Some(event) = cluster_events.next() => {
3852                        messages.push(Message::ClusterEvent(event))
3853                    },
3854                    // See [`mz_controller::Controller::Controller::ready`] for notes
3855                    // on why this is cancel-safe.
3856                    // Receive a single command.
3857                    () = self.controller.ready() => {
3858                        // NOTE: We don't get a `Readiness` back from `ready()`
3859                        // because the controller wants to keep it and it's not
3860                        // trivially `Clone` or `Copy`. Hence this accessor.
3861                        let controller = match self.controller.get_readiness() {
3862                            Readiness::Storage => ControllerReadiness::Storage,
3863                            Readiness::Compute => ControllerReadiness::Compute,
3864                            Readiness::Metrics(_) => ControllerReadiness::Metrics,
3865                            Readiness::Internal(_) => ControllerReadiness::Internal,
3866                            Readiness::NotReady => unreachable!("just signaled as ready"),
3867                        };
3868                        messages.push(Message::ControllerReady { controller });
3869                    }
3870                    // See [`appends::GroupCommitWaiter`] for notes on why this is cancel safe.
3871                    // Receive a single command.
3872                    permit = group_commit_rx.ready() => {
3873                        // If we happen to have batched exactly one user write, use
3874                        // that span so the `emit_trace_id_notice` hooks up.
3875                        // Otherwise, the best we can do is invent a new root span
3876                        // and make it follow from all the Spans in the pending
3877                        // writes.
3878                        let user_write_spans = self.pending_writes.iter().flat_map(|x| match x {
3879                            PendingWriteTxn::User{span, ..} => Some(span),
3880                            PendingWriteTxn::System{..} => None,
3881                        });
3882                        let span = match user_write_spans.exactly_one() {
3883                            Ok(span) => span.clone(),
3884                            Err(user_write_spans) => {
3885                                let span = info_span!(parent: None, "group_commit_notify");
3886                                for s in user_write_spans {
3887                                    span.follows_from(s);
3888                                }
3889                                span
3890                            }
3891                        };
3892                        messages.push(Message::GroupCommitInitiate(span, Some(permit)));
3893                    },
3894                    // `recv_many()` on `UnboundedReceiver` is cancellation safe:
3895                    // https://docs.rs/tokio/1.38.0/tokio/sync/mpsc/struct.UnboundedReceiver.html#cancel-safety-1
3896                    // Receive a batch of commands.
3897                    count = cmd_rx.recv_many(&mut cmd_messages, MESSAGE_BATCH) => {
3898                        if count == 0 {
3899                            break;
3900                        } else {
3901                            messages.extend(cmd_messages.drain(..).map(
3902                                |(otel_ctx, cmd)| Message::Command(otel_ctx, cmd),
3903                            ));
3904                        }
3905                    },
3906                    // `recv()` on `UnboundedReceiver` is cancellation safe:
3907                    // https://docs.rs/tokio/1.38.0/tokio/sync/mpsc/struct.UnboundedReceiver.html#cancel-safety
3908                    // Receive a single command.
3909                    Some(pending_read_txn) = strict_serializable_reads_rx.recv() => {
3910                        let mut pending_read_txns = vec![pending_read_txn];
3911                        while let Ok(pending_read_txn) = strict_serializable_reads_rx.try_recv() {
3912                            pending_read_txns.push(pending_read_txn);
3913                        }
3914                        for (conn_id, pending_read_txn) in pending_read_txns {
3915                            let prev = self
3916                                .pending_linearize_read_txns
3917                                .insert(conn_id, pending_read_txn);
3918                            soft_assert_or_log!(
3919                                prev.is_none(),
3920                                "connections can not have multiple concurrent reads, prev: {prev:?}"
3921                            )
3922                        }
3923                        messages.push(Message::LinearizeReads);
3924                    }
3925                    // `tick()` on `Interval` is cancel-safe:
3926                    // https://docs.rs/tokio/1.19.2/tokio/time/struct.Interval.html#cancel-safety
3927                    // Receive a single command.
3928                    _ = self.advance_timelines_interval.tick() => {
3929                        let span = info_span!(parent: None, "coord::advance_timelines_interval");
3930                        span.follows_from(Span::current());
3931
3932                        // Group commit sends an `AdvanceTimelines` message when
3933                        // done, which is what downgrades read holds. In
3934                        // read-only mode we send this message directly because
3935                        // we're not doing group commits.
3936                        if self.controller.read_only() {
3937                            messages.push(Message::AdvanceTimelines);
3938                        } else {
3939                            messages.push(Message::GroupCommitInitiate(span, None));
3940                        }
3941                    },
3942                    // Re-check pending strict serializable reads. Deliberately
3943                    // placed below the group commit branches above: a re-check
3944                    // only makes a read ready if the timestamp oracle has
3945                    // advanced, and the oracle only advances via group commit, so
3946                    // this must never win over (and thereby starve) group commit.
3947                    // `Notify` coalesces re-arms into a single wakeup, so even
3948                    // when a pending read sits just behind the oracle (re-armed
3949                    // sub-millisecond), the lower branches (including the idle
3950                    // watchdog) stay reachable. See the pin above for why the
3951                    // future is persisted rather than recreated per iteration.
3952                    () = linearize_reads_notified.as_mut() => {
3953                        linearize_reads_notified.set(linearize_reads_notify.notified());
3954                        messages.push(Message::LinearizeReads);
3955                    }
3956                    // `tick()` on `Interval` is cancel-safe:
3957                    // https://docs.rs/tokio/1.19.2/tokio/time/struct.Interval.html#cancel-safety
3958                    // Receive a single command.
3959                    _ = self.check_cluster_scheduling_policies_interval.tick() => {
3960                        messages.push(Message::CheckSchedulingPolicies);
3961                    },
3962
3963                    // `tick()` on `Interval` is cancel-safe:
3964                    // https://docs.rs/tokio/1.19.2/tokio/time/struct.Interval.html#cancel-safety
3965                    // Receive a single command.
3966                    _ = self.caught_up_check_interval.tick() => {
3967                        // We do this directly on the main loop instead of
3968                        // firing off a message. We are still in read-only mode,
3969                        // so optimizing for latency, not blocking the main loop
3970                        // is not that important.
3971                        self.maybe_check_caught_up().await;
3972
3973                        continue;
3974                    },
3975
3976                    // Process the idle metric at the lowest priority to sample queue non-idle time.
3977                    // `recv()` on `Receiver` is cancellation safe:
3978                    // https://docs.rs/tokio/1.8.0/tokio/sync/mpsc/struct.Receiver.html#cancel-safety
3979                    // Receive a single command.
3980                    timer = idle_rx.recv() => {
3981                        timer.expect("does not drop").observe_duration();
3982                        self.metrics
3983                            .message_handling
3984                            .with_label_values(&["watchdog"])
3985                            .observe(0.0);
3986                        continue;
3987                    }
3988                };
3989
3990                // Observe the number of messages we're processing at once.
3991                message_batch.observe(f64::cast_lossy(messages.len()));
3992
3993                for msg in messages.drain(..) {
3994                    // All message processing functions trace. Start a parent span
3995                    // for them to make it easy to find slow messages.
3996                    let msg_kind = msg.kind();
3997                    let span = span!(
3998                        target: "mz_adapter::coord::handle_message_loop",
3999                        Level::INFO,
4000                        "coord::handle_message",
4001                        kind = msg_kind
4002                    );
4003                    let otel_context = span.context().span().span_context().clone();
4004
4005                    // Record the last kind of message in case we get stuck. For
4006                    // execute commands, we additionally stash the user's SQL,
4007                    // statement, so we can log it in case we get stuck.
4008                    *last_message.lock().expect("poisoned") = LastMessage {
4009                        kind: msg_kind,
4010                        stmt: match &msg {
4011                            Message::Command(
4012                                _,
4013                                Command::Execute {
4014                                    portal_name,
4015                                    session,
4016                                    ..
4017                                },
4018                            ) => session
4019                                .get_portal_unverified(portal_name)
4020                                .and_then(|p| p.stmt.as_ref().map(Arc::clone)),
4021                            _ => None,
4022                        },
4023                    };
4024
4025                    let start = Instant::now();
4026                    self.handle_message(msg).instrument(span).await;
4027                    let duration = start.elapsed();
4028
4029                    self.metrics
4030                        .message_handling
4031                        .with_label_values(&[msg_kind])
4032                        .observe(duration.as_secs_f64());
4033
4034                    // If something is _really_ slow, print a trace id for debugging, if OTEL is enabled.
4035                    if duration > warn_threshold {
4036                        let trace_id = otel_context.is_valid().then(|| otel_context.trace_id());
4037                        tracing::error!(
4038                            ?msg_kind,
4039                            ?trace_id,
4040                            ?duration,
4041                            "very slow coordinator message"
4042                        );
4043                    }
4044                }
4045            }
4046            // Try and cleanup as a best effort. There may be some async tasks out there holding a
4047            // reference that prevents us from cleaning up.
4048            if let Some(catalog) = Arc::into_inner(self.catalog) {
4049                catalog.expire().await;
4050            }
4051        }
4052        .boxed_local()
4053    }
4054
4055    /// Obtain a read-only Catalog reference.
4056    fn catalog(&self) -> &Catalog {
4057        &self.catalog
4058    }
4059
4060    /// Obtain a read-only Catalog snapshot, suitable for giving out to
4061    /// non-Coordinator thread tasks.
4062    fn owned_catalog(&self) -> Arc<Catalog> {
4063        Arc::clone(&self.catalog)
4064    }
4065
4066    /// Obtain a handle to the optimizer metrics, suitable for giving
4067    /// out to non-Coordinator thread tasks.
4068    fn optimizer_metrics(&self) -> OptimizerMetrics {
4069        self.optimizer_metrics.clone()
4070    }
4071
4072    /// Obtain a writeable Catalog reference.
4073    fn catalog_mut(&mut self) -> &mut Catalog {
4074        // make_mut will cause any other Arc references (from owned_catalog) to
4075        // continue to be valid by cloning the catalog, putting it in a new Arc,
4076        // which lives at self._catalog. If there are no other Arc references,
4077        // then no clone is made, and it returns a reference to the existing
4078        // object. This makes this method and owned_catalog both very cheap: at
4079        // most one clone per catalog mutation, but only if there's a read-only
4080        // reference to it.
4081        Arc::make_mut(&mut self.catalog)
4082    }
4083
4084    /// Refills the user ID pool by allocating IDs from the catalog.
4085    ///
4086    /// Requests `max(min_count, batch_size)` IDs so the pool is never
4087    /// under-filled relative to the configured batch size.
4088    async fn refill_user_id_pool(&mut self, min_count: u64) -> Result<(), AdapterError> {
4089        let batch_size = USER_ID_POOL_BATCH_SIZE.get(self.catalog().system_config().dyncfgs());
4090        let to_allocate = min_count.max(u64::from(batch_size));
4091        let id_ts = self.get_catalog_write_ts().await;
4092        let ids = self.catalog().allocate_user_ids(to_allocate, id_ts).await?;
4093        if let (Some((first_id, _)), Some((last_id, _))) = (ids.first(), ids.last()) {
4094            let start = match first_id {
4095                CatalogItemId::User(id) => *id,
4096                other => {
4097                    return Err(AdapterError::Internal(format!(
4098                        "expected User CatalogItemId, got {other:?}"
4099                    )));
4100                }
4101            };
4102            let end = match last_id {
4103                CatalogItemId::User(id) => *id + 1, // exclusive upper bound
4104                other => {
4105                    return Err(AdapterError::Internal(format!(
4106                        "expected User CatalogItemId, got {other:?}"
4107                    )));
4108                }
4109            };
4110            self.user_id_pool.refill(start, end);
4111        } else {
4112            return Err(AdapterError::Internal(
4113                "catalog returned no user IDs".into(),
4114            ));
4115        }
4116        Ok(())
4117    }
4118
4119    /// Allocates a single user ID, refilling the pool from the catalog if needed.
4120    async fn allocate_user_id(&mut self) -> Result<(CatalogItemId, GlobalId), AdapterError> {
4121        if let Some(id) = self.user_id_pool.allocate() {
4122            return Ok((CatalogItemId::User(id), GlobalId::User(id)));
4123        }
4124        self.refill_user_id_pool(1).await?;
4125        let id = self.user_id_pool.allocate().expect("ID pool just refilled");
4126        Ok((CatalogItemId::User(id), GlobalId::User(id)))
4127    }
4128
4129    /// Allocates `count` user IDs, refilling the pool from the catalog if needed.
4130    async fn allocate_user_ids(
4131        &mut self,
4132        count: u64,
4133    ) -> Result<Vec<(CatalogItemId, GlobalId)>, AdapterError> {
4134        if self.user_id_pool.remaining() < count {
4135            self.refill_user_id_pool(count).await?;
4136        }
4137        let raw_ids = self
4138            .user_id_pool
4139            .allocate_many(count)
4140            .expect("pool has enough IDs after refill");
4141        Ok(raw_ids
4142            .into_iter()
4143            .map(|id| (CatalogItemId::User(id), GlobalId::User(id)))
4144            .collect())
4145    }
4146
4147    /// Obtain a reference to the coordinator's connection context.
4148    fn connection_context(&self) -> &ConnectionContext {
4149        self.controller.connection_context()
4150    }
4151
4152    /// Obtain a reference to the coordinator's secret reader, in an `Arc`.
4153    fn secrets_reader(&self) -> &Arc<dyn SecretsReader> {
4154        &self.connection_context().secrets_reader
4155    }
4156
4157    /// Publishes a notice message to all sessions.
4158    ///
4159    /// TODO(parkmycar): This code is dead, but is a nice parallel to [`Coordinator::broadcast_notice_tx`]
4160    /// so we keep it around.
4161    #[allow(dead_code)]
4162    pub(crate) fn broadcast_notice(&self, notice: AdapterNotice) {
4163        for meta in self.active_conns.values() {
4164            let _ = meta.notice_tx.send(notice.clone());
4165        }
4166    }
4167
4168    /// Returns a closure that will publish a notice to all sessions that were active at the time
4169    /// this method was called.
4170    pub(crate) fn broadcast_notice_tx(
4171        &self,
4172    ) -> Box<dyn FnOnce(AdapterNotice) -> () + Send + 'static> {
4173        let senders: Vec<_> = self
4174            .active_conns
4175            .values()
4176            .map(|meta| meta.notice_tx.clone())
4177            .collect();
4178        Box::new(move |notice| {
4179            for tx in senders {
4180                let _ = tx.send(notice.clone());
4181            }
4182        })
4183    }
4184
4185    pub(crate) fn active_conns(&self) -> &BTreeMap<ConnectionId, ConnMeta> {
4186        &self.active_conns
4187    }
4188
4189    #[instrument(level = "debug")]
4190    pub(crate) fn retire_execution(
4191        &mut self,
4192        reason: StatementEndedExecutionReason,
4193        ctx_extra: ExecuteContextExtra,
4194    ) {
4195        if let Some(uuid) = ctx_extra.retire() {
4196            self.end_statement_execution(uuid, reason);
4197        }
4198    }
4199
4200    /// Creates a new dataflow builder from the catalog and indexes in `self`.
4201    #[instrument(level = "debug")]
4202    pub fn dataflow_builder(&self, instance: ComputeInstanceId) -> DataflowBuilder<'_> {
4203        let compute = self
4204            .instance_snapshot(instance)
4205            .expect("compute instance does not exist");
4206        DataflowBuilder::new(self.catalog().state(), compute)
4207    }
4208
4209    /// Return a reference-less snapshot to the indicated compute instance.
4210    pub fn instance_snapshot(
4211        &self,
4212        id: ComputeInstanceId,
4213    ) -> Result<ComputeInstanceSnapshot, InstanceMissing> {
4214        ComputeInstanceSnapshot::new(&self.controller, id)
4215    }
4216
4217    /// Call into the compute controller to install a finalized dataflow, and
4218    /// initialize the read policies for its exported readable objects.
4219    ///
4220    /// # Panics
4221    ///
4222    /// Panics if dataflow creation fails.
4223    pub(crate) async fn ship_dataflow(
4224        &mut self,
4225        dataflow: DataflowDescription<LirRelationExpr>,
4226        instance: ComputeInstanceId,
4227        target_replica: Option<ReplicaId>,
4228    ) {
4229        self.try_ship_dataflow(dataflow, instance, target_replica)
4230            .await
4231            .unwrap_or_terminate("dataflow creation cannot fail");
4232    }
4233
4234    /// Call into the compute controller to install a finalized dataflow, and
4235    /// initialize the read policies for its exported readable objects.
4236    pub(crate) async fn try_ship_dataflow(
4237        &mut self,
4238        dataflow: DataflowDescription<LirRelationExpr>,
4239        instance: ComputeInstanceId,
4240        target_replica: Option<ReplicaId>,
4241    ) -> Result<(), DataflowCreationError> {
4242        // We must only install read policies for indexes, not for sinks.
4243        // Sinks are write-only compute collections that don't have read policies.
4244        let export_ids = dataflow.exported_index_ids().collect();
4245
4246        self.controller
4247            .compute
4248            .create_dataflow(instance, dataflow, target_replica)?;
4249
4250        self.initialize_compute_read_policies(export_ids, instance, CompactionWindow::Default)
4251            .await;
4252
4253        Ok(())
4254    }
4255
4256    /// Call into the compute controller to allow writes to the specified IDs
4257    /// from the specified instance. Calling this function multiple times and
4258    /// calling it on a read-only instance has no effect.
4259    pub(crate) fn allow_writes(&mut self, instance: ComputeInstanceId, id: GlobalId) {
4260        self.controller
4261            .compute
4262            .allow_writes(instance, id)
4263            .unwrap_or_terminate("allow_writes cannot fail");
4264    }
4265
4266    /// Like `ship_dataflow`, but also await on builtin table updates.
4267    pub(crate) async fn ship_dataflow_and_notice_builtin_table_updates(
4268        &mut self,
4269        dataflow: DataflowDescription<LirRelationExpr>,
4270        instance: ComputeInstanceId,
4271        notice_builtin_updates_fut: Option<BuiltinTableAppendNotify>,
4272        target_replica: Option<ReplicaId>,
4273    ) {
4274        if let Some(notice_builtin_updates_fut) = notice_builtin_updates_fut {
4275            let ship_dataflow_fut = self.ship_dataflow(dataflow, instance, target_replica);
4276            let ((), ()) =
4277                futures::future::join(notice_builtin_updates_fut, ship_dataflow_fut).await;
4278        } else {
4279            self.ship_dataflow(dataflow, instance, target_replica).await;
4280        }
4281    }
4282
4283    /// Install a _watch set_ in the controller that is automatically associated with the given
4284    /// connection id. The watchset will be automatically cleared if the connection terminates
4285    /// before the watchset completes.
4286    pub fn install_compute_watch_set(
4287        &mut self,
4288        conn_id: ConnectionId,
4289        objects: BTreeSet<GlobalId>,
4290        t: Timestamp,
4291        state: WatchSetResponse,
4292    ) -> Result<(), CollectionLookupError> {
4293        let ws_id = self.controller.install_compute_watch_set(objects, t)?;
4294        self.connection_watch_sets
4295            .entry(conn_id.clone())
4296            .or_default()
4297            .insert(ws_id);
4298        self.installed_watch_sets.insert(ws_id, (conn_id, state));
4299        Ok(())
4300    }
4301
4302    /// Install a _watch set_ in the controller that is automatically associated with the given
4303    /// connection id. The watchset will be automatically cleared if the connection terminates
4304    /// before the watchset completes.
4305    pub fn install_storage_watch_set(
4306        &mut self,
4307        conn_id: ConnectionId,
4308        objects: BTreeSet<GlobalId>,
4309        t: Timestamp,
4310        state: WatchSetResponse,
4311    ) -> Result<(), CollectionMissing> {
4312        let ws_id = self.controller.install_storage_watch_set(objects, t)?;
4313        self.connection_watch_sets
4314            .entry(conn_id.clone())
4315            .or_default()
4316            .insert(ws_id);
4317        self.installed_watch_sets.insert(ws_id, (conn_id, state));
4318        Ok(())
4319    }
4320
4321    /// Cancels pending watchsets associated with the provided connection id.
4322    pub fn cancel_pending_watchsets(&mut self, conn_id: &ConnectionId) {
4323        if let Some(ws_ids) = self.connection_watch_sets.remove(conn_id) {
4324            for ws_id in ws_ids {
4325                self.installed_watch_sets.remove(&ws_id);
4326            }
4327        }
4328    }
4329
4330    /// Returns the state of the [`Coordinator`] formatted as JSON.
4331    ///
4332    /// The returned value is not guaranteed to be stable and may change at any point in time.
4333    pub async fn dump(&self) -> Result<serde_json::Value, anyhow::Error> {
4334        // Note: We purposefully use the `Debug` formatting for the value of all fields in the
4335        // returned object as a tradeoff between usability and stability. `serde_json` will fail
4336        // to serialize an object if the keys aren't strings, so `Debug` formatting the values
4337        // prevents a future unrelated change from silently breaking this method.
4338
4339        let global_timelines: BTreeMap<_, _> = self
4340            .global_timelines
4341            .iter()
4342            .map(|(timeline, state)| (timeline.to_string(), format!("{state:?}")))
4343            .collect();
4344        let active_conns: BTreeMap<_, _> = self
4345            .active_conns
4346            .iter()
4347            .map(|(id, meta)| (id.unhandled().to_string(), format!("{meta:?}")))
4348            .collect();
4349        let txn_read_holds: BTreeMap<_, _> = self
4350            .txn_read_holds
4351            .iter()
4352            .map(|(id, capability)| (id.unhandled().to_string(), format!("{capability:?}")))
4353            .collect();
4354        let pending_peeks: BTreeMap<_, _> = self
4355            .pending_peeks
4356            .iter()
4357            .map(|(id, peek)| (id.to_string(), format!("{peek:?}")))
4358            .collect();
4359        let client_pending_peeks: BTreeMap<_, _> = self
4360            .client_pending_peeks
4361            .iter()
4362            .map(|(id, peek)| {
4363                let peek: BTreeMap<_, _> = peek
4364                    .iter()
4365                    .map(|(uuid, storage_id)| (uuid.to_string(), storage_id))
4366                    .collect();
4367                (id.to_string(), peek)
4368            })
4369            .collect();
4370        let pending_linearize_read_txns: BTreeMap<_, _> = self
4371            .pending_linearize_read_txns
4372            .iter()
4373            .map(|(id, read_txn)| (id.unhandled().to_string(), format!("{read_txn:?}")))
4374            .collect();
4375
4376        Ok(serde_json::json!({
4377            "global_timelines": global_timelines,
4378            "active_conns": active_conns,
4379            "txn_read_holds": txn_read_holds,
4380            "pending_peeks": pending_peeks,
4381            "client_pending_peeks": client_pending_peeks,
4382            "pending_linearize_read_txns": pending_linearize_read_txns,
4383            "controller": self.controller.dump().await?,
4384        }))
4385    }
4386
4387    /// Prune all storage usage events from the [`MZ_STORAGE_USAGE_BY_SHARD`] table that are older
4388    /// than `retention_period`.
4389    ///
4390    /// This method will read the entire contents of [`MZ_STORAGE_USAGE_BY_SHARD`] into memory
4391    /// which can be expensive.
4392    ///
4393    /// DO NOT call this method outside of startup. The safety of reading at the current oracle read
4394    /// timestamp and then writing at whatever the current write timestamp is (instead of
4395    /// `read_ts + 1`) relies on the fact that there are no outstanding writes during startup.
4396    ///
4397    /// Group commit, which this method uses to write the retractions, has builtin fencing, and we
4398    /// never commit retractions to [`MZ_STORAGE_USAGE_BY_SHARD`] outside of this method, which is
4399    /// only called once during startup. So we don't have to worry about double/invalid retractions.
4400    async fn prune_storage_usage_events_on_startup(&self, retention_period: Duration) {
4401        let item_id = self
4402            .catalog()
4403            .resolve_builtin_table(&MZ_STORAGE_USAGE_BY_SHARD);
4404        let global_id = self.catalog.get_entry(&item_id).latest_global_id();
4405        let read_ts = self.get_local_read_ts().await;
4406        let current_contents_fut = self
4407            .controller
4408            .storage_collections
4409            .snapshot(global_id, read_ts);
4410        let internal_cmd_tx = self.internal_cmd_tx.clone();
4411        spawn(|| "storage_usage_prune", async move {
4412            let mut current_contents = current_contents_fut
4413                .await
4414                .unwrap_or_terminate("cannot fail to fetch snapshot");
4415            differential_dataflow::consolidation::consolidate(&mut current_contents);
4416
4417            let cutoff_ts = u128::from(read_ts).saturating_sub(retention_period.as_millis());
4418            let mut expired = Vec::new();
4419            for (row, diff) in current_contents {
4420                assert_eq!(
4421                    diff, 1,
4422                    "consolidated contents should not contain retractions: ({row:#?}, {diff:#?})"
4423                );
4424                // This logic relies on the definition of `mz_storage_usage_by_shard` not changing.
4425                let collection_timestamp = row
4426                    .unpack()
4427                    .get(3)
4428                    .expect("definition of mz_storage_by_shard changed")
4429                    .unwrap_timestamptz();
4430                let collection_timestamp = collection_timestamp.timestamp_millis();
4431                let collection_timestamp: u128 = collection_timestamp
4432                    .try_into()
4433                    .expect("all collections happen after Jan 1 1970");
4434                if collection_timestamp < cutoff_ts {
4435                    debug!("pruning storage event {row:?}");
4436                    let builtin_update = BuiltinTableUpdate::row(item_id, row, Diff::MINUS_ONE);
4437                    expired.push(builtin_update);
4438                }
4439            }
4440
4441            // main thread has shut down.
4442            let _ = internal_cmd_tx.send(Message::StorageUsagePrune(expired));
4443        });
4444    }
4445
4446    /// Retracts `mz_object_arrangement_size_history` rows older than the
4447    /// `arrangement_size_history_retention_period` dyncfg.
4448    ///
4449    /// Must only run at startup: it reads at the oracle read timestamp and
4450    /// writes retractions at the current write timestamp, which is only safe
4451    /// when no other writes are in flight. See [the equivalent storage-usage
4452    /// pruner](Self::prune_storage_usage_events_on_startup) for the same
4453    /// reasoning.
4454    async fn prune_arrangement_sizes_history_on_startup(&self) {
4455        // The catalog server is not writable in read-only mode.
4456        if self.controller.read_only() {
4457            return;
4458        }
4459
4460        let retention_period = mz_adapter_types::dyncfgs::ARRANGEMENT_SIZE_HISTORY_RETENTION_PERIOD
4461            .get(self.catalog().system_config().dyncfgs());
4462        let item_id = self
4463            .catalog()
4464            .resolve_builtin_table(&mz_catalog::builtin::MZ_OBJECT_ARRANGEMENT_SIZE_HISTORY);
4465        let global_id = self.catalog.get_entry(&item_id).latest_global_id();
4466        let read_ts = self.get_local_read_ts().await;
4467        let current_contents_fut = self
4468            .controller
4469            .storage_collections
4470            .snapshot(global_id, read_ts);
4471        let internal_cmd_tx = self.internal_cmd_tx.clone();
4472        spawn(|| "arrangement_sizes_history_prune", async move {
4473            let mut current_contents = current_contents_fut
4474                .await
4475                .unwrap_or_terminate("cannot fail to fetch snapshot");
4476            differential_dataflow::consolidation::consolidate(&mut current_contents);
4477
4478            let cutoff_ts = u128::from(read_ts).saturating_sub(retention_period.as_millis());
4479            let expired =
4480                arrangement_sizes_expired_retractions(current_contents, cutoff_ts, item_id);
4481
4482            // TODO(arrangement-sizes): when the writeable-catalog-server
4483            // plumbing in https://github.com/MaterializeInc/materialize/pull/35436
4484            // lands, retract directly on `mz_catalog_server`.
4485            let _ = internal_cmd_tx.send(Message::ArrangementSizesPrune(expired));
4486        });
4487    }
4488
4489    fn current_credit_consumption_rate(&self) -> Numeric {
4490        self.catalog()
4491            .user_cluster_replicas()
4492            .filter_map(|replica| match &replica.config.location {
4493                ReplicaLocation::Managed(location) => Some(location.size_for_billing()),
4494                ReplicaLocation::Unmanaged(_) => None,
4495            })
4496            .map(|size| {
4497                self.catalog()
4498                    .cluster_replica_sizes()
4499                    .0
4500                    .get(size)
4501                    .expect("location size is validated against the cluster replica sizes")
4502                    .credits_per_hour
4503            })
4504            .sum()
4505    }
4506}
4507
4508/// Returns retraction updates for rows in a consolidated
4509/// `mz_object_arrangement_size_history` snapshot whose `collection_timestamp`
4510/// (column 3) is strictly before `cutoff_ts`.
4511///
4512/// Panics if any input row has `diff != 1`: the caller must consolidate first,
4513/// and a consolidated history table should never contain retractions because
4514/// the only source of retractions is this function itself.
4515fn arrangement_sizes_expired_retractions(
4516    rows: impl IntoIterator<Item = (mz_repr::Row, i64)>,
4517    cutoff_ts: u128,
4518    item_id: CatalogItemId,
4519) -> Vec<BuiltinTableUpdate> {
4520    let mut expired = Vec::new();
4521    for (row, diff) in rows {
4522        assert_eq!(
4523            diff, 1,
4524            "consolidated contents should not contain retractions: ({row:#?}, {diff:#?})"
4525        );
4526        let collection_timestamp = row
4527            .unpack()
4528            .get(3)
4529            .expect("definition of mz_object_arrangement_size_history changed")
4530            .unwrap_timestamptz()
4531            .timestamp_millis();
4532        let collection_timestamp: u128 = collection_timestamp
4533            .try_into()
4534            .expect("all collections happen after Jan 1 1970");
4535        if collection_timestamp < cutoff_ts {
4536            expired.push(BuiltinTableUpdate::row(item_id, row, Diff::MINUS_ONE));
4537        }
4538    }
4539    expired
4540}
4541
4542#[cfg(test)]
4543impl Coordinator {
4544    #[allow(dead_code)]
4545    async fn verify_ship_dataflow_no_error(
4546        &mut self,
4547        dataflow: DataflowDescription<LirRelationExpr>,
4548    ) {
4549        // `ship_dataflow_new` is not allowed to have a `Result` return because this function is
4550        // called after `catalog_transact`, after which no errors are allowed. This test exists to
4551        // prevent us from incorrectly teaching those functions how to return errors (which has
4552        // happened twice and is the motivation for this test).
4553
4554        // An arbitrary compute instance ID to satisfy the function calls below. Note that
4555        // this only works because this function will never run.
4556        let compute_instance = ComputeInstanceId::user(1).expect("1 is a valid ID");
4557
4558        let _: () = self.ship_dataflow(dataflow, compute_instance, None).await;
4559    }
4560}
4561
4562/// Contains information about the last message the [`Coordinator`] processed.
4563struct LastMessage {
4564    kind: &'static str,
4565    stmt: Option<Arc<Statement<Raw>>>,
4566}
4567
4568impl LastMessage {
4569    /// Returns a redacted version of the statement that is safe for logs.
4570    fn stmt_to_string(&self) -> Cow<'static, str> {
4571        self.stmt
4572            .as_ref()
4573            .map(|stmt| stmt.to_ast_string_redacted().into())
4574            .unwrap_or(Cow::Borrowed("<none>"))
4575    }
4576}
4577
4578impl fmt::Debug for LastMessage {
4579    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
4580        f.debug_struct("LastMessage")
4581            .field("kind", &self.kind)
4582            .field("stmt", &self.stmt_to_string())
4583            .finish()
4584    }
4585}
4586
4587impl Drop for LastMessage {
4588    fn drop(&mut self) {
4589        // Only print the last message if we're currently panicking, otherwise we'd spam our logs.
4590        if std::thread::panicking() {
4591            // If we're panicking theres no guarantee `tracing` still works, so print to stderr.
4592            eprintln!("Coordinator panicking, dumping last message\n{self:?}",);
4593        }
4594    }
4595}
4596
4597/// Serves the coordinator based on the provided configuration.
4598///
4599/// For a high-level description of the coordinator, see the [crate
4600/// documentation](crate).
4601///
4602/// Returns a handle to the coordinator and a client to communicate with the
4603/// coordinator.
4604///
4605/// BOXED FUTURE: As of Nov 2023 the returned Future from this function was 42KB. This would
4606/// get stored on the stack which is bad for runtime performance, and blow up our stack usage.
4607/// Because of that we purposefully move this Future onto the heap (i.e. Box it).
4608pub fn serve(
4609    Config {
4610        controller_config,
4611        controller_envd_epoch,
4612        mut storage,
4613        timestamp_oracle_url,
4614        unsafe_mode,
4615        all_features,
4616        build_info,
4617        environment_id,
4618        metrics_registry,
4619        now,
4620        secrets_controller,
4621        cloud_resource_controller,
4622        cluster_replica_sizes,
4623        builtin_system_cluster_config,
4624        builtin_catalog_server_cluster_config,
4625        builtin_probe_cluster_config,
4626        builtin_support_cluster_config,
4627        builtin_analytics_cluster_config,
4628        system_parameter_defaults,
4629        availability_zones,
4630        storage_usage_client,
4631        storage_usage_collection_interval,
4632        storage_usage_retention_period,
4633        segment_client,
4634        egress_addresses,
4635        aws_account_id,
4636        aws_privatelink_availability_zones,
4637        connection_context,
4638        connection_limit_callback,
4639        remote_system_parameters,
4640        webhook_concurrency_limit,
4641        http_host_name,
4642        tracing_handle,
4643        read_only_controllers,
4644        caught_up_trigger: clusters_caught_up_trigger,
4645        helm_chart_version,
4646        license_key,
4647        external_login_password_mz_system,
4648        force_builtin_schema_migration,
4649    }: Config,
4650) -> BoxFuture<'static, Result<(Handle, Client), AdapterError>> {
4651    async move {
4652        let coord_start = Instant::now();
4653        info!("startup: coordinator init: beginning");
4654        info!("startup: coordinator init: preamble beginning");
4655
4656        // Initializing the builtins can be an expensive process and consume a lot of memory. We
4657        // forcibly initialize it early while the stack is relatively empty to avoid stack
4658        // overflows later.
4659        let _builtins = LazyLock::force(&BUILTINS_STATIC);
4660
4661        let (cmd_tx, cmd_rx) = mpsc::unbounded_channel();
4662        let (internal_cmd_tx, internal_cmd_rx) = mpsc::unbounded_channel();
4663        let (strict_serializable_reads_tx, strict_serializable_reads_rx) =
4664            mpsc::unbounded_channel();
4665
4666        // Validate and process availability zones.
4667        if !availability_zones.iter().all_unique() {
4668            coord_bail!("availability zones must be unique");
4669        }
4670
4671        let aws_principal_context = match (
4672            aws_account_id,
4673            connection_context.aws_external_id_prefix.clone(),
4674        ) {
4675            (Some(aws_account_id), Some(aws_external_id_prefix)) => Some(AwsPrincipalContext {
4676                aws_account_id,
4677                aws_external_id_prefix,
4678            }),
4679            _ => None,
4680        };
4681
4682        let aws_privatelink_availability_zones = aws_privatelink_availability_zones
4683            .map(|azs_vec| BTreeSet::from_iter(azs_vec.iter().cloned()));
4684
4685        info!(
4686            "startup: coordinator init: preamble complete in {:?}",
4687            coord_start.elapsed()
4688        );
4689        let oracle_init_start = Instant::now();
4690        info!("startup: coordinator init: timestamp oracle init beginning");
4691
4692        let timestamp_oracle_config = timestamp_oracle_url
4693            .map(|url| TimestampOracleConfig::from_url(&url, &metrics_registry))
4694            .transpose()?;
4695        let mut initial_timestamps =
4696            get_initial_oracle_timestamps(&timestamp_oracle_config).await?;
4697
4698        // Insert an entry for the `EpochMilliseconds` timeline if one doesn't exist,
4699        // which will ensure that the timeline is initialized since it's required
4700        // by the system.
4701        initial_timestamps
4702            .entry(Timeline::EpochMilliseconds)
4703            .or_insert_with(mz_repr::Timestamp::minimum);
4704        let mut timestamp_oracles = BTreeMap::new();
4705        for (timeline, initial_timestamp) in initial_timestamps {
4706            Coordinator::ensure_timeline_state_with_initial_time(
4707                &timeline,
4708                initial_timestamp,
4709                now.clone(),
4710                timestamp_oracle_config.clone(),
4711                &mut timestamp_oracles,
4712                read_only_controllers,
4713            )
4714            .await;
4715        }
4716
4717        // Opening the durable catalog uses one or more timestamps without communicating with
4718        // the timestamp oracle. Here we make sure to apply the catalog upper with the timestamp
4719        // oracle to linearize future operations with opening the catalog.
4720        let catalog_upper = storage.current_upper().await;
4721        // Choose a time at which to boot. This is used, for example, to prune
4722        // old storage usage data or migrate audit log entries.
4723        //
4724        // This time is usually the current system time, but with protection
4725        // against backwards time jumps, even across restarts.
4726        let epoch_millis_oracle = &timestamp_oracles
4727            .get(&Timeline::EpochMilliseconds)
4728            .expect("inserted above")
4729            .oracle;
4730
4731        let mut boot_ts = if read_only_controllers {
4732            let read_ts = epoch_millis_oracle.read_ts().await;
4733            std::cmp::max(read_ts, catalog_upper)
4734        } else {
4735            // Getting/applying a write timestamp bumps the write timestamp in the
4736            // oracle, which we're not allowed in read-only mode.
4737            epoch_millis_oracle.apply_write(catalog_upper).await;
4738            epoch_millis_oracle.write_ts().await.timestamp
4739        };
4740
4741        info!(
4742            "startup: coordinator init: timestamp oracle init complete in {:?}",
4743            oracle_init_start.elapsed()
4744        );
4745
4746        let catalog_open_start = Instant::now();
4747        info!("startup: coordinator init: catalog open beginning");
4748        let persist_client = controller_config
4749            .persist_clients
4750            .open(controller_config.persist_location.clone())
4751            .await
4752            .context("opening persist client")?;
4753        let builtin_item_migration_config =
4754            BuiltinItemMigrationConfig {
4755                persist_client: persist_client.clone(),
4756                read_only: read_only_controllers,
4757                force_migration: force_builtin_schema_migration,
4758            }
4759        ;
4760        let OpenCatalogResult {
4761            mut catalog,
4762            migrated_storage_collections_0dt,
4763            new_builtin_collections,
4764            builtin_table_updates,
4765            cached_global_exprs,
4766            uncached_local_exprs,
4767        } = Catalog::open(mz_catalog::config::Config {
4768            storage,
4769            metrics_registry: &metrics_registry,
4770            state: mz_catalog::config::StateConfig {
4771                unsafe_mode,
4772                all_features,
4773                build_info,
4774                environment_id: environment_id.clone(),
4775                read_only: read_only_controllers,
4776                now: now.clone(),
4777                boot_ts: boot_ts.clone(),
4778                skip_migrations: false,
4779                cluster_replica_sizes,
4780                builtin_system_cluster_config,
4781                builtin_catalog_server_cluster_config,
4782                builtin_probe_cluster_config,
4783                builtin_support_cluster_config,
4784                builtin_analytics_cluster_config,
4785                system_parameter_defaults,
4786                remote_system_parameters,
4787                availability_zones,
4788                egress_addresses,
4789                aws_principal_context,
4790                aws_privatelink_availability_zones,
4791                connection_context,
4792                http_host_name,
4793                builtin_item_migration_config,
4794                persist_client: persist_client.clone(),
4795                enable_expression_cache_override: None,
4796                helm_chart_version,
4797                external_login_password_mz_system,
4798                license_key: license_key.clone(),
4799            },
4800        })
4801        .await?;
4802
4803        // Opening the catalog uses one or more timestamps, so push the boot timestamp up to the
4804        // current catalog upper.
4805        let catalog_upper = catalog.current_upper().await;
4806        boot_ts = std::cmp::max(boot_ts, catalog_upper);
4807
4808        if !read_only_controllers {
4809            epoch_millis_oracle.apply_write(boot_ts).await;
4810        }
4811
4812        info!(
4813            "startup: coordinator init: catalog open complete in {:?}",
4814            catalog_open_start.elapsed()
4815        );
4816
4817        let coord_thread_start = Instant::now();
4818        info!("startup: coordinator init: coordinator thread start beginning");
4819
4820        let session_id = catalog.config().session_id;
4821        let start_instant = catalog.config().start_instant;
4822
4823        // In order for the coordinator to support Rc and Refcell types, it cannot be
4824        // sent across threads. Spawn it in a thread and have this parent thread wait
4825        // for bootstrap completion before proceeding.
4826        let (bootstrap_tx, bootstrap_rx) = oneshot::channel();
4827        let handle = TokioHandle::current();
4828
4829        let metrics = Metrics::register_into(&metrics_registry);
4830        let metrics_clone = metrics.clone();
4831        let optimizer_metrics = OptimizerMetrics::register_into(
4832            &metrics_registry,
4833            catalog.system_config().optimizer_e2e_latency_warning_threshold(),
4834        );
4835        let segment_client_clone = segment_client.clone();
4836        let coord_now = now.clone();
4837        let advance_timelines_interval =
4838            tokio::time::interval(catalog.system_config().default_timestamp_interval());
4839        let mut check_scheduling_policies_interval = tokio::time::interval(
4840            catalog
4841                .system_config()
4842                .cluster_check_scheduling_policies_interval(),
4843        );
4844        check_scheduling_policies_interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
4845
4846        let clusters_caught_up_check_interval = if read_only_controllers {
4847            let dyncfgs = catalog.system_config().dyncfgs();
4848            let interval = WITH_0DT_DEPLOYMENT_CAUGHT_UP_CHECK_INTERVAL.get(dyncfgs);
4849
4850            let mut interval = tokio::time::interval(interval);
4851            interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
4852            interval
4853        } else {
4854            // When not in read-only mode, we don't do hydration checks. But we
4855            // still have to provide _some_ interval. This is large enough that
4856            // it doesn't matter.
4857            //
4858            // TODO(aljoscha): We cannot use Duration::MAX right now because of
4859            // https://github.com/tokio-rs/tokio/issues/6634. Use that once it's
4860            // fixed for good.
4861            let mut interval = tokio::time::interval(Duration::from_secs(60 * 60));
4862            interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
4863            interval
4864        };
4865
4866        let clusters_caught_up_check =
4867            clusters_caught_up_trigger.map(|trigger| {
4868                let mut exclude_collections: BTreeSet<GlobalId> =
4869                    new_builtin_collections.into_iter().collect();
4870
4871                // Migrated MVs can't make progress in read-only mode. Exclude them and all their
4872                // transitive dependents.
4873                //
4874                // TODO: Consider sending `allow_writes` for the dataflows of migrated MVs, which
4875                //       would allow them to make progress even in read-only mode. This doesn't
4876                //       work for MVs based on `mz_catalog_raw`, if the leader's version is less
4877                //       than v26.17, since before that version the catalog shard's frontier wasn't
4878                //       kept up-to-date with the current time. So this workaround has to remain in
4879                //       place upgrades from a version less than v26.17 are no longer supported.
4880                let mut todo: Vec<_> = migrated_storage_collections_0dt
4881                    .iter()
4882                    .filter(|id| {
4883                        catalog.state().get_entry(id).is_materialized_view()
4884                    })
4885                    .copied()
4886                    .collect();
4887                while let Some(item_id) = todo.pop() {
4888                    let entry = catalog.state().get_entry(&item_id);
4889                    exclude_collections.extend(entry.global_ids());
4890                    todo.extend_from_slice(entry.used_by());
4891                }
4892
4893                CaughtUpCheckContext {
4894                    trigger,
4895                    exclude_collections,
4896                    cluster_stability: BTreeMap::new(),
4897                }
4898            });
4899
4900        if let Some(TimestampOracleConfig::Postgres(pg_config)) =
4901            timestamp_oracle_config.as_ref()
4902        {
4903            // Apply settings from system vars as early as possible because some
4904            // of them are locked in right when an oracle is first opened!
4905            let pg_timestamp_oracle_params =
4906                flags::timestamp_oracle_config(catalog.system_config());
4907            pg_timestamp_oracle_params.apply(pg_config);
4908        }
4909
4910        // Register a callback so whenever the MAX_CONNECTIONS or SUPERUSER_RESERVED_CONNECTIONS
4911        // system variables change, we update our connection limits.
4912        let connection_limit_callback: Arc<dyn Fn(&SystemVars) + Send + Sync> =
4913            Arc::new(move |system_vars: &SystemVars| {
4914                let limit: u64 = system_vars.max_connections().cast_into();
4915                let superuser_reserved: u64 =
4916                    system_vars.superuser_reserved_connections().cast_into();
4917
4918                // If superuser_reserved > max_connections, prefer max_connections.
4919                //
4920                // In this scenario all normal users would be locked out because all connections
4921                // would be reserved for superusers so complain if this is the case.
4922                let superuser_reserved = if superuser_reserved >= limit {
4923                    tracing::warn!(
4924                        "superuser_reserved ({superuser_reserved}) is greater than max connections ({limit})!"
4925                    );
4926                    limit
4927                } else {
4928                    superuser_reserved
4929                };
4930
4931                (connection_limit_callback)(limit, superuser_reserved);
4932            });
4933        catalog.system_config_mut().register_callback(
4934            &mz_sql::session::vars::MAX_CONNECTIONS,
4935            Arc::clone(&connection_limit_callback),
4936        );
4937        catalog.system_config_mut().register_callback(
4938            &mz_sql::session::vars::SUPERUSER_RESERVED_CONNECTIONS,
4939            connection_limit_callback,
4940        );
4941
4942        let (group_commit_tx, group_commit_rx) = appends::notifier();
4943
4944        let parent_span = tracing::Span::current();
4945        let thread = thread::Builder::new()
4946            // The Coordinator thread tends to keep a lot of data on its stack. To
4947            // prevent a stack overflow we allocate a stack three times as big as the default
4948            // stack.
4949            .stack_size(3 * stack::STACK_SIZE)
4950            .name("coordinator".to_string())
4951            .spawn(move || {
4952                let span = info_span!(parent: parent_span, "coord::coordinator").entered();
4953
4954                let controller = handle
4955                    .block_on({
4956                        catalog.initialize_controller(
4957                            controller_config,
4958                            controller_envd_epoch,
4959                            read_only_controllers,
4960                        )
4961                    })
4962                    .unwrap_or_terminate("failed to initialize storage_controller");
4963                // Initializing the controller uses one or more timestamps, so push the boot timestamp up to the
4964                // current catalog upper.
4965                let catalog_upper = handle.block_on(catalog.current_upper());
4966                boot_ts = std::cmp::max(boot_ts, catalog_upper);
4967                if !read_only_controllers {
4968                    let epoch_millis_oracle = &timestamp_oracles
4969                        .get(&Timeline::EpochMilliseconds)
4970                        .expect("inserted above")
4971                        .oracle;
4972                    handle.block_on(epoch_millis_oracle.apply_write(boot_ts));
4973                }
4974
4975                let catalog = Arc::new(catalog);
4976
4977                let caching_secrets_reader = CachingSecretsReader::new(secrets_controller.reader());
4978                let mut coord = Coordinator {
4979                    controller,
4980                    catalog,
4981                    internal_cmd_tx,
4982                    group_commit_tx,
4983                    reconcile_now: Arc::new(Notify::new()),
4984                    strict_serializable_reads_tx,
4985                    linearize_reads_notify: Arc::new(Notify::new()),
4986                    global_timelines: timestamp_oracles,
4987                    transient_id_gen: Arc::new(TransientIdGen::new()),
4988                    active_conns: BTreeMap::new(),
4989                    txn_read_holds: Default::default(),
4990                    pending_peeks: BTreeMap::new(),
4991                    client_pending_peeks: BTreeMap::new(),
4992                    pending_linearize_read_txns: BTreeMap::new(),
4993                    serialized_ddl: LockedVecDeque::new(),
4994                    active_compute_sinks: BTreeMap::new(),
4995                    active_webhooks: BTreeMap::new(),
4996                    active_copies: BTreeMap::new(),
4997                    connection_cancel_watches: BTreeMap::new(),
4998                    introspection_subscribes: BTreeMap::new(),
4999                    write_locks: BTreeMap::new(),
5000                    deferred_write_ops: BTreeMap::new(),
5001                    pending_writes: Vec::new(),
5002                    advance_timelines_interval,
5003                    secrets_controller,
5004                    caching_secrets_reader,
5005                    cloud_resource_controller,
5006                    storage_usage_client,
5007                    storage_usage_collection_interval,
5008                    segment_client,
5009                    metrics,
5010                    catalog_info_metrics_registry: metrics_registry.clone(),
5011                    scoped_frontend: None,
5012                    optimizer_metrics,
5013                    tracing_handle,
5014                    statement_logging: StatementLogging::new(coord_now.clone()),
5015                    webhook_concurrency_limit,
5016                    timestamp_oracle_config,
5017                    check_cluster_scheduling_policies_interval: check_scheduling_policies_interval,
5018                    cluster_scheduling_decisions: BTreeMap::new(),
5019                    caught_up_check_interval: clusters_caught_up_check_interval,
5020                    caught_up_check: clusters_caught_up_check,
5021                    installed_watch_sets: BTreeMap::new(),
5022                    connection_watch_sets: BTreeMap::new(),
5023                    cluster_replica_statuses: ClusterReplicaStatuses::new(),
5024                    read_only_controllers,
5025                    buffered_builtin_table_updates: Some(Vec::new()),
5026                    license_key,
5027                    user_id_pool: IdPool::empty(),
5028                    persist_client,
5029                };
5030                let bootstrap = handle.block_on(async {
5031                    coord
5032                        .bootstrap(
5033                            boot_ts,
5034                            migrated_storage_collections_0dt,
5035                            builtin_table_updates,
5036                            cached_global_exprs,
5037                            uncached_local_exprs,
5038                        )
5039                        .await?;
5040                    coord
5041                        .controller
5042                        .remove_orphaned_replicas(
5043                            coord.catalog().get_next_user_replica_id().await?,
5044                            coord.catalog().get_next_system_replica_id().await?,
5045                        )
5046                        .await
5047                        .map_err(AdapterError::Orchestrator)?;
5048
5049                    if let Some(retention_period) = storage_usage_retention_period {
5050                        coord
5051                            .prune_storage_usage_events_on_startup(retention_period)
5052                            .await;
5053                    }
5054
5055                    coord.prune_arrangement_sizes_history_on_startup().await;
5056
5057                    Ok(())
5058                });
5059                let ok = bootstrap.is_ok();
5060                drop(span);
5061                bootstrap_tx
5062                    .send(bootstrap)
5063                    .expect("bootstrap_rx is not dropped until it receives this message");
5064                if ok {
5065                    handle.block_on(coord.serve(
5066                        internal_cmd_rx,
5067                        strict_serializable_reads_rx,
5068                        cmd_rx,
5069                        group_commit_rx,
5070                    ));
5071                }
5072            })
5073            .expect("failed to create coordinator thread");
5074        match bootstrap_rx
5075            .await
5076            .expect("bootstrap_tx always sends a message or panics/halts")
5077        {
5078            Ok(()) => {
5079                info!(
5080                    "startup: coordinator init: coordinator thread start complete in {:?}",
5081                    coord_thread_start.elapsed()
5082                );
5083                info!(
5084                    "startup: coordinator init: complete in {:?}",
5085                    coord_start.elapsed()
5086                );
5087                let handle = Handle {
5088                    session_id,
5089                    start_instant,
5090                    _thread: thread.join_on_drop(),
5091                };
5092                let client = Client::new(
5093                    build_info,
5094                    cmd_tx,
5095                    metrics_clone,
5096                    now,
5097                    environment_id,
5098                    segment_client_clone,
5099                );
5100                Ok((handle, client))
5101            }
5102            Err(e) => Err(e),
5103        }
5104    }
5105    .boxed()
5106}
5107
5108// Determines and returns the highest timestamp for each timeline, for all known
5109// timestamp oracle implementations.
5110//
5111// Initially, we did this so that we can switch between implementations of
5112// timestamp oracle, but now we also do this to determine a monotonic boot
5113// timestamp, a timestamp that does not regress across reboots.
5114//
5115// This mostly works, but there can be linearizability violations, because there
5116// is no central moment where we do distributed coordination for all oracle
5117// types. Working around this seems prohibitively hard, maybe even impossible so
5118// we have to live with this window of potential violations during the upgrade
5119// window (which is the only point where we should switch oracle
5120// implementations).
5121async fn get_initial_oracle_timestamps(
5122    timestamp_oracle_config: &Option<TimestampOracleConfig>,
5123) -> Result<BTreeMap<Timeline, Timestamp>, AdapterError> {
5124    let mut initial_timestamps = BTreeMap::new();
5125
5126    if let Some(config) = timestamp_oracle_config {
5127        let oracle_timestamps = config.get_all_timelines().await?;
5128
5129        let debug_msg = || {
5130            oracle_timestamps
5131                .iter()
5132                .map(|(timeline, ts)| format!("{:?} -> {}", timeline, ts))
5133                .join(", ")
5134        };
5135        info!(
5136            "current timestamps from the timestamp oracle: {}",
5137            debug_msg()
5138        );
5139
5140        for (timeline, ts) in oracle_timestamps {
5141            let entry = initial_timestamps
5142                .entry(Timeline::from_str(&timeline).expect("could not parse timeline"));
5143
5144            entry
5145                .and_modify(|current_ts| *current_ts = std::cmp::max(*current_ts, ts))
5146                .or_insert(ts);
5147        }
5148    } else {
5149        info!("no timestamp oracle configured!");
5150    };
5151
5152    let debug_msg = || {
5153        initial_timestamps
5154            .iter()
5155            .map(|(timeline, ts)| format!("{:?}: {}", timeline, ts))
5156            .join(", ")
5157    };
5158    info!("initial oracle timestamps: {}", debug_msg());
5159
5160    Ok(initial_timestamps)
5161}
5162
5163#[instrument]
5164pub async fn load_remote_system_parameters(
5165    storage: &mut Box<dyn OpenableDurableCatalogState>,
5166    system_parameter_sync_config: Option<SystemParameterSyncConfig>,
5167    system_parameter_sync_timeout: Duration,
5168) -> Result<Option<BTreeMap<String, String>>, AdapterError> {
5169    if let Some(system_parameter_sync_config) = system_parameter_sync_config {
5170        tracing::info!("parameter sync on boot: start sync");
5171
5172        // We intentionally block initial startup, potentially forever,
5173        // on initializing LaunchDarkly. This may seem scary, but the
5174        // alternative is even scarier. Over time, we expect that the
5175        // compiled-in default values for the system parameters will
5176        // drift substantially from the defaults configured in
5177        // LaunchDarkly, to the point that starting an environment
5178        // without loading the latest values from LaunchDarkly will
5179        // result in running an untested configuration.
5180        //
5181        // Note this only applies during initial startup. Restarting
5182        // after we've synced once only blocks for a maximum of
5183        // `FRONTEND_SYNC_TIMEOUT` on LaunchDarkly, as it seems
5184        // reasonable to assume that the last-synced configuration was
5185        // valid enough.
5186        //
5187        // This philosophy appears to provide a good balance between not
5188        // running untested configurations in production while also not
5189        // making LaunchDarkly a "tier 1" dependency for existing
5190        // environments.
5191        //
5192        // If this proves to be an issue, we could seek to address the
5193        // configuration drift in a different way--for example, by
5194        // writing a script that runs in CI nightly and checks for
5195        // deviation between the compiled Rust code and LaunchDarkly.
5196        //
5197        // If it is absolutely necessary to bring up a new environment
5198        // while LaunchDarkly is down, the following manual mitigation
5199        // can be performed:
5200        //
5201        //    1. Edit the environmentd startup parameters to omit the
5202        //       LaunchDarkly configuration.
5203        //    2. Boot environmentd.
5204        //    3. Use the catalog-debug tool to run `edit config "{\"key\":\"system_config_synced\"}" "{\"value\": 1}"`.
5205        //    4. Adjust any other parameters as necessary to avoid
5206        //       running a nonstandard configuration in production.
5207        //    5. Edit the environmentd startup parameters to restore the
5208        //       LaunchDarkly configuration, for when LaunchDarkly comes
5209        //       back online.
5210        //    6. Reboot environmentd.
5211        let mut params = SynchronizedParameters::new(SystemVars::default());
5212        let frontend_sync = async {
5213            let frontend = SystemParameterFrontend::from(&system_parameter_sync_config).await?;
5214            frontend.pull(&mut params);
5215            let ops = params
5216                .modified()
5217                .into_iter()
5218                .map(|param| {
5219                    let name = param.name;
5220                    let value = param.value;
5221                    tracing::info!(name, value, initial = true, "sync parameter");
5222                    (name, value)
5223                })
5224                .collect();
5225            tracing::info!("parameter sync on boot: end sync");
5226            Ok(Some(ops))
5227        };
5228        if !storage.has_system_config_synced_once().await? {
5229            frontend_sync.await
5230        } else {
5231            match mz_ore::future::timeout(system_parameter_sync_timeout, frontend_sync).await {
5232                Ok(ops) => Ok(ops),
5233                Err(TimeoutError::Inner(e)) => Err(e),
5234                Err(TimeoutError::DeadlineElapsed) => {
5235                    tracing::info!("parameter sync on boot: sync has timed out");
5236                    Ok(None)
5237                }
5238            }
5239        }
5240    } else {
5241        Ok(None)
5242    }
5243}
5244
5245#[derive(Debug)]
5246pub enum WatchSetResponse {
5247    StatementDependenciesReady(StatementLoggingId, StatementLifecycleEvent),
5248    AlterSinkReady(AlterSinkReadyContext),
5249    AlterMaterializedViewReady(AlterMaterializedViewReadyContext),
5250}
5251
5252#[derive(Debug)]
5253pub struct AlterSinkReadyContext {
5254    ctx: Option<ExecuteContext>,
5255    otel_ctx: OpenTelemetryContext,
5256    plan: AlterSinkPlan,
5257    plan_validity: PlanValidity,
5258    read_hold: ReadHolds,
5259}
5260
5261impl AlterSinkReadyContext {
5262    fn ctx(&mut self) -> &mut ExecuteContext {
5263        self.ctx.as_mut().expect("only cleared on drop")
5264    }
5265
5266    fn retire(mut self, result: Result<ExecuteResponse, AdapterError>) {
5267        self.ctx
5268            .take()
5269            .expect("only cleared on drop")
5270            .retire(result);
5271    }
5272}
5273
5274impl Drop for AlterSinkReadyContext {
5275    fn drop(&mut self) {
5276        if let Some(ctx) = self.ctx.take() {
5277            ctx.retire(Err(AdapterError::Canceled));
5278        }
5279    }
5280}
5281
5282#[derive(Debug)]
5283pub struct AlterMaterializedViewReadyContext {
5284    ctx: Option<ExecuteContext>,
5285    otel_ctx: OpenTelemetryContext,
5286    plan: plan::AlterMaterializedViewApplyReplacementPlan,
5287    plan_validity: PlanValidity,
5288}
5289
5290impl AlterMaterializedViewReadyContext {
5291    fn ctx(&mut self) -> &mut ExecuteContext {
5292        self.ctx.as_mut().expect("only cleared on drop")
5293    }
5294
5295    fn retire(mut self, result: Result<ExecuteResponse, AdapterError>) {
5296        self.ctx
5297            .take()
5298            .expect("only cleared on drop")
5299            .retire(result);
5300    }
5301}
5302
5303impl Drop for AlterMaterializedViewReadyContext {
5304    fn drop(&mut self) {
5305        if let Some(ctx) = self.ctx.take() {
5306            ctx.retire(Err(AdapterError::Canceled));
5307        }
5308    }
5309}
5310
5311/// A struct for tracking the ownership of a lock and a VecDeque to store to-be-done work after the
5312/// lock is freed.
5313#[derive(Debug)]
5314struct LockedVecDeque<T> {
5315    items: VecDeque<T>,
5316    lock: Arc<tokio::sync::Mutex<()>>,
5317}
5318
5319impl<T> LockedVecDeque<T> {
5320    pub fn new() -> Self {
5321        Self {
5322            items: VecDeque::new(),
5323            lock: Arc::new(tokio::sync::Mutex::new(())),
5324        }
5325    }
5326
5327    pub fn try_lock_owned(&self) -> Result<OwnedMutexGuard<()>, tokio::sync::TryLockError> {
5328        Arc::clone(&self.lock).try_lock_owned()
5329    }
5330
5331    pub fn is_empty(&self) -> bool {
5332        self.items.is_empty()
5333    }
5334
5335    pub fn push_back(&mut self, value: T) {
5336        self.items.push_back(value)
5337    }
5338
5339    pub fn pop_front(&mut self) -> Option<T> {
5340        self.items.pop_front()
5341    }
5342
5343    pub fn remove(&mut self, index: usize) -> Option<T> {
5344        self.items.remove(index)
5345    }
5346
5347    pub fn iter(&self) -> std::collections::vec_deque::Iter<'_, T> {
5348        self.items.iter()
5349    }
5350}
5351
5352#[derive(Debug)]
5353struct DeferredPlanStatement {
5354    ctx: ExecuteContext,
5355    ps: PlanStatement,
5356}
5357
5358#[derive(Debug)]
5359enum PlanStatement {
5360    Statement {
5361        stmt: Arc<Statement<Raw>>,
5362        params: Params,
5363    },
5364    Plan {
5365        plan: mz_sql::plan::Plan,
5366        resolved_ids: ResolvedIds,
5367        sql_impl_resolved_ids: ResolvedIds,
5368    },
5369}
5370
5371#[derive(Debug, Error)]
5372pub enum NetworkPolicyError {
5373    #[error("Access denied for address {0}")]
5374    AddressDenied(IpAddr),
5375    #[error("Access denied missing IP address")]
5376    MissingIp,
5377}
5378
5379pub(crate) fn validate_ip_with_policy_rules(
5380    ip: &IpAddr,
5381    rules: &Vec<NetworkPolicyRule>,
5382) -> Result<(), NetworkPolicyError> {
5383    // At the moment we're not handling action or direction
5384    // as those are only able to be "allow" and "ingress" respectively
5385    if rules.iter().any(|r| r.address.0.contains(ip)) {
5386        Ok(())
5387    } else {
5388        Err(NetworkPolicyError::AddressDenied(ip.clone()))
5389    }
5390}
5391
5392pub(crate) fn infer_sql_type_for_catalog(
5393    hir_expr: &HirRelationExpr,
5394    mir_expr: &MirRelationExpr,
5395) -> SqlRelationType {
5396    let mut typ = hir_expr.top_level_typ();
5397    typ.backport_nullability_and_keys(&mir_expr.typ());
5398    typ
5399}
5400
5401#[cfg(test)]
5402mod id_pool_tests {
5403    use super::IdPool;
5404
5405    #[mz_ore::test]
5406    fn test_empty_pool() {
5407        let mut pool = IdPool::empty();
5408        assert_eq!(pool.remaining(), 0);
5409        assert_eq!(pool.allocate(), None);
5410        assert_eq!(pool.allocate_many(1), None);
5411    }
5412
5413    #[mz_ore::test]
5414    fn test_allocate_single() {
5415        let mut pool = IdPool::empty();
5416        pool.refill(10, 13);
5417        assert_eq!(pool.remaining(), 3);
5418        assert_eq!(pool.allocate(), Some(10));
5419        assert_eq!(pool.allocate(), Some(11));
5420        assert_eq!(pool.allocate(), Some(12));
5421        assert_eq!(pool.remaining(), 0);
5422        assert_eq!(pool.allocate(), None);
5423    }
5424
5425    #[mz_ore::test]
5426    fn test_allocate_many() {
5427        let mut pool = IdPool::empty();
5428        pool.refill(100, 105);
5429        assert_eq!(pool.allocate_many(3), Some(vec![100, 101, 102]));
5430        assert_eq!(pool.remaining(), 2);
5431        // Not enough remaining for 3 more.
5432        assert_eq!(pool.allocate_many(3), None);
5433        // But 2 works.
5434        assert_eq!(pool.allocate_many(2), Some(vec![103, 104]));
5435        assert_eq!(pool.remaining(), 0);
5436    }
5437
5438    #[mz_ore::test]
5439    fn test_allocate_many_zero() {
5440        let mut pool = IdPool::empty();
5441        pool.refill(1, 5);
5442        assert_eq!(pool.allocate_many(0), Some(vec![]));
5443        assert_eq!(pool.remaining(), 4);
5444    }
5445
5446    #[mz_ore::test]
5447    fn test_refill_resets_pool() {
5448        let mut pool = IdPool::empty();
5449        pool.refill(0, 2);
5450        assert_eq!(pool.allocate(), Some(0));
5451        // Refill before exhaustion replaces the range.
5452        pool.refill(50, 52);
5453        assert_eq!(pool.allocate(), Some(50));
5454        assert_eq!(pool.allocate(), Some(51));
5455        assert_eq!(pool.allocate(), None);
5456    }
5457
5458    #[mz_ore::test]
5459    fn test_mixed_allocate_and_allocate_many() {
5460        let mut pool = IdPool::empty();
5461        pool.refill(0, 10);
5462        assert_eq!(pool.allocate(), Some(0));
5463        assert_eq!(pool.allocate_many(3), Some(vec![1, 2, 3]));
5464        assert_eq!(pool.allocate(), Some(4));
5465        assert_eq!(pool.remaining(), 5);
5466    }
5467
5468    #[mz_ore::test]
5469    #[should_panic(expected = "invalid pool range")]
5470    fn test_refill_invalid_range_panics() {
5471        let mut pool = IdPool::empty();
5472        pool.refill(10, 5);
5473    }
5474}
5475
5476#[cfg(test)]
5477mod arrangement_sizes_pruner_tests {
5478    use mz_repr::catalog_item_id::CatalogItemId;
5479    use mz_repr::{Datum, Row};
5480
5481    use super::arrangement_sizes_expired_retractions;
5482
5483    // Pack a row shaped like `mz_object_arrangement_size_history`: the pruner
5484    // only cares about column 3 (`collection_timestamp`), but we stuff the
5485    // other three columns with realistic values so shape changes would fail.
5486    fn history_row(ts_ms: i64) -> Row {
5487        let dt = mz_ore::now::to_datetime(ts_ms.try_into().expect("non-negative"));
5488        Row::pack_slice(&[
5489            Datum::String("r1"),
5490            Datum::String("u1"),
5491            Datum::Int64(123),
5492            Datum::TimestampTz(dt.try_into().expect("fits in TimestampTz")),
5493        ])
5494    }
5495
5496    fn item_id() -> CatalogItemId {
5497        // Any CatalogItemId will do; tests don't dispatch on it.
5498        CatalogItemId::User(42)
5499    }
5500
5501    #[mz_ore::test]
5502    fn empty_input_produces_no_retractions() {
5503        let out = arrangement_sizes_expired_retractions(Vec::new(), 1_000, item_id());
5504        assert!(out.is_empty());
5505    }
5506
5507    #[mz_ore::test]
5508    fn retracts_only_rows_strictly_before_cutoff() {
5509        // Mixes both sides of the filter and includes a row at exactly
5510        // the cutoff timestamp to pin down the strict-less-than boundary.
5511        let rows = vec![
5512            (history_row(100), 1),
5513            (history_row(500), 1),
5514            (history_row(1_000), 1), // at cutoff: kept (strict <)
5515            (history_row(5_000), 1),
5516        ];
5517        let out = arrangement_sizes_expired_retractions(rows, 1_000, item_id());
5518        assert_eq!(out.len(), 2);
5519    }
5520
5521    #[mz_ore::test]
5522    #[should_panic(expected = "consolidated contents should not contain retractions")]
5523    fn retraction_in_input_panics() {
5524        let rows = vec![(history_row(100), -1)];
5525        let _ = arrangement_sizes_expired_retractions(rows, 1_000, item_id());
5526    }
5527}