mz_catalog/durable/

persist.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

#[cfg(test)]
mod tests;

use std::cmp::max;
use std::collections::{BTreeMap, VecDeque};
use std::fmt::Debug;
use std::str::FromStr;
use std::sync::{Arc, LazyLock};
use std::time::{Duration, Instant};

use async_trait::async_trait;
use differential_dataflow::lattice::Lattice;
use futures::{FutureExt, StreamExt};
use itertools::Itertools;
use mz_audit_log::VersionedEvent;
use mz_ore::cast::CastFrom;
use mz_ore::metrics::MetricsFutureExt;
use mz_ore::now::EpochMillis;
use mz_ore::{
    soft_assert_eq_no_log, soft_assert_eq_or_log, soft_assert_ne_or_log, soft_assert_no_log,
    soft_assert_or_log, soft_panic_or_log,
};
use mz_persist_client::cfg::USE_CRITICAL_SINCE_CATALOG;
use mz_persist_client::cli::admin::{CATALOG_FORCE_COMPACTION_FUEL, CATALOG_FORCE_COMPACTION_WAIT};
use mz_persist_client::critical::{CriticalReaderId, Opaque, SinceHandle};
use mz_persist_client::error::UpperMismatch;
use mz_persist_client::read::{Listen, ListenEvent, ReadHandle};
use mz_persist_client::write::WriteHandle;
use mz_persist_client::{Diagnostics, PersistClient, ShardId};
use mz_persist_types::codec_impls::UnitSchema;
use mz_proto::TryFromProtoError;
use mz_repr::Diff;
use mz_storage_client::controller::PersistEpoch;
use mz_storage_types::StorageDiff;
use mz_storage_types::sources::SourceData;
use sha2::Digest;
use timely::progress::{Antichain, Timestamp as TimelyTimestamp};
use tracing::{debug, info, warn};
use uuid::Uuid;

use crate::durable::debug::{Collection, CollectionType, DebugCatalogState, Trace};
use crate::durable::error::FenceError;
use crate::durable::initialize::{
    ENABLE_0DT_DEPLOYMENT_PANIC_AFTER_TIMEOUT, SYSTEM_CONFIG_SYNCED_KEY, USER_VERSION_KEY,
    WITH_0DT_DEPLOYMENT_DDL_CHECK_INTERVAL, WITH_0DT_DEPLOYMENT_MAX_WAIT,
};
use crate::durable::metrics::Metrics;
use crate::durable::objects::serialization::RustType;
use crate::durable::objects::state_update::{
    IntoStateUpdateKindJson, StateUpdate, StateUpdateKind, StateUpdateKindJson,
    TryIntoStateUpdateKind,
};
use crate::durable::objects::{AuditLogKey, FenceToken, Snapshot};
use crate::durable::transaction::TransactionBatch;
use crate::durable::upgrade::upgrade;
use crate::durable::{
    BootstrapArgs, CATALOG_CONTENT_VERSION_KEY, CatalogError, DryRunTransaction,
    DurableCatalogError, DurableCatalogState, Epoch, OpenableDurableCatalogState,
    ReadOnlyDurableCatalogState, Transaction, initialize, persist_desc,
};
use crate::memory;

/// New-type used to represent timestamps in persist.
pub(crate) type Timestamp = mz_repr::Timestamp;

/// The minimum value of an epoch.
const MIN_EPOCH: Epoch = Epoch::new(1).expect("1 is non-zero");

/// Human readable catalog shard name.
const CATALOG_SHARD_NAME: &str = "catalog";

/// [`CriticalReaderId`] for the catalog shard's own since hold, separate from
/// [`PersistClient::CONTROLLER_CRITICAL_SINCE`].
static CATALOG_CRITICAL_SINCE: LazyLock<CriticalReaderId> = LazyLock::new(|| {
    "c55555555-6666-7777-8888-999999999999"
        .parse()
        .expect("valid CriticalReaderId")
});

/// Seed used to generate the persist shard ID for the catalog.
const CATALOG_SEED: usize = 1;
/// Legacy seed used to generate the persist shard ID for the upgrade shard. DO NOT REUSE.
const _UPGRADE_SEED: usize = 2;
/// Legacy seed used to generate the persist shard ID for builtin table migrations. DO NOT REUSE.
pub const _BUILTIN_MIGRATION_SEED: usize = 3;
/// Legacy seed used to generate the persist shard ID for the expression cache. DO NOT REUSE.
pub const _EXPRESSION_CACHE_SEED: usize = 4;

/// Durable catalog mode that dictates the effect of mutable operations.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub(crate) enum Mode {
    /// Mutable operations are prohibited.
    Readonly,
    /// Mutable operations have an effect in-memory, but aren't persisted durably.
    Savepoint,
    /// Mutable operations have an effect in-memory and durably.
    Writable,
}

/// Enum representing the fenced state of the catalog.
#[derive(Debug)]
pub(crate) enum FenceableToken {
    /// The catalog is still initializing and learning about previously written fence tokens. This
    /// state can be fenced if it encounters a larger deploy generation.
    Initializing {
        /// The largest fence token durably written to the catalog, if any.
        durable_token: Option<FenceToken>,
        /// This process's deploy generation.
        current_deploy_generation: Option<u64>,
    },
    /// The current token has not been fenced.
    Unfenced { current_token: FenceToken },
    /// The current token has been fenced.
    Fenced {
        current_token: FenceToken,
        fence_token: FenceToken,
    },
}

impl FenceableToken {
    /// Returns a new token.
    fn new(current_deploy_generation: Option<u64>) -> Self {
        Self::Initializing {
            durable_token: None,
            current_deploy_generation,
        }
    }

    /// Returns the current token if it is not fenced, otherwise returns an error.
    fn validate(&self) -> Result<Option<FenceToken>, FenceError> {
        match self {
            FenceableToken::Initializing { durable_token, .. } => Ok(durable_token.clone()),
            FenceableToken::Unfenced { current_token, .. } => Ok(Some(current_token.clone())),
            FenceableToken::Fenced {
                current_token,
                fence_token,
            } => {
                assert!(
                    fence_token > current_token,
                    "must be fenced by higher token; current={current_token:?}, fence={fence_token:?}"
                );
                if fence_token.deploy_generation > current_token.deploy_generation {
                    Err(FenceError::DeployGeneration {
                        current_generation: current_token.deploy_generation,
                        fence_generation: fence_token.deploy_generation,
                    })
                } else {
                    assert!(
                        fence_token.epoch > current_token.epoch,
                        "must be fenced by higher token; current={current_token:?}, fence={fence_token:?}"
                    );
                    Err(FenceError::Epoch {
                        current_epoch: current_token.epoch,
                        fence_epoch: fence_token.epoch,
                    })
                }
            }
        }
    }

    /// Returns the current token.
    fn token(&self) -> Option<FenceToken> {
        match self {
            FenceableToken::Initializing { durable_token, .. } => durable_token.clone(),
            FenceableToken::Unfenced { current_token, .. } => Some(current_token.clone()),
            FenceableToken::Fenced { current_token, .. } => Some(current_token.clone()),
        }
    }

    /// Returns `Err` if `token` fences out `self`, `Ok` otherwise.
    fn maybe_fence(&mut self, token: FenceToken) -> Result<(), FenceError> {
        match self {
            FenceableToken::Initializing {
                durable_token,
                current_deploy_generation,
                ..
            } => {
                match durable_token {
                    Some(durable_token) => {
                        *durable_token = max(durable_token.clone(), token.clone());
                    }
                    None => {
                        *durable_token = Some(token.clone());
                    }
                }
                if let Some(current_deploy_generation) = current_deploy_generation {
                    if *current_deploy_generation < token.deploy_generation {
                        *self = FenceableToken::Fenced {
                            current_token: FenceToken {
                                deploy_generation: *current_deploy_generation,
                                epoch: token.epoch,
                            },
                            fence_token: token,
                        };
                        self.validate()?;
                    }
                }
            }
            FenceableToken::Unfenced { current_token } => {
                if *current_token < token {
                    *self = FenceableToken::Fenced {
                        current_token: current_token.clone(),
                        fence_token: token,
                    };
                    self.validate()?;
                }
            }
            FenceableToken::Fenced { .. } => {
                self.validate()?;
            }
        }

        Ok(())
    }

    /// Returns a [`FenceableToken::Unfenced`] token and the updates to the catalog required to
    /// transition to the `Unfenced` state if `self` is [`FenceableToken::Initializing`], otherwise
    /// returns `None`.
    fn generate_unfenced_token(
        &self,
        mode: Mode,
    ) -> Result<Option<(Vec<(StateUpdateKind, Diff)>, FenceableToken)>, DurableCatalogError> {
        let (durable_token, current_deploy_generation) = match self {
            FenceableToken::Initializing {
                durable_token,
                current_deploy_generation,
            } => (durable_token.clone(), current_deploy_generation.clone()),
            FenceableToken::Unfenced { .. } | FenceableToken::Fenced { .. } => return Ok(None),
        };

        let mut fence_updates = Vec::with_capacity(2);

        if let Some(durable_token) = &durable_token {
            fence_updates.push((
                StateUpdateKind::FenceToken(durable_token.clone()),
                Diff::MINUS_ONE,
            ));
        }

        let current_deploy_generation = current_deploy_generation
            .or_else(|| durable_token.as_ref().map(|token| token.deploy_generation))
            // We cannot initialize a catalog without a deploy generation.
            .ok_or(DurableCatalogError::Uninitialized)?;
        let mut current_epoch = durable_token
            .map(|token| token.epoch)
            .unwrap_or(MIN_EPOCH)
            .get();
        // Only writable catalogs attempt to increment the epoch.
        if matches!(mode, Mode::Writable) {
            current_epoch = current_epoch + 1;
        }
        let current_epoch = Epoch::new(current_epoch).expect("known to be non-zero");
        let current_token = FenceToken {
            deploy_generation: current_deploy_generation,
            epoch: current_epoch,
        };

        fence_updates.push((
            StateUpdateKind::FenceToken(current_token.clone()),
            Diff::ONE,
        ));

        let current_fenceable_token = FenceableToken::Unfenced { current_token };

        Ok(Some((fence_updates, current_fenceable_token)))
    }
}

/// An error that can occur while executing [`PersistHandle::compare_and_append`].
#[derive(Debug, thiserror::Error)]
pub(crate) enum CompareAndAppendError {
    #[error(transparent)]
    Fence(#[from] FenceError),
    /// Catalog encountered an upper mismatch when trying to write to the catalog: another
    /// writer moved the upper between our snapshot of it and the write. Handled by the conflict
    /// classification in the commit and advance paths (rebase over empty progress, surface
    /// content conflicts as out-of-sync).
    #[error(
        "expected catalog upper {expected_upper:?} did not match actual catalog upper {actual_upper:?}"
    )]
    UpperMismatch {
        expected_upper: Timestamp,
        actual_upper: Timestamp,
    },
}

