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