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