impl CompareAndAppendError {
    pub(crate) fn unwrap_fence_error(self) -> FenceError {
        match self {
            CompareAndAppendError::Fence(e) => e,
            e @ CompareAndAppendError::UpperMismatch { .. } => {
                panic!("unexpected upper mismatch: {e:?}")
            }
        }
    }
}

impl From<UpperMismatch<Timestamp>> for CompareAndAppendError {
    fn from(upper_mismatch: UpperMismatch<Timestamp>) -> Self {
        Self::UpperMismatch {
            expected_upper: antichain_to_timestamp(upper_mismatch.expected),
            actual_upper: antichain_to_timestamp(upper_mismatch.current),
        }
    }
}

pub(crate) trait ApplyUpdate<T: IntoStateUpdateKindJson> {
    /// Process and apply `update`.
    ///
    /// Returns `Some` if `update` should be cached in memory and `None` otherwise.
    fn apply_update(
        &mut self,
        update: StateUpdate<T>,
        current_fence_token: &mut FenceableToken,
        metrics: &Arc<Metrics>,
    ) -> Result<Option<StateUpdate<T>>, FenceError>;
}

/// A handle for interacting with the persist catalog shard.
///
/// The catalog shard is used in multiple different contexts, for example pre-open and post-open,
/// but for all contexts the majority of the durable catalog's behavior is identical. This struct
/// implements those behaviors that are identical while allowing the user to specify the different
/// behaviors via generic parameters.
///
/// The behavior of the durable catalog can be different along one of two axes. The first is the
/// format of each individual update, i.e. raw binary, the current protobuf version, previous
/// protobuf versions, etc. The second axis is what to do with each individual update, for example
/// before opening we cache all config updates but don't cache them after opening. These behaviors
/// are customizable via the `T: TryIntoStateUpdateKind` and `U: ApplyUpdate<T>` generic parameters
/// respectively.
#[derive(Debug)]
pub(crate) struct PersistHandle<T: TryIntoStateUpdateKind, U: ApplyUpdate<T>> {
    /// The [`Mode`] that this catalog was opened in.
    pub(crate) mode: Mode,
    /// Since handle to control compaction.
    since_handle: SinceHandle<SourceData, (), Timestamp, StorageDiff>,
    /// Write handle to persist.
    write_handle: WriteHandle<SourceData, (), Timestamp, StorageDiff>,
    /// Listener to catalog changes.
    listen: Listen<SourceData, (), Timestamp, StorageDiff>,
    /// Handle for connecting to persist.
    persist_client: PersistClient,
    /// Catalog shard ID.
    shard_id: ShardId,
    /// Cache of the most recent catalog snapshot.
    ///
    /// We use a tuple instead of [`StateUpdate`] to make consolidation easier.
    pub(crate) snapshot: Vec<(T, Timestamp, Diff)>,
    /// Applies custom processing, filtering, and fencing for each individual update.
    update_applier: U,
    /// The current upper of the persist shard.
    pub(crate) upper: Timestamp,
    /// The fence token of the catalog, if one exists.
    fenceable_token: FenceableToken,
    /// The semantic version of the current binary.
    catalog_content_version: semver::Version,
    /// Flag to indicate if bootstrap is complete.
    bootstrap_complete: bool,
    /// Metrics for the persist catalog.
    metrics: Arc<Metrics>,
    /// Snapshot size at the last amortized consolidation, used by
    /// [`Self::maybe_consolidate`] to decide when to consolidate. Initialized
    /// lazily on the first call.
    size_at_last_consolidation: Option<usize>,
    /// Counts raw updates applied to this handle.
    ///
    /// This distinguishes empty upper progress from content. Memory updates are insufficient
    /// because kinds such as ID allocators do not produce them.
    updates_applied: u64,
}

impl<T: TryIntoStateUpdateKind, U: ApplyUpdate<T>> PersistHandle<T, U> {
    /// Fetch the current upper of the catalog state.
    #[mz_ore::instrument]
    async fn current_upper(&mut self) -> Timestamp {
        match self.mode {
            Mode::Writable | Mode::Readonly => {
                let upper = self.write_handle.fetch_recent_upper().await;
                antichain_to_timestamp(upper.clone())
            }
            Mode::Savepoint => self.upper,
        }
    }

    /// Appends `updates` iff the current global upper of the catalog is `self.upper`.
    ///
    /// Returns the next upper used to commit the transaction.
    #[mz_ore::instrument]
    pub(crate) async fn compare_and_append<S: IntoStateUpdateKindJson>(
        &mut self,
        updates: Vec<(S, Diff)>,
        commit_ts: Timestamp,
    ) -> Result<Timestamp, CompareAndAppendError> {
        let updates = updates.into_iter().map(|(kind, diff)| {
            let kind: StateUpdateKindJson = kind.into();
            (
                (Into::<SourceData>::into(kind), ()),
                commit_ts,
                diff.into_inner(),
            )
        });
        let next_upper = commit_ts.step_forward();
        // Upper mismatches are classified by the commit and advance callers.
        self.compare_and_append_inner(updates, next_upper).await?;

        self.sync(next_upper).await?;
        Ok(next_upper)
    }

    /// Compare-and-append `updates` to the catalog shard, advancing the upper to `next_upper`.
    ///
    /// On success, updating `self.upper` is left to the caller. The caller can thus decide whether
    /// or not it needs to sync the catalog.
    ///
    /// # Panics
    ///
    /// Panics if not in `Writable` mode.
    /// Panics if `next_upper` is not greater than `self.upper`.
    async fn compare_and_append_inner(
        &mut self,
        updates: impl IntoIterator<Item = ((SourceData, ()), Timestamp, StorageDiff)>,
        next_upper: Timestamp,
    ) -> Result<(), CompareAndAppendError> {
        assert_eq!(self.mode, Mode::Writable);
        assert!(
            next_upper > self.upper,
            "next_upper ({next_upper}) not greater than current upper ({})",
            self.upper,
        );

        let res = self
            .write_handle
            .compare_and_append(
                updates,
                Antichain::from_elem(self.upper),
                Antichain::from_elem(next_upper),
            )
            .await
            .expect("invalid usage");

        if let Err(e @ UpperMismatch { .. }) = res {
            // Most likely we were fenced out.
            // Sync to the current upper to detect that.
            self.sync_to_current_upper().await?;
            return Err(e.into());
        }

        // Lag the shard's upper by 1 to keep it readable.
        let downgrade_to = Antichain::from_elem(next_upper.saturating_sub(1));

        // The since handle gives us the ability to fence out other downgraders using an opaque token.
        // (See the method documentation for details.)
        // That's not needed here, so we use the since handle's opaque token to avoid any comparison
        // failures.
        let opaque = self.since_handle.opaque().clone();
        let downgrade = self
            .since_handle
            .maybe_compare_and_downgrade_since(&opaque, (&opaque, &downgrade_to))
            .await;
        if let Some(Err(e)) = downgrade {
            soft_panic_or_log!("found opaque value {e:?}, but expected {opaque:?}");
        }

        Ok(())
    }

    /// Accepts an upper mismatch caused only by empty progress.
    ///
    /// `updates_applied_before` must be captured before the compare-and-append, which synchronizes
    /// this handle on mismatch.
    fn classify_upper_mismatch(
        &self,
        updates_applied_before: u64,
        actual_upper: Timestamp,
    ) -> Result<(), DurableCatalogError> {
        if self.updates_applied != updates_applied_before {
            Err(DurableCatalogError::CatalogOutOfSync {
                update_count: usize::cast_from(self.updates_applied - updates_applied_before),
                upper: actual_upper,
            })
        } else {
            Ok(())
        }
    }

    /// Generates an iterator of [`StateUpdate`] that contain all unconsolidated updates to the
    /// catalog state up to, and including, `as_of`.
    #[mz_ore::instrument]
    async fn snapshot_unconsolidated(&mut self) -> Vec<StateUpdate<StateUpdateKind>> {
        let current_upper = self.current_upper().await;

        let mut snapshot = Vec::new();
        let mut read_handle = self.read_handle().await;
        let as_of = as_of(&read_handle, current_upper);
        let mut stream = Box::pin(
            // We use `snapshot_and_stream` because it guarantees unconsolidated output.
            read_handle
                .snapshot_and_stream(Antichain::from_elem(as_of))
                .await
                .expect("we have advanced the restart_as_of by the since"),
        );
        while let Some(update) = stream.next().await {
            snapshot.push(update)
        }
        read_handle.expire().await;
        snapshot
            .into_iter()
            .map(Into::<StateUpdate<StateUpdateKindJson>>::into)
            .map(|state_update| state_update.try_into().expect("kind decoding error"))
            .collect()
    }

    /// Listen and apply all updates that are currently in persist.
    ///
    /// Returns an error if this instance has been fenced out.
    #[mz_ore::instrument]
    pub(crate) async fn sync_to_current_upper(&mut self) -> Result<(), FenceError> {
        let upper = self.current_upper().await;
        self.sync(upper).await
    }

    /// Listen and apply all updates up to `target_upper`.
    ///
    /// Returns an error if this instance has been fenced out.
    #[mz_ore::instrument(level = "debug")]
    pub(crate) async fn sync(&mut self, target_upper: Timestamp) -> Result<(), FenceError> {
        self.metrics.syncs.inc();
        let histogram = self.metrics.sync_latency_seconds.clone();
        self.sync_inner(target_upper)
            .wall_time()
            .observe(histogram)
            .await
    }

