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