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