    #[mz_ore::instrument(level = "debug")]
    async fn sync_inner(&mut self, target_upper: Timestamp) -> Result<(), FenceError> {
        self.fenceable_token.validate()?;

        // Savepoint catalogs do not yet know how to update themselves in response to concurrent
        // writes from writer catalogs.
        if self.mode == Mode::Savepoint {
            self.upper = max(self.upper, target_upper);
            return Ok(());
        }

        let mut updates: BTreeMap<_, Vec<_>> = BTreeMap::new();
        let updates_applied_before = self.updates_applied;

        // Reset the amortized consolidation tracker so it picks up the
        // current snapshot size as its baseline.
        self.size_at_last_consolidation = None;

        while self.upper < target_upper {
            let listen_events = self.listen.fetch_next().await;
            for listen_event in listen_events {
                match listen_event {
                    ListenEvent::Progress(upper) => {
                        debug!("synced up to {upper:?}");
                        self.upper = antichain_to_timestamp(upper);
                        // Attempt to apply updates in batches of a single timestamp. If another
                        // catalog wrote a fence token at one timestamp and then updates in a new
                        // format at a later timestamp, then we want to apply the fence token
                        // before attempting to deserialize the new updates.
                        while let Some((ts, updates)) = updates.pop_first() {
                            assert!(ts < self.upper, "expected {} < {}", ts, self.upper);
                            let updates = updates.into_iter().map(
                                |update: StateUpdate<StateUpdateKindJson>| {
                                    let kind =
                                        T::try_from(update.kind).expect("kind decoding error");
                                    StateUpdate {
                                        kind,
                                        ts: update.ts,
                                        diff: update.diff,
                                    }
                                },
                            );
                            self.apply_updates(updates)?;
                            self.maybe_consolidate();
                        }
                    }
                    ListenEvent::Updates(batch_updates) => {
                        for update in batch_updates {
                            let update: StateUpdate<StateUpdateKindJson> = update.into();
                            updates.entry(update.ts).or_default().push(update);
                        }
                    }
                }
            }
        }
        assert_eq!(updates, BTreeMap::new(), "all updates should be applied");
        // Only consolidate when there are actual updates.
        if self.updates_applied != updates_applied_before {
            self.consolidate();
        }
        Ok(())
    }

    /// Apply a batch of updates and then consolidate the snapshot. This is the
    /// typical entry point for callers that apply updates in a single batch.
    ///
    /// For hot loops that apply updates across many timestamps (e.g., `sync_inner`),
    /// use `apply_updates` directly and call `consolidate()` periodically (e.g.,
    /// on snapshot doubling) to bound memory while staying amortized O(N log N).
    pub(crate) fn apply_updates_and_consolidate(
        &mut self,
        updates: impl IntoIterator<Item = StateUpdate<T>>,
    ) -> Result<(), FenceError> {
        self.apply_updates(updates)?;
        self.consolidate();
        Ok(())
    }

    /// Apply a batch of updates to the catalog state without consolidating.
    ///
    /// Does NOT consolidate the snapshot afterward. If you are calling this once,
    /// prefer `apply_updates_and_consolidate`. This method exists for loops that
    /// call it many times — consolidating per call would be O(K * N log N) instead
    /// of O(N log N). Callers should consolidate periodically (e.g., on snapshot
    /// doubling) to bound memory.
    #[mz_ore::instrument(level = "debug")]
    fn apply_updates(
        &mut self,
        updates: impl IntoIterator<Item = StateUpdate<T>>,
    ) -> Result<(), FenceError> {
        let mut updates: Vec<_> = updates
            .into_iter()
            .map(|StateUpdate { kind, ts, diff }| (kind, ts, diff))
            .collect();

        // This helps guarantee that for a single key, there is at most a single retraction and a
        // single insertion per timestamp. Otherwise, we would need to match the retractions and
        // insertions up by value and manually figure out what the end value should be.
        differential_dataflow::consolidation::consolidate_updates(&mut updates);

        // Updates must be applied in timestamp order. Within a timestamp retractions must be
        // applied before insertions, or we might end up retracting the wrong value.
        updates.sort_by(|(_, ts1, diff1), (_, ts2, diff2)| ts1.cmp(ts2).then(diff1.cmp(diff2)));

        let mut errors = Vec::new();

        for (kind, ts, diff) in updates {
            if diff != Diff::ONE && diff != Diff::MINUS_ONE {
                panic!("invalid update in consolidated trace: ({kind:?}, {ts:?}, {diff:?})");
            }
            self.updates_applied += 1;

            match self.update_applier.apply_update(
                StateUpdate { kind, ts, diff },
                &mut self.fenceable_token,
                &self.metrics,
            ) {
                Ok(Some(StateUpdate { kind, ts, diff })) => self.snapshot.push((kind, ts, diff)),
                Ok(None) => {}
                // Instead of returning immediately, we accumulate all the errors and return the one
                // with the most information.
                Err(err) => errors.push(err),
            }
        }

        // Track the high-water mark of the unconsolidated snapshot size.
        let len = i64::try_from(self.snapshot.len()).unwrap_or(i64::MAX);
        if len > self.metrics.snapshot_max_entries.get() {
            self.metrics.snapshot_max_entries.set(len);
        }

        errors.sort();
        if let Some(err) = errors.into_iter().next() {
            return Err(err);
        }

        Ok(())
    }

    /// Consolidate the snapshot if it has at least doubled in size since the
    /// last consolidation. This amortizes the O(N log N) consolidation cost
    /// over many small updates, keeping the total work O(N log N) rather than
    /// O(K * N log N) for K timestamps.
    fn maybe_consolidate(&mut self) {
        let threshold = *self
            .size_at_last_consolidation
            // Use a minimum of 8 to avoid consolidating on every update when
            // the snapshot is small or empty (since 0 * 2 = 0).
            .get_or_insert_with(|| max(self.snapshot.len(), 8));
        if self.snapshot.len() >= threshold * 2 {
            self.consolidate();
            self.size_at_last_consolidation = Some(self.snapshot.len());
        }
    }

    #[mz_ore::instrument]
    pub(crate) fn consolidate(&mut self) {
        self.metrics.snapshot_consolidations.inc();
        soft_assert_no_log!(
            self.snapshot
                .windows(2)
                .all(|updates| updates[0].1 <= updates[1].1),
            "snapshot should be sorted by timestamp, {:#?}",
            self.snapshot
        );

        let new_ts = self
            .snapshot
            .last()
            .map(|(_, ts, _)| *ts)
            .unwrap_or_else(Timestamp::minimum);
        for (_, ts, _) in &mut self.snapshot {
            *ts = new_ts;
        }
        differential_dataflow::consolidation::consolidate_updates(&mut self.snapshot);
    }

    /// Execute and return the results of `f` on the current catalog trace.
    ///
    /// Will return an error if the catalog has been fenced out.
    async fn with_trace<R>(
        &mut self,
        f: impl FnOnce(&Vec<(T, Timestamp, Diff)>) -> Result<R, CatalogError>,
    ) -> Result<R, CatalogError> {
        self.sync_to_current_upper().await?;
        f(&self.snapshot)
    }

    /// Open a read handle to the catalog.
    async fn read_handle(&self) -> ReadHandle<SourceData, (), Timestamp, StorageDiff> {
        self.persist_client
            .open_leased_reader(
                self.shard_id,
                Arc::new(persist_desc()),
                Arc::new(UnitSchema::default()),
                Diagnostics {
                    shard_name: CATALOG_SHARD_NAME.to_string(),
                    handle_purpose: "openable durable catalog state temporary reader".to_string(),
                },
                USE_CRITICAL_SINCE_CATALOG.get(self.persist_client.dyncfgs()),
            )
            .await
            .expect("invalid usage")
    }

    /// Politely releases all external resources that can only be released in an async context.
    async fn expire(self: Box<Self>) {
        self.write_handle.expire().await;
        self.listen.expire().await;
    }
}

impl<U: ApplyUpdate<StateUpdateKind>> PersistHandle<StateUpdateKind, U> {
    /// Execute and return the results of `f` on the current catalog snapshot.
    ///
    /// Will return an error if the catalog has been fenced out.
    async fn with_snapshot<T>(
        &mut self,
        f: impl FnOnce(Snapshot) -> Result<T, CatalogError>,
    ) -> Result<T, CatalogError> {
        fn apply<K, V>(map: &mut BTreeMap<K, V>, key: &K, value: &V, diff: Diff)
        where
            K: Ord + Clone,
            V: Ord + Clone + Debug,
        {
            let key = key.clone();
            let value = value.clone();
            if diff == Diff::ONE {
                let prev = map.insert(key, value);
                assert_eq!(
                    prev, None,
                    "values must be explicitly retracted before inserting a new value"
                );
            } else if diff == Diff::MINUS_ONE {
                let prev = map.remove(&key);
                assert_eq!(
                    prev,
                    Some(value),
                    "retraction does not match existing value"
                );
            }
        }

        self.with_trace(|trace| {
            let mut snapshot = Snapshot::empty();
            for (kind, ts, diff) in trace {
                let diff = *diff;
                if diff != Diff::ONE && diff != Diff::MINUS_ONE {
                    panic!("invalid update in consolidated trace: ({kind:?}, {ts:?}, {diff:?})");
                }

                match kind {
                    StateUpdateKind::AuditLog(_key, ()) => {
                        // Ignore for snapshots.
                    }
                    StateUpdateKind::Cluster(key, value) => {
                        apply(&mut snapshot.clusters, key, value, diff);
                    }
                    StateUpdateKind::ClusterReplica(key, value) => {
                        apply(&mut snapshot.cluster_replicas, key, value, diff);
                    }
                    StateUpdateKind::ClusterReplicaSize(key, value) => {
                        apply(&mut snapshot.cluster_replica_sizes, key, value, diff);
                    }
                    StateUpdateKind::Comment(key, value) => {
                        apply(&mut snapshot.comments, key, value, diff);
                    }
                    StateUpdateKind::IndexColumns(key, value) => {
                        apply(&mut snapshot.index_columns, key, value, diff);
                    }
                    StateUpdateKind::ItemColumns(key, value) => {
                        apply(&mut snapshot.item_columns, key, value, diff);
                    }
                    StateUpdateKind::Config(key, value) => {
                        apply(&mut snapshot.configs, key, value, diff);
                    }
                    StateUpdateKind::Database(key, value) => {
                        apply(&mut snapshot.databases, key, value, diff);
                    }
                    StateUpdateKind::DefaultPrivilege(key, value) => {
                        apply(&mut snapshot.default_privileges, key, value, diff);
                    }
                    StateUpdateKind::FenceToken(_token) => {
                        // Ignore for snapshots.
                    }
                    StateUpdateKind::IdAllocator(key, value) => {
                        apply(&mut snapshot.id_allocator, key, value, diff);
                    }
                    StateUpdateKind::IntrospectionSourceIndex(key, value) => {
                        apply(&mut snapshot.introspection_sources, key, value, diff);
                    }
                    StateUpdateKind::Item(key, value) => {
                        apply(&mut snapshot.items, key, value, diff);
                    }
                    StateUpdateKind::NetworkPolicy(key, value) => {
                        apply(&mut snapshot.network_policies, key, value, diff);
                    }
                    StateUpdateKind::Role(key, value) => {
                        apply(&mut snapshot.roles, key, value, diff);
                    }
                    StateUpdateKind::Schema(key, value) => {
                        apply(&mut snapshot.schemas, key, value, diff);
                    }
                    StateUpdateKind::Setting(key, value) => {
                        apply(&mut snapshot.settings, key, value, diff);
                    }
                    StateUpdateKind::SourceReferences(key, value) => {
                        apply(&mut snapshot.source_references, key, value, diff);
                    }
                    StateUpdateKind::SystemConfiguration(key, value) => {
                        apply(&mut snapshot.system_configurations, key, value, diff);
                    }
                    StateUpdateKind::ClusterSystemConfiguration(key, value) => {
                        apply(
                            &mut snapshot.cluster_system_configurations,
                            key,
                            value,
                            diff,
                        );
                    }
                    StateUpdateKind::ReplicaSystemConfiguration(key, value) => {
                        apply(
                            &mut snapshot.replica_system_configurations,
                            key,
                            value,
                            diff,
                        );
                    }
                    StateUpdateKind::SystemObjectMapping(key, value) => {
                        apply(&mut snapshot.system_object_mappings, key, value, diff);
                    }
                    StateUpdateKind::SystemPrivilege(key, value) => {
                        apply(&mut snapshot.system_privileges, key, value, diff);
                    }
                    StateUpdateKind::StorageCollectionMetadata(key, value) => {
                        apply(&mut snapshot.storage_collection_metadata, key, value, diff);
                    }
                    StateUpdateKind::UnfinalizedShard(key, ()) => {
                        apply(&mut snapshot.unfinalized_shards, key, &(), diff);
                    }
                    StateUpdateKind::TxnWalShard((), value) => {
                        apply(&mut snapshot.txn_wal_shard, &(), value, diff);
                    }
                    StateUpdateKind::RoleAuth(key, value) => {
                        apply(&mut snapshot.role_auth, key, value, diff);
                    }
                }
            }
            f(snapshot)
        })
        .await
    }

    /// Generates an iterator of [`StateUpdate`] that contain all updates to the catalog
    /// state.
    ///
    /// The output is fetched directly from persist instead of the in-memory cache.
    ///
    /// The output is consolidated and sorted by timestamp in ascending order.
    #[mz_ore::instrument(level = "debug")]
    async fn persist_snapshot(&self) -> impl Iterator<Item = StateUpdate> + DoubleEndedIterator {
        let mut read_handle = self.read_handle().await;
        let as_of = as_of(&read_handle, self.upper);
        let snapshot = snapshot_binary(&mut read_handle, as_of, &self.metrics)
            .await
            .map(|update| update.try_into().expect("kind decoding error"));
        read_handle.expire().await;
        snapshot
    }
}

/// Applies updates for an unopened catalog.
#[derive(Debug)]
pub(crate) struct UnopenedCatalogStateInner {
    /// A cache of the config collection of the catalog.
    configs: BTreeMap<String, u64>,
    /// A cache of the settings collection of the catalog.
    settings: BTreeMap<String, String>,
}

impl UnopenedCatalogStateInner {
    fn new() -> UnopenedCatalogStateInner {
        UnopenedCatalogStateInner {
            configs: BTreeMap::new(),
            settings: BTreeMap::new(),
        }
    }
}

impl ApplyUpdate<StateUpdateKindJson> for UnopenedCatalogStateInner {
    fn apply_update(
        &mut self,
        update: StateUpdate<StateUpdateKindJson>,
        current_fence_token: &mut FenceableToken,
        _metrics: &Arc<Metrics>,
    ) -> Result<Option<StateUpdate<StateUpdateKindJson>>, FenceError> {
        if !update.kind.is_audit_log() && update.kind.is_always_deserializable() {
            let kind = TryInto::try_into(&update.kind).expect("kind is known to be deserializable");
            match (kind, update.diff) {
                (StateUpdateKind::Config(key, value), Diff::ONE) => {
                    let prev = self.configs.insert(key.key, value.value);
                    assert_eq!(
                        prev, None,
                        "values must be explicitly retracted before inserting a new value"
                    );
                }
                (StateUpdateKind::Config(key, value), Diff::MINUS_ONE) => {
                    let prev = self.configs.remove(&key.key);
                    assert_eq!(
                        prev,
                        Some(value.value),
                        "retraction does not match existing value"
                    );
                }
                (StateUpdateKind::Setting(key, value), Diff::ONE) => {
                    let prev = self.settings.insert(key.name, value.value);
                    assert_eq!(
                        prev, None,
                        "values must be explicitly retracted before inserting a new value"
                    );
                }
                (StateUpdateKind::Setting(key, value), Diff::MINUS_ONE) => {
                    let prev = self.settings.remove(&key.name);
                    assert_eq!(
                        prev,
                        Some(value.value),
                        "retraction does not match existing value"
                    );
                }
                (StateUpdateKind::FenceToken(fence_token), Diff::ONE) => {
                    current_fence_token.maybe_fence(fence_token)?;
                }
                _ => {}
            }
        }

        Ok(Some(update))
    }
}

/// A Handle to an unopened catalog stored in persist. The unopened catalog can serve `Config` data,
/// `Setting` data, or the current epoch. All other catalog data may be un-migrated and should not
/// be read until the catalog has been opened. The [`UnopenedPersistCatalogState`] is responsible
/// for opening the catalog, see [`OpenableDurableCatalogState::open`] for more details.
///
/// Production users should call [`Self::expire`] before dropping an [`UnopenedPersistCatalogState`]
/// so that it can expire its leases. If/when rust gets AsyncDrop, this will be done automatically.
pub(crate) type UnopenedPersistCatalogState =
    PersistHandle<StateUpdateKindJson, UnopenedCatalogStateInner>;

impl UnopenedPersistCatalogState {
    /// Create a new [`UnopenedPersistCatalogState`] to the catalog state associated with
    /// `organization_id`.
    ///
    /// All usages of the persist catalog must go through this function. That includes the
    /// catalog-debug tool, the adapter's catalog, etc.
    #[mz_ore::instrument]
    pub(crate) async fn new(
        persist_client: PersistClient,
        organization_id: Uuid,
        version: semver::Version,
        deploy_generation: Option<u64>,
        metrics: Arc<Metrics>,
    ) -> Result<UnopenedPersistCatalogState, DurableCatalogError> {
        let catalog_shard_id = shard_id(organization_id, CATALOG_SEED);
        debug!(?catalog_shard_id, "new persist backed catalog state");

        // Check the catalog shard version to ensure that we are compatible with the persist
        // data format. This lets us return an error gracefully, rather than panicking later in
        // persist.
        let version_in_catalog_shard =
            fetch_catalog_shard_version(&persist_client, catalog_shard_id).await;
        if let Some(version_in_catalog_shard) = version_in_catalog_shard {
            if !mz_persist_client::cfg::code_can_write_data(&version, &version_in_catalog_shard) {
                return Err(DurableCatalogError::IncompatiblePersistVersion {
                    found_version: version_in_catalog_shard,
                    catalog_version: version,
                });
            }
        }

        let open_handles_start = Instant::now();
        info!("startup: envd serve: catalog init: open handles beginning");
        let since_handle = persist_client
            .open_critical_since(
                catalog_shard_id,
                CATALOG_CRITICAL_SINCE.clone(),
                Opaque::encode(&i64::MIN),
                Diagnostics {
                    shard_name: CATALOG_SHARD_NAME.to_string(),
                    handle_purpose: "durable catalog state critical since".to_string(),
                },
            )
            .await
            .expect("invalid usage");

        let (mut write_handle, mut read_handle) = persist_client
            .open(
                catalog_shard_id,
                Arc::new(persist_desc()),
                Arc::new(UnitSchema::default()),
                Diagnostics {
                    shard_name: CATALOG_SHARD_NAME.to_string(),
                    handle_purpose: "durable catalog state handles".to_string(),
                },
                USE_CRITICAL_SINCE_CATALOG.get(persist_client.dyncfgs()),
            )
            .await
            .expect("invalid usage");
        info!(
            "startup: envd serve: catalog init: open handles complete in {:?}",
            open_handles_start.elapsed()
        );

        // Commit an empty write at the minimum timestamp so the catalog is always readable.
        let upper = {
            const EMPTY_UPDATES: &[((SourceData, ()), Timestamp, StorageDiff)] = &[];
            let upper = Antichain::from_elem(Timestamp::minimum());
            let next_upper = Timestamp::minimum().step_forward();
            match write_handle
                .compare_and_append(EMPTY_UPDATES, upper, Antichain::from_elem(next_upper))
                .await
                .expect("invalid usage")
            {
                Ok(()) => next_upper,
                Err(mismatch) => antichain_to_timestamp(mismatch.current),
            }
        };

        let snapshot_start = Instant::now();
        info!("startup: envd serve: catalog init: snapshot beginning");
        let as_of = as_of(&read_handle, upper);
        let snapshot: Vec<_> = snapshot_binary(&mut read_handle, as_of, &metrics)
            .await
            .map(|StateUpdate { kind, ts, diff }| (kind, ts, diff))
            .collect();
        let listen = read_handle
            .listen(Antichain::from_elem(as_of))
            .await
            .expect("invalid usage");
        info!(
            "startup: envd serve: catalog init: snapshot complete in {:?}",
            snapshot_start.elapsed()
        );

        let mut handle = UnopenedPersistCatalogState {
            // Unopened catalogs are always writeable until they're opened in an explicit mode.
            mode: Mode::Writable,
            since_handle,
            write_handle,
            listen,
            persist_client,
            shard_id: catalog_shard_id,
            // Initialize empty in-memory state.
            snapshot: Vec::new(),
            update_applier: UnopenedCatalogStateInner::new(),
            upper,
            fenceable_token: FenceableToken::new(deploy_generation),
            catalog_content_version: version,
            bootstrap_complete: false,
            metrics,
            size_at_last_consolidation: None,
            updates_applied: 0,
        };
        // If the snapshot is not consolidated, and we see multiple epoch values while applying the
        // updates, then we might accidentally fence ourselves out.
        soft_assert_no_log!(
            snapshot.iter().all(|(_, _, diff)| *diff == Diff::ONE),
            "snapshot should be consolidated: {snapshot:#?}"
        );

        let apply_start = Instant::now();
        info!("startup: envd serve: catalog init: apply updates beginning");
        let updates = snapshot
            .into_iter()
            .map(|(kind, ts, diff)| StateUpdate { kind, ts, diff });
        handle.apply_updates_and_consolidate(updates)?;
        info!(
            "startup: envd serve: catalog init: apply updates complete in {:?}",
            apply_start.elapsed()
        );

        // Validate that the binary version of the current process is not less than any binary
        // version that has written to the catalog.
        // This condition is only checked once, right here. If a new process comes along with a
        // higher version, it must fence this process out with one of the existing fencing
        // mechanisms.
        if let Some(found_version) = handle.get_catalog_content_version().await? {
            // Use cmp_precedence() to ignore build metadata per SemVer 2.0.0 spec
            if handle
                .catalog_content_version
                .cmp_precedence(&found_version)
                == std::cmp::Ordering::Less
            {
                return Err(DurableCatalogError::IncompatiblePersistVersion {
                    found_version,
                    catalog_version: handle.catalog_content_version,
                });
            }
        }

        Ok(handle)
    }

    #[mz_ore::instrument]
    async fn open_inner(
        mut self,
        mode: Mode,
        initial_ts: Timestamp,
        bootstrap_args: &BootstrapArgs,
    ) -> Result<Box<dyn DurableCatalogState>, CatalogError> {
        // It would be nice to use `initial_ts` here, but it comes from the system clock, not the
        // timestamp oracle.
        let mut commit_ts = self.upper;
        self.mode = mode;

        // Validate the current deploy generation.
        match (&self.mode, &self.fenceable_token) {
            (_, FenceableToken::Unfenced { .. } | FenceableToken::Fenced { .. }) => {
                return Err(DurableCatalogError::Internal(
                    "catalog should not have fenced before opening".to_string(),
                )
                .into());
            }
            (
                Mode::Writable | Mode::Savepoint,
                FenceableToken::Initializing {
                    current_deploy_generation: None,
                    ..
                },
            ) => {
                return Err(DurableCatalogError::Internal(format!(
                    "cannot open in mode '{:?}' without a deploy generation",
                    self.mode,
                ))
                .into());
            }
            _ => {}
        }

        let read_only = matches!(self.mode, Mode::Readonly);

        // Fence out previous catalogs.
        loop {
            self.sync_to_current_upper().await?;
            commit_ts = max(commit_ts, self.upper);
            let (fence_updates, current_fenceable_token) = self
                .fenceable_token
                .generate_unfenced_token(self.mode)?
                .ok_or_else(|| {
                    DurableCatalogError::Internal(
                        "catalog should not have fenced before opening".to_string(),
                    )
                })?;
            debug!(
                ?self.upper,
                ?self.fenceable_token,
                ?current_fenceable_token,
                "fencing previous catalogs"
            );
            if matches!(self.mode, Mode::Writable) {
                match self
                    .compare_and_append(fence_updates.clone(), commit_ts)
                    .await
                {
                    Ok(upper) => {
                        commit_ts = upper;
                    }
                    Err(CompareAndAppendError::Fence(e)) => return Err(e.into()),
                    Err(e @ CompareAndAppendError::UpperMismatch { .. }) => {
                        warn!("catalog write failed due to upper mismatch, retrying: {e:?}");
                        continue;
                    }
                }
            }
            self.fenceable_token = current_fenceable_token;
            break;
        }

        if matches!(self.mode, Mode::Writable) {
            // One-time migration: The catalog previously used `CONTROLLER_CRITICAL_SINCE` for its
            // since handle. Now it uses its own `CATALOG_CRITICAL_SINCE`, to free
            // `CONTROLLER_CRITICAL_SINCE` up for the storage controller. The catalog and
            // controller handles differ in the `Opaque` codec, so we need a migration.
            //
            // TODO: Remove this once we don't support upgrading from v26 anymore.
            let mut controller_handle = self
                .persist_client
                .open_critical_since::<SourceData, (), Timestamp, StorageDiff>(
                    self.shard_id,
                    PersistClient::CONTROLLER_CRITICAL_SINCE,
                    Opaque::encode(&i64::MIN),
                    Diagnostics {
                        shard_name: CATALOG_SHARD_NAME.to_string(),
                        handle_purpose: "durable catalog state critical since (migration)"
                            .to_string(),
                    },
                )
                .await
                .expect("invalid usage");

            let since = controller_handle.since().clone();
            let res = controller_handle
                .compare_and_downgrade_since(
                    &Opaque::encode(&i64::MIN),
                    (&Opaque::encode(&PersistEpoch::default()), &since),
                )
                .await;
            match res {
                Ok(_) => info!("migrated Opaque of catalog since handle"),
                Err(_) => { /* critical since was already migrated */ }
            }
        }

        let is_initialized = self.is_initialized_inner();
        if !matches!(self.mode, Mode::Writable) && !is_initialized {
            return Err(CatalogError::Durable(DurableCatalogError::NotWritable(
                format!(
                    "catalog tables do not exist; will not create in {:?} mode",
                    self.mode
                ),
            )));
        }
        soft_assert_ne_or_log!(self.upper, Timestamp::minimum());

        // Audit log entries are served from `mz_internal.mz_catalog_raw` via
        // the `mz_audit_events` materialized view, so they do not need to live
        // in the in-memory catalog snapshot. Drop them here and only keep the
        // count for metrics.
        let (audit_logs, snapshot): (Vec<_>, Vec<_>) = self
            .snapshot
            .into_iter()
            .partition(|(update, _, _)| update.is_audit_log());
        self.snapshot = snapshot;
        let audit_log_count = audit_logs.iter().map(|(_, _, diff)| diff).sum::<Diff>();
        drop(audit_logs);

        // Perform data migrations.
        if is_initialized && !read_only {
            commit_ts = upgrade(&mut self, commit_ts).await?;
        }

        debug!(
            ?is_initialized,
            ?self.upper,
            "initializing catalog state"
        );
        let mut catalog = PersistCatalogState {
            mode: self.mode,
            since_handle: self.since_handle,
            write_handle: self.write_handle,
            listen: self.listen,
            persist_client: self.persist_client,
            shard_id: self.shard_id,
            upper: self.upper,
            fenceable_token: self.fenceable_token,
            // Initialize empty in-memory state.
            snapshot: Vec::new(),
            update_applier: CatalogStateInner::new(),
            catalog_content_version: self.catalog_content_version,
            bootstrap_complete: false,
            metrics: self.metrics,
            size_at_last_consolidation: None,
            updates_applied: 0,
        };
        catalog.metrics.collection_entries.reset();
        // Normally, `collection_entries` is updated in `apply_updates`. The audit log updates skip
        // over that function so we manually update it here.
        catalog
            .metrics
            .collection_entries
            .with_label_values(&[&CollectionType::AuditLog.to_string()])
            .add(audit_log_count.into_inner());
        let updates = self.snapshot.into_iter().map(|(kind, ts, diff)| {
            let kind = TryIntoStateUpdateKind::try_into(kind).expect("kind decoding error");
            StateUpdate { kind, ts, diff }
        });
        catalog.apply_updates_and_consolidate(updates)?;

        let catalog_content_version = catalog.catalog_content_version.to_string();
        let txn = if is_initialized {
            let mut txn = catalog.transaction_unchecked().await?;

            // Ad-hoc migration: Initialize the `migration_version` expected by adapter to be
            // present in existing catalogs.
            //
            // Note: Need to exclude read-only catalog mode here, because in that mode all
            // transactions are expected to be no-ops.
            // TODO: remove this once we only support upgrades from version >= 0.164
            if txn.get_setting("migration_version".into()).is_none() && mode != Mode::Readonly {
                let old_version = txn.get_catalog_content_version();
                txn.set_setting("migration_version".into(), old_version.map(Into::into))?;
            }

            // Opening the catalog with write intent fences out every previous
            // catalog owner, so all sessions served by previous owners are
            // dead. Reclaim the temporary items they owned here, before
            // anything else reads the catalog.
            //
            // NOTE: This only reclaims on the fence, which covers a
            // single-writer world where every crash is followed by some
            // process's writable open. Once several serving envds run
            // concurrently, a peer crash triggers no fence here, so
            // reclaiming its sessions' items needs the durable
            // envd-heartbeat mechanism described in the durable temporary
            // objects design doc.
            if mode != Mode::Readonly {
                txn.remove_ephemeral_items();
            }

            txn.set_catalog_content_version(catalog_content_version)?;
            txn
        } else {
            soft_assert_eq_no_log!(
                catalog
                    .snapshot
                    .iter()
                    .filter(|(kind, _, _)| !matches!(kind, StateUpdateKind::FenceToken(_)))
                    .count(),
                0,
                "trace should not contain any updates for an uninitialized catalog: {:#?}",
                catalog.snapshot
            );

            let mut txn = catalog.transaction_unchecked().await?;
            initialize::initialize(
                &mut txn,
                bootstrap_args,
                initial_ts.into(),
                catalog_content_version,
            )
            .await?;
            txn
        };

        if read_only {
            let (txn_batch, _) = txn.into_parts()?;
            // The upper here doesn't matter because we are only applying the updates in memory.
            let updates = StateUpdate::from_txn_batch_ts(txn_batch, catalog.upper);
            catalog.apply_updates_and_consolidate(updates)?;
        } else {
            txn.commit_internal(commit_ts).await?;
        }

        if matches!(catalog.mode, Mode::Writable) {
            let write_handle = catalog
                .persist_client
                .open_writer::<SourceData, (), Timestamp, i64>(
                    catalog.write_handle.shard_id(),
                    Arc::new(persist_desc()),
                    Arc::new(UnitSchema::default()),
                    Diagnostics {
                        shard_name: CATALOG_SHARD_NAME.to_string(),
                        handle_purpose: "compact catalog".to_string(),
                    },
                )
                .await
                .expect("invalid usage");
            let fuel = CATALOG_FORCE_COMPACTION_FUEL.handle(catalog.persist_client.dyncfgs());
            let wait = CATALOG_FORCE_COMPACTION_WAIT.handle(catalog.persist_client.dyncfgs());
            // We're going to gradually turn this on via dyncfgs. Run it in a task so that it
            // doesn't block startup.
            let _task = mz_ore::task::spawn(|| "catalog::force_shard_compaction", async move {
                let () =
                    mz_persist_client::cli::admin::dangerous_force_compaction_and_break_pushdown(
                        &write_handle,
                        || fuel.get(),
                        || wait.get(),
                    )
                    .await;
            });
        }

        Ok(Box::new(catalog))
    }

    /// Reports if the catalog state has been initialized.
    ///
    /// NOTE: This is the answer as of the last call to [`PersistHandle::sync`] or [`PersistHandle::sync_to_current_upper`],
    /// not necessarily what is currently in persist.
    #[mz_ore::instrument]
    fn is_initialized_inner(&self) -> bool {
        !self.update_applier.configs.is_empty()
    }

    /// Get the current value of config `key`.
    ///
    /// Some configs need to be read before the catalog is opened for bootstrapping.
    #[mz_ore::instrument]
    async fn get_current_config(&mut self, key: &str) -> Result<Option<u64>, DurableCatalogError> {
        self.sync_to_current_upper().await?;
        Ok(self.update_applier.configs.get(key).cloned())
    }

    /// Get the user version of this instance.
    ///
    /// The user version is used to determine if a migration is needed.
    #[mz_ore::instrument]
    pub(crate) async fn get_user_version(&mut self) -> Result<Option<u64>, DurableCatalogError> {
        self.get_current_config(USER_VERSION_KEY).await
    }

    /// Get the current value of setting `name`.
    ///
    /// Some settings need to be read before the catalog is opened for bootstrapping.
    #[mz_ore::instrument]
    async fn get_current_setting(
        &mut self,
        name: &str,
    ) -> Result<Option<String>, DurableCatalogError> {
        self.sync_to_current_upper().await?;
        Ok(self.update_applier.settings.get(name).cloned())
    }

    /// Get the catalog content version.
    ///
    /// The catalog content version is the semantic version of the most recent binary that wrote to
    /// the catalog.
    #[mz_ore::instrument]
    async fn get_catalog_content_version(
        &mut self,
    ) -> Result<Option<semver::Version>, DurableCatalogError> {
        let version = self
            .get_current_setting(CATALOG_CONTENT_VERSION_KEY)
            .await?;
        let version = version.map(|version| version.parse().expect("invalid version persisted"));
        Ok(version)
    }
}

#[async_trait]
impl OpenableDurableCatalogState for UnopenedPersistCatalogState {
    #[mz_ore::instrument]
    async fn open_savepoint(
        mut self: Box<Self>,
        initial_ts: Timestamp,
        bootstrap_args: &BootstrapArgs,
    ) -> Result<Box<dyn DurableCatalogState>, CatalogError> {
        self.open_inner(Mode::Savepoint, initial_ts, bootstrap_args)
            .boxed()
            .await
    }

    #[mz_ore::instrument]
    async fn open_read_only(
        mut self: Box<Self>,
        bootstrap_args: &BootstrapArgs,
    ) -> Result<Box<dyn DurableCatalogState>, CatalogError> {
        self.open_inner(Mode::Readonly, EpochMillis::MIN.into(), bootstrap_args)
            .boxed()
            .await
    }

    #[mz_ore::instrument]
    async fn open(
        mut self: Box<Self>,
        initial_ts: Timestamp,
        bootstrap_args: &BootstrapArgs,
    ) -> Result<Box<dyn DurableCatalogState>, CatalogError> {
        self.open_inner(Mode::Writable, initial_ts, bootstrap_args)
            .boxed()
            .await
    }

    #[mz_ore::instrument(level = "debug")]
    async fn open_debug(mut self: Box<Self>) -> Result<DebugCatalogState, CatalogError> {
        Ok(DebugCatalogState(*self))
    }

    #[mz_ore::instrument]
    async fn is_initialized(&mut self) -> Result<bool, CatalogError> {
        self.sync_to_current_upper().await?;
        Ok(self.is_initialized_inner())
    }

    #[mz_ore::instrument]
    async fn epoch(&mut self) -> Result<Epoch, CatalogError> {
        self.sync_to_current_upper().await?;
        self.fenceable_token
            .validate()?
            .map(|token| token.epoch)
            .ok_or(CatalogError::Durable(DurableCatalogError::Uninitialized))
    }

    #[mz_ore::instrument]
    async fn get_deployment_generation(&mut self) -> Result<u64, CatalogError> {
        self.sync_to_current_upper().await?;
        self.fenceable_token
            .token()
            .map(|token| token.deploy_generation)
            .ok_or(CatalogError::Durable(DurableCatalogError::Uninitialized))
    }

    #[mz_ore::instrument(level = "debug")]
    async fn get_0dt_deployment_max_wait(&mut self) -> Result<Option<Duration>, CatalogError> {
        let value = self
            .get_current_config(WITH_0DT_DEPLOYMENT_MAX_WAIT)
            .await?;
        match value {
            None => Ok(None),
            Some(millis) => Ok(Some(Duration::from_millis(millis))),
        }
    }

    #[mz_ore::instrument(level = "debug")]
    async fn get_0dt_deployment_ddl_check_interval(
        &mut self,
    ) -> Result<Option<Duration>, CatalogError> {
        let value = self
            .get_current_config(WITH_0DT_DEPLOYMENT_DDL_CHECK_INTERVAL)
            .await?;
        match value {
            None => Ok(None),
            Some(millis) => Ok(Some(Duration::from_millis(millis))),
        }
    }

    #[mz_ore::instrument(level = "debug")]
    async fn get_enable_0dt_deployment_panic_after_timeout(
        &mut self,
    ) -> Result<Option<bool>, CatalogError> {
        let value = self
            .get_current_config(ENABLE_0DT_DEPLOYMENT_PANIC_AFTER_TIMEOUT)
            .await?;
        match value {
            None => Ok(None),
            Some(0) => Ok(Some(false)),
            Some(1) => Ok(Some(true)),
            Some(v) => Err(
                DurableCatalogError::from(TryFromProtoError::UnknownEnumVariant(format!(
                    "{v} is not a valid boolean value"
                )))
                .into(),
            ),
        }
    }

    #[mz_ore::instrument]
    async fn has_system_config_synced_once(&mut self) -> Result<bool, DurableCatalogError> {
        self.get_current_config(SYSTEM_CONFIG_SYNCED_KEY)
            .await
            .map(|value| value.map(|value| value > 0).unwrap_or(false))
    }

    #[mz_ore::instrument]
    async fn trace_unconsolidated(&mut self) -> Result<Trace, CatalogError> {
        self.sync_to_current_upper().await?;
        if self.is_initialized_inner() {
            let snapshot = self.snapshot_unconsolidated().await;
            Ok(Trace::from_snapshot(snapshot))
        } else {
            Err(CatalogError::Durable(DurableCatalogError::Uninitialized))
        }
    }

    #[mz_ore::instrument]
    async fn trace_consolidated(&mut self) -> Result<Trace, CatalogError> {
        self.sync_to_current_upper().await?;
        if self.is_initialized_inner() {
            let snapshot = self.current_snapshot().await?;
            Ok(Trace::from_snapshot(snapshot))
        } else {
            Err(CatalogError::Durable(DurableCatalogError::Uninitialized))
        }
    }

    #[mz_ore::instrument(level = "debug")]
    async fn expire(self: Box<Self>) {
        self.expire().await
    }
}

/// Applies updates for an opened catalog.
#[derive(Debug)]
struct CatalogStateInner {
    /// A trace of all catalog updates that can be consumed by some higher layer.
    updates: VecDeque<memory::objects::StateUpdate>,
}

impl CatalogStateInner {
    fn new() -> CatalogStateInner {
        CatalogStateInner {
            updates: VecDeque::new(),
        }
    }
}

impl ApplyUpdate<StateUpdateKind> for CatalogStateInner {
    fn apply_update(
        &mut self,
        update: StateUpdate<StateUpdateKind>,
        current_fence_token: &mut FenceableToken,
        metrics: &Arc<Metrics>,
    ) -> Result<Option<StateUpdate<StateUpdateKind>>, FenceError> {
        if let Some(collection_type) = update.kind.collection_type() {
            metrics
                .collection_entries
                .with_label_values(&[&collection_type.to_string()])
                .add(update.diff.into_inner());
        }

        {
            let update: Option<memory::objects::StateUpdate> = (&update)
                .try_into()
                .expect("invalid persisted update: {update:#?}");
            if let Some(update) = update {
                self.updates.push_back(update);
            }
        }

        match (update.kind, update.diff) {
            (StateUpdateKind::AuditLog(_, ()), _) => Ok(None),
            // Nothing to due for fence token retractions but wait for the next insertion.
            (StateUpdateKind::FenceToken(_), Diff::MINUS_ONE) => Ok(None),
            (StateUpdateKind::FenceToken(token), Diff::ONE) => {
                current_fence_token.maybe_fence(token)?;
                Ok(None)
            }
            (kind, diff) => Ok(Some(StateUpdate {
                kind,
                ts: update.ts,
                diff,
            })),
        }
    }
}

/// A durable store of the catalog state using Persist as an implementation. The durable store can
/// serve any catalog data and transactionally modify catalog data.
///
/// Production users should call [`Self::expire`] before dropping a [`PersistCatalogState`]
/// so that it can expire its leases. If/when rust gets AsyncDrop, this will be done automatically.
type PersistCatalogState = PersistHandle<StateUpdateKind, CatalogStateInner>;

impl PersistHandle<StateUpdateKind, CatalogStateInner> {
    /// Creates a transaction without validating pending catalog updates.
    async fn transaction_unchecked(&mut self) -> Result<Transaction<'_>, CatalogError> {
        self.metrics.transactions_started.inc();
        let snapshot = self.snapshot().await?;
        let commit_ts = self.upper;
        Transaction::new(self, snapshot, commit_ts)
    }
}

#[async_trait]
impl ReadOnlyDurableCatalogState for PersistCatalogState {
    fn epoch(&self) -> Epoch {
        self.fenceable_token
            .token()
            .expect("opened catalog state must have an epoch")
            .epoch
    }

    fn metrics(&self) -> &Metrics {
        &self.metrics
    }

    #[mz_ore::instrument(level = "debug")]
    async fn expire(self: Box<Self>) {
        self.expire().await
    }

    fn is_bootstrap_complete(&self) -> bool {
        self.bootstrap_complete
    }

    async fn get_audit_logs(&mut self) -> Result<Vec<VersionedEvent>, CatalogError> {
        self.sync_to_current_upper().await?;
        let audit_logs: Vec<_> = self
            .persist_snapshot()
            .await
            .filter_map(
                |StateUpdate {
                     kind,
                     ts: _,
                     diff: _,
                 }| match kind {
                    StateUpdateKind::AuditLog(key, ()) => Some(key),
                    _ => None,
                },
            )
            .collect();
        let mut audit_logs: Vec<_> = audit_logs
            .into_iter()
            .map(RustType::from_proto)
            .map_ok(|key: AuditLogKey| key.event)
            .collect::<Result<_, _>>()?;
        audit_logs.sort_by(|a, b| a.sortable_id().cmp(&b.sortable_id()));
        Ok(audit_logs)
    }

    #[mz_ore::instrument(level = "debug")]
    async fn get_next_id(&mut self, id_type: &str) -> Result<u64, CatalogError> {
        self.with_trace(|trace| {
            Ok(trace
                .into_iter()
                .rev()
                .filter_map(|(kind, _, _)| match kind {
                    StateUpdateKind::IdAllocator(key, value) if key.name == id_type => {
                        Some(value.next_id)
                    }
                    _ => None,
                })
                .next()
                .expect("must exist"))
        })
        .await
    }

    #[mz_ore::instrument(level = "debug")]
    async fn get_deployment_generation(&mut self) -> Result<u64, CatalogError> {
        self.sync_to_current_upper().await?;
        Ok(self
            .fenceable_token
            .token()
            .expect("opened catalogs must have a token")
            .deploy_generation)
    }

    #[mz_ore::instrument(level = "debug")]
    async fn snapshot(&mut self) -> Result<Snapshot, CatalogError> {
        self.with_snapshot(Ok).await
    }

    #[mz_ore::instrument(level = "debug")]
    async fn sync_to_current_updates(
        &mut self,
    ) -> Result<Vec<memory::objects::StateUpdate>, CatalogError> {
        let upper = self.current_upper().await;
        self.sync_updates(upper).await
    }

    #[mz_ore::instrument(level = "debug")]
    async fn sync_updates(
        &mut self,
        target_upper: mz_repr::Timestamp,
    ) -> Result<Vec<memory::objects::StateUpdate>, CatalogError> {
        self.sync(target_upper).await?;
        let mut updates = Vec::new();
        while let Some(update) = self.update_applier.updates.front() {
            if update.ts >= target_upper {
                break;
            }

            let update = self
                .update_applier
                .updates
                .pop_front()
                .expect("peeked above");
            updates.push(update);
        }
        Ok(updates)
    }

    #[mz_ore::instrument(level = "debug")]
    async fn ensure_not_out_of_sync(
        &mut self,
        target_upper: Timestamp,
    ) -> Result<(), CatalogError> {
        self.sync(target_upper).await?;
        let update_count = self
            .update_applier
            .updates
            .iter()
            .take_while(|update| update.ts < target_upper)
            .count();
        if update_count == 0 {
            Ok(())
        } else {
            Err(DurableCatalogError::CatalogOutOfSync {
                update_count,
                upper: target_upper,
            }
            .into())
        }
    }

    async fn current_upper(&mut self) -> Timestamp {
        self.current_upper().await
    }
}

#[async_trait]
#[allow(mismatched_lifetime_syntaxes)]
impl DurableCatalogState for PersistCatalogState {
    fn is_read_only(&self) -> bool {
        matches!(self.mode, Mode::Readonly)
    }

    fn is_savepoint(&self) -> bool {
        matches!(self.mode, Mode::Savepoint)
    }

    async fn mark_bootstrap_complete(&mut self) {
        self.bootstrap_complete = true;
        if matches!(self.mode, Mode::Writable) {
            self.since_handle
                .upgrade_version()
                .await
                .expect("invalid usage")
        }
    }

    #[mz_ore::instrument(level = "debug")]
    async fn transaction(&mut self) -> Result<Transaction, CatalogError> {
        let mut txn = self.transaction_unchecked().await?;
        txn.ensure_not_out_of_sync().await?;
        Ok(txn)
    }

    fn transaction_from_snapshot(
        &mut self,
        snapshot: Snapshot,
    ) -> Result<DryRunTransaction, CatalogError> {
        let commit_ts = self.upper;
        Transaction::new(self, snapshot, commit_ts).map(DryRunTransaction::new)
    }

    #[mz_ore::instrument(level = "debug")]
    async fn allocate_id(
        &mut self,
        id_type: &str,
        amount: u64,
        commit_ts: Timestamp,
    ) -> Result<Vec<u64>, CatalogError> {
        let start = Instant::now();
        if amount == 0 {
            return Ok(Vec::new());
        }
        let mut txn = self.transaction_unchecked().await?;
        let ids = txn.get_and_increment_id_by(id_type.to_string(), amount)?;
        txn.commit_internal(commit_ts).await?;
        self.metrics
            .allocate_id_seconds
            .observe(start.elapsed().as_secs_f64());
        Ok(ids)
    }

    #[mz_ore::instrument(level = "debug")]
    async fn commit_transaction(
        &mut self,
        txn_batch: TransactionBatch,
        commit_ts: Timestamp,
    ) -> Result<Timestamp, CatalogError> {
        async fn commit_transaction_inner(
            catalog: &mut PersistCatalogState,
            txn_batch: TransactionBatch,
            commit_ts: Timestamp,
        ) -> Result<Timestamp, CatalogError> {
            // If the transaction is empty then we don't error, even in read-only mode.
            // This is mostly for legacy reasons (i.e. with enough elbow grease this
            // behavior can be changed without breaking any fundamental assumptions).
            if catalog.mode == Mode::Readonly {
                let updates: Vec<_> = StateUpdate::from_txn_batch(txn_batch).collect();
                if !updates.is_empty() {
                    let collection_types: Vec<_> = updates
                        .iter()
                        .filter_map(|u| u.0.collection_type())
                        .collect();
                    return Err(DurableCatalogError::NotWritable(format!(
                        "cannot commit a transaction in a read-only catalog: \
                         {} updates across collections: {collection_types:?}",
                        updates.len(),
                    ))
                    .into());
                }
                return Ok(catalog.upper);
            }

            // The handle is mutex-protected from transaction creation through commit, so only a
            // local programming error can change its cached upper here.
            assert_eq!(
                catalog.upper, txn_batch.upper,
                "the handle was mutated mid-transaction"
            );

            let updates: Vec<_> = StateUpdate::from_txn_batch(txn_batch).collect();
            debug!("committing updates: {updates:?}");

            // Empty upper progress does not invalidate the transaction.
            let mut commit_ts = max(commit_ts, catalog.upper);

            let next_upper = match catalog.mode {
                Mode::Writable => loop {
                    let updates_applied_before = catalog.updates_applied;
                    match catalog.compare_and_append(updates.clone(), commit_ts).await {
                        Ok(next_upper) => break next_upper,
                        Err(CompareAndAppendError::Fence(e)) => return Err(e.into()),
                        Err(CompareAndAppendError::UpperMismatch { actual_upper, .. }) => {
                            // The mismatch synchronized the handle. Retry only if it applied no
                            // content.
                            catalog
                                .classify_upper_mismatch(updates_applied_before, actual_upper)?;
                            commit_ts = max(commit_ts, catalog.upper);
                        }
                    }
                },
                Mode::Savepoint => {
                    let updates = updates.into_iter().map(|(kind, diff)| StateUpdate {
                        kind,
                        ts: commit_ts,
                        diff,
                    });
                    catalog.apply_updates_and_consolidate(updates)?;
                    catalog.upper = commit_ts.step_forward();
                    catalog.upper
                }
                Mode::Readonly => unreachable!("handled above"),
            };

            Ok(next_upper)
        }
        self.metrics.transaction_commits.inc();
        let histogram = self.metrics.transaction_commit_latency_seconds.clone();
        commit_transaction_inner(self, txn_batch, commit_ts)
            .wall_time()
            .observe(histogram)
            .await
    }

    #[mz_ore::instrument(level = "debug")]
    async fn advance_upper(&mut self, new_upper: Timestamp) -> Result<(), CatalogError> {
        loop {
            if self.upper >= new_upper {
                // This does not consult Persist. It only revalidates a fence already cached by
                // this handle, since a sync can advance `upper` while recording the fence.
                self.fenceable_token.validate()?;
                return Ok(());
            }

            match self.mode {
                Mode::Writable => {}
                Mode::Savepoint => {
                    self.upper = new_upper;
                    return Ok(());
                }
                Mode::Readonly => {
                    return Err(DurableCatalogError::NotWritable(
                        "cannot advance upper of a read-only catalog".into(),
                    )
                    .into());
                }
            }

            let updates_applied_before = self.updates_applied;
            match self.compare_and_append_inner([], new_upper).await {
                Ok(()) => {
                    self.upper = new_upper;
                    // No sync needed since no data was written.
                    return Ok(());
                }
                Err(CompareAndAppendError::Fence(e)) => return Err(e.into()),
                Err(CompareAndAppendError::UpperMismatch { actual_upper, .. }) => {
                    // The mismatch synchronized the handle. Retry only if it applied no content.
                    self.classify_upper_mismatch(updates_applied_before, actual_upper)?;
                }
            }
        }
    }

    fn shard_id(&self) -> ShardId {
        self.shard_id
    }
}

/// Deterministically generate a shard ID for the given `organization_id` and `seed`.
pub fn shard_id(organization_id: Uuid, seed: usize) -> ShardId {
    let hash = sha2::Sha256::digest(format!("{organization_id}{seed}")).to_vec();
    soft_assert_eq_or_log!(hash.len(), 32, "SHA256 returns 32 bytes (256 bits)");
    let uuid = Uuid::from_slice(&hash[0..16]).expect("from_slice accepts exactly 16 bytes");
    ShardId::from_str(&format!("s{uuid}")).expect("known to be valid")
}

/// Generates a timestamp for reading from `read_handle` that is as fresh as possible, given
/// `upper`.
fn as_of(
    read_handle: &ReadHandle<SourceData, (), Timestamp, StorageDiff>,
    upper: Timestamp,
) -> Timestamp {
    let since = read_handle.since().clone();
    let mut as_of = upper.checked_sub(1).unwrap_or_else(|| {
        panic!("catalog persist shard should be initialize, found upper: {upper:?}")
    });
    // We only downgrade the since after writing, and we always set the since to one less than the
    // upper.
    soft_assert_or_log!(
        since.less_equal(&as_of),
        "since={since:?}, as_of={as_of:?}; since must be less than or equal to as_of"
    );
    // This should be a no-op if the assert above passes, however if it doesn't then we'd like to
    // continue with a correct timestamp instead of entering a panic loop.
    as_of.advance_by(since.borrow());
    as_of
}

/// Fetch the persist version of the catalog shard, if one exists. A version will not
/// exist if we are creating a brand-new environment.
async fn fetch_catalog_shard_version(
    persist_client: &PersistClient,
    catalog_shard_id: ShardId,
) -> Option<semver::Version> {
    let shard_state = persist_client
        .inspect_shard::<Timestamp>(&catalog_shard_id)
        .await
        .ok()?;
    let json_state = serde_json::to_value(shard_state).expect("state serialization error");
    let json_version = json_state
        .get("applier_version")
        .cloned()
        .expect("missing applier_version");
    let version = serde_json::from_value(json_version).expect("version deserialization error");
    Some(version)
}

/// Generates an iterator of [`StateUpdate`] that contain all updates to the catalog
/// state up to, and including, `as_of`.
///
/// The output is consolidated and sorted by timestamp in ascending order.
#[mz_ore::instrument(level = "debug")]
async fn snapshot_binary(
    read_handle: &mut ReadHandle<SourceData, (), Timestamp, StorageDiff>,
    as_of: Timestamp,
    metrics: &Arc<Metrics>,
) -> impl Iterator<Item = StateUpdate<StateUpdateKindJson>> + DoubleEndedIterator + use<> {
    metrics.snapshots_taken.inc();
    let histogram = metrics.snapshot_latency_seconds.clone();
    snapshot_binary_inner(read_handle, as_of)
        .wall_time()
        .observe(histogram)
        .await
}

/// Generates an iterator of [`StateUpdate`] that contain all updates to the catalog
/// state up to, and including, `as_of`.
///
/// The output is consolidated and sorted by timestamp in ascending order.
#[mz_ore::instrument(level = "debug")]
async fn snapshot_binary_inner(
    read_handle: &mut ReadHandle<SourceData, (), Timestamp, StorageDiff>,
    as_of: Timestamp,
) -> impl Iterator<Item = StateUpdate<StateUpdateKindJson>> + DoubleEndedIterator + use<> {
    let snapshot = read_handle
        .snapshot_and_fetch(Antichain::from_elem(as_of))
        .await
        .expect("we have advanced the restart_as_of by the since");
    soft_assert_no_log!(
        snapshot.iter().all(|(_, _, diff)| *diff == 1),
        "snapshot_and_fetch guarantees a consolidated result: {snapshot:#?}"
    );
    snapshot
        .into_iter()
        .map(Into::<StateUpdate<StateUpdateKindJson>>::into)
        .sorted_by(|a, b| Ord::cmp(&b.ts, &a.ts))
}

/// Convert an [`Antichain<Timestamp>`] to a [`Timestamp`].
///
/// The correctness of this function relies on [`Timestamp`] being totally ordered and never
/// finalizing the catalog shard.
pub(crate) fn antichain_to_timestamp(antichain: Antichain<Timestamp>) -> Timestamp {
    antichain
        .into_option()
        .expect("we use a totally ordered time and never finalize the shard")
}

// Debug methods used by the catalog-debug tool.

impl Trace {
    /// Generates a [`Trace`] from snapshot.
    fn from_snapshot(snapshot: impl IntoIterator<Item = StateUpdate>) -> Trace {
        let mut trace = Trace::new();
        for StateUpdate { kind, ts, diff } in snapshot {
            match kind {
                StateUpdateKind::AuditLog(k, v) => trace.audit_log.values.push(((k, v), ts, diff)),
                StateUpdateKind::Cluster(k, v) => trace.clusters.values.push(((k, v), ts, diff)),
                StateUpdateKind::ClusterReplica(k, v) => {
                    trace.cluster_replicas.values.push(((k, v), ts, diff))
                }
                StateUpdateKind::ClusterReplicaSize(k, v) => {
                    trace.cluster_replica_sizes.values.push(((k, v), ts, diff))
                }
                StateUpdateKind::Comment(k, v) => trace.comments.values.push(((k, v), ts, diff)),
                StateUpdateKind::IndexColumns(k, v) => {
                    trace.index_columns.values.push(((k, v), ts, diff))
                }
                StateUpdateKind::ItemColumns(k, v) => {
                    trace.item_columns.values.push(((k, v), ts, diff))
                }
                StateUpdateKind::Config(k, v) => trace.configs.values.push(((k, v), ts, diff)),
                StateUpdateKind::Database(k, v) => trace.databases.values.push(((k, v), ts, diff)),
                StateUpdateKind::DefaultPrivilege(k, v) => {
                    trace.default_privileges.values.push(((k, v), ts, diff))
                }
                StateUpdateKind::FenceToken(_) => {
                    // Fence token not included in trace.
                }
                StateUpdateKind::IdAllocator(k, v) => {
                    trace.id_allocator.values.push(((k, v), ts, diff))
                }
                StateUpdateKind::IntrospectionSourceIndex(k, v) => {
                    trace.introspection_sources.values.push(((k, v), ts, diff))
                }
                StateUpdateKind::Item(k, v) => trace.items.values.push(((k, v), ts, diff)),
                StateUpdateKind::NetworkPolicy(k, v) => {
                    trace.network_policies.values.push(((k, v), ts, diff))
                }
                StateUpdateKind::Role(k, v) => trace.roles.values.push(((k, v), ts, diff)),
                StateUpdateKind::Schema(k, v) => trace.schemas.values.push(((k, v), ts, diff)),
                StateUpdateKind::Setting(k, v) => trace.settings.values.push(((k, v), ts, diff)),
                StateUpdateKind::SourceReferences(k, v) => {
                    trace.source_references.values.push(((k, v), ts, diff))
                }
                StateUpdateKind::SystemConfiguration(k, v) => {
                    trace.system_configurations.values.push(((k, v), ts, diff))
                }
                StateUpdateKind::ClusterSystemConfiguration(k, v) => trace
                    .cluster_system_configurations
                    .values
                    .push(((k, v), ts, diff)),
                StateUpdateKind::ReplicaSystemConfiguration(k, v) => trace
                    .replica_system_configurations
                    .values
                    .push(((k, v), ts, diff)),
                StateUpdateKind::SystemObjectMapping(k, v) => {
                    trace.system_object_mappings.values.push(((k, v), ts, diff))
                }
                StateUpdateKind::SystemPrivilege(k, v) => {
                    trace.system_privileges.values.push(((k, v), ts, diff))
                }
                StateUpdateKind::StorageCollectionMetadata(k, v) => trace
                    .storage_collection_metadata
                    .values
                    .push(((k, v), ts, diff)),
                StateUpdateKind::UnfinalizedShard(k, ()) => {
                    trace.unfinalized_shards.values.push(((k, ()), ts, diff))
                }
                StateUpdateKind::TxnWalShard((), v) => {
                    trace.txn_wal_shard.values.push((((), v), ts, diff))
                }
                StateUpdateKind::RoleAuth(k, v) => trace.role_auth.values.push(((k, v), ts, diff)),
            }
        }
        trace
    }
}

impl UnopenedPersistCatalogState {
    /// Manually update value of `key` in collection `T` to `value`.
    #[mz_ore::instrument]
    pub(crate) async fn debug_edit<T: Collection>(
        &mut self,
        key: T::Key,
        value: T::Value,
    ) -> Result<Option<T::Value>, CatalogError>
    where
        T::Key: PartialEq + Eq + Debug + Clone,
        T::Value: Debug + Clone,
    {
        let prev_value = loop {
            let key = key.clone();
            let value = value.clone();
            let snapshot = self.current_snapshot().await?;
            let trace = Trace::from_snapshot(snapshot);
            let collection_trace = T::collection_trace(trace);
            let prev_values: Vec<_> = collection_trace
                .values
                .into_iter()
                .filter(|((k, _), _, diff)| {
                    soft_assert_eq_or_log!(*diff, Diff::ONE, "trace is consolidated");
                    &key == k
                })
                .collect();

            let prev_value = match &prev_values[..] {
                [] => None,
                [((_, v), _, _)] => Some(v.clone()),
                prev_values => panic!("multiple values found for key {key:?}: {prev_values:?}"),
            };

            let mut updates: Vec<_> = prev_values
                .into_iter()
                .map(|((k, v), _, _)| (T::update(k, v), Diff::MINUS_ONE))
                .collect();
            updates.push((T::update(key, value), Diff::ONE));
            // We must fence out all other catalogs, if we haven't already, since we are writing.
            match self.fenceable_token.generate_unfenced_token(self.mode)? {
                Some((fence_updates, current_fenceable_token)) => {
                    updates.extend(fence_updates.clone());
                    match self.compare_and_append(updates, self.upper).await {
                        Ok(_) => {
                            self.fenceable_token = current_fenceable_token;
                            break prev_value;
                        }
                        Err(CompareAndAppendError::Fence(e)) => return Err(e.into()),
                        Err(e @ CompareAndAppendError::UpperMismatch { .. }) => {
                            warn!("catalog write failed due to upper mismatch, retrying: {e:?}");
                            continue;
                        }
                    }
                }
                None => {
                    self.compare_and_append(updates, self.upper)
                        .await
                        .map_err(|e| e.unwrap_fence_error())?;
                    break prev_value;
                }
            }
        };
        Ok(prev_value)
    }

    /// Manually delete `key` from collection `T`.
    #[mz_ore::instrument]
    pub(crate) async fn debug_delete<T: Collection>(
        &mut self,
        key: T::Key,
    ) -> Result<(), CatalogError>
    where
        T::Key: PartialEq + Eq + Debug + Clone,
        T::Value: Debug,
    {
        loop {
            let key = key.clone();
            let snapshot = self.current_snapshot().await?;
            let trace = Trace::from_snapshot(snapshot);
            let collection_trace = T::collection_trace(trace);
            let mut retractions: Vec<_> = collection_trace
                .values
                .into_iter()
                .filter(|((k, _), _, diff)| {
                    soft_assert_eq_or_log!(*diff, Diff::ONE, "trace is consolidated");
                    &key == k
                })
                .map(|((k, v), _, _)| (T::update(k, v), Diff::MINUS_ONE))
                .collect();

            // We must fence out all other catalogs, if we haven't already, since we are writing.
            match self.fenceable_token.generate_unfenced_token(self.mode)? {
                Some((fence_updates, current_fenceable_token)) => {
                    retractions.extend(fence_updates.clone());
                    match self.compare_and_append(retractions, self.upper).await {
                        Ok(_) => {
                            self.fenceable_token = current_fenceable_token;
                            break;
                        }
                        Err(CompareAndAppendError::Fence(e)) => return Err(e.into()),
                        Err(e @ CompareAndAppendError::UpperMismatch { .. }) => {
                            warn!("catalog write failed due to upper mismatch, retrying: {e:?}");
                            continue;
                        }
                    }
                }
                None => {
                    self.compare_and_append(retractions, self.upper)
                        .await
                        .map_err(|e| e.unwrap_fence_error())?;
                    break;
                }
            }
        }
        Ok(())
    }

    /// Generates a [`Vec<StateUpdate>`] that contain all updates to the catalog
    /// state.
    ///
    /// The output is consolidated and sorted by timestamp in ascending order and the current upper.
    async fn current_snapshot(
        &mut self,
    ) -> Result<impl IntoIterator<Item = StateUpdate> + '_, CatalogError> {
        self.sync_to_current_upper().await?;
        self.consolidate();
        Ok(self.snapshot.iter().cloned().map(|(kind, ts, diff)| {
            let kind = TryIntoStateUpdateKind::try_into(kind).expect("kind decoding error");
            StateUpdate { kind, ts, diff }
        }))
    }
}