Skip to main content

mz_catalog/durable/
persist.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10#[cfg(test)]
11mod tests;
12
13use std::cmp::max;
14use std::collections::{BTreeMap, VecDeque};
15use std::fmt::Debug;
16use std::str::FromStr;
17use std::sync::{Arc, LazyLock};
18use std::time::{Duration, Instant};
19
20use async_trait::async_trait;
21use differential_dataflow::lattice::Lattice;
22use futures::{FutureExt, StreamExt};
23use itertools::Itertools;
24use mz_audit_log::VersionedEvent;
25use mz_ore::cast::CastFrom;
26use mz_ore::metrics::MetricsFutureExt;
27use mz_ore::now::EpochMillis;
28use mz_ore::{
29    soft_assert_eq_no_log, soft_assert_eq_or_log, soft_assert_ne_or_log, soft_assert_no_log,
30    soft_assert_or_log, soft_panic_or_log,
31};
32use mz_persist_client::cfg::USE_CRITICAL_SINCE_CATALOG;
33use mz_persist_client::cli::admin::{CATALOG_FORCE_COMPACTION_FUEL, CATALOG_FORCE_COMPACTION_WAIT};
34use mz_persist_client::critical::{CriticalReaderId, Opaque, SinceHandle};
35use mz_persist_client::error::UpperMismatch;
36use mz_persist_client::read::{Listen, ListenEvent, ReadHandle};
37use mz_persist_client::write::WriteHandle;
38use mz_persist_client::{Diagnostics, PersistClient, ShardId};
39use mz_persist_types::codec_impls::UnitSchema;
40use mz_proto::{RustType, TryFromProtoError};
41use mz_repr::Diff;
42use mz_storage_client::controller::PersistEpoch;
43use mz_storage_types::StorageDiff;
44use mz_storage_types::sources::SourceData;
45use sha2::Digest;
46use timely::progress::{Antichain, Timestamp as TimelyTimestamp};
47use tracing::{debug, info, warn};
48use uuid::Uuid;
49
50use crate::durable::debug::{Collection, CollectionType, DebugCatalogState, Trace};
51use crate::durable::error::FenceError;
52use crate::durable::initialize::{
53    ENABLE_0DT_DEPLOYMENT_PANIC_AFTER_TIMEOUT, SYSTEM_CONFIG_SYNCED_KEY, USER_VERSION_KEY,
54    WITH_0DT_DEPLOYMENT_DDL_CHECK_INTERVAL, WITH_0DT_DEPLOYMENT_MAX_WAIT,
55};
56use crate::durable::metrics::Metrics;
57use crate::durable::objects::state_update::{
58    IntoStateUpdateKindJson, StateUpdate, StateUpdateKind, StateUpdateKindJson,
59    TryIntoStateUpdateKind,
60};
61use crate::durable::objects::{AuditLogKey, FenceToken, Snapshot};
62use crate::durable::transaction::TransactionBatch;
63use crate::durable::upgrade::upgrade;
64use crate::durable::{
65    BootstrapArgs, CATALOG_CONTENT_VERSION_KEY, CatalogError, DryRunTransaction,
66    DurableCatalogError, DurableCatalogState, Epoch, OpenableDurableCatalogState,
67    ReadOnlyDurableCatalogState, Transaction, initialize, persist_desc,
68};
69use crate::memory;
70
71/// New-type used to represent timestamps in persist.
72pub(crate) type Timestamp = mz_repr::Timestamp;
73
74/// The minimum value of an epoch.
75const MIN_EPOCH: Epoch = Epoch::new(1).expect("1 is non-zero");
76
77/// Human readable catalog shard name.
78const CATALOG_SHARD_NAME: &str = "catalog";
79
80/// [`CriticalReaderId`] for the catalog shard's own since hold, separate from
81/// [`PersistClient::CONTROLLER_CRITICAL_SINCE`].
82static CATALOG_CRITICAL_SINCE: LazyLock<CriticalReaderId> = LazyLock::new(|| {
83    "c55555555-6666-7777-8888-999999999999"
84        .parse()
85        .expect("valid CriticalReaderId")
86});
87
88/// Seed used to generate the persist shard ID for the catalog.
89const CATALOG_SEED: usize = 1;
90/// Legacy seed used to generate the persist shard ID for the upgrade shard. DO NOT REUSE.
91const _UPGRADE_SEED: usize = 2;
92/// Legacy seed used to generate the persist shard ID for builtin table migrations. DO NOT REUSE.
93pub const _BUILTIN_MIGRATION_SEED: usize = 3;
94/// Legacy seed used to generate the persist shard ID for the expression cache. DO NOT REUSE.
95pub const _EXPRESSION_CACHE_SEED: usize = 4;
96
97/// Durable catalog mode that dictates the effect of mutable operations.
98#[derive(Debug, Copy, Clone, Eq, PartialEq)]
99pub(crate) enum Mode {
100    /// Mutable operations are prohibited.
101    Readonly,
102    /// Mutable operations have an effect in-memory, but aren't persisted durably.
103    Savepoint,
104    /// Mutable operations have an effect in-memory and durably.
105    Writable,
106}
107
108/// Enum representing the fenced state of the catalog.
109#[derive(Debug)]
110pub(crate) enum FenceableToken {
111    /// The catalog is still initializing and learning about previously written fence tokens. This
112    /// state can be fenced if it encounters a larger deploy generation.
113    Initializing {
114        /// The largest fence token durably written to the catalog, if any.
115        durable_token: Option<FenceToken>,
116        /// This process's deploy generation.
117        current_deploy_generation: Option<u64>,
118    },
119    /// The current token has not been fenced.
120    Unfenced { current_token: FenceToken },
121    /// The current token has been fenced.
122    Fenced {
123        current_token: FenceToken,
124        fence_token: FenceToken,
125    },
126}
127
128impl FenceableToken {
129    /// Returns a new token.
130    fn new(current_deploy_generation: Option<u64>) -> Self {
131        Self::Initializing {
132            durable_token: None,
133            current_deploy_generation,
134        }
135    }
136
137    /// Returns the current token if it is not fenced, otherwise returns an error.
138    fn validate(&self) -> Result<Option<FenceToken>, FenceError> {
139        match self {
140            FenceableToken::Initializing { durable_token, .. } => Ok(durable_token.clone()),
141            FenceableToken::Unfenced { current_token, .. } => Ok(Some(current_token.clone())),
142            FenceableToken::Fenced {
143                current_token,
144                fence_token,
145            } => {
146                assert!(
147                    fence_token > current_token,
148                    "must be fenced by higher token; current={current_token:?}, fence={fence_token:?}"
149                );
150                if fence_token.deploy_generation > current_token.deploy_generation {
151                    Err(FenceError::DeployGeneration {
152                        current_generation: current_token.deploy_generation,
153                        fence_generation: fence_token.deploy_generation,
154                    })
155                } else {
156                    assert!(
157                        fence_token.epoch > current_token.epoch,
158                        "must be fenced by higher token; current={current_token:?}, fence={fence_token:?}"
159                    );
160                    Err(FenceError::Epoch {
161                        current_epoch: current_token.epoch,
162                        fence_epoch: fence_token.epoch,
163                    })
164                }
165            }
166        }
167    }
168
169    /// Returns the current token.
170    fn token(&self) -> Option<FenceToken> {
171        match self {
172            FenceableToken::Initializing { durable_token, .. } => durable_token.clone(),
173            FenceableToken::Unfenced { current_token, .. } => Some(current_token.clone()),
174            FenceableToken::Fenced { current_token, .. } => Some(current_token.clone()),
175        }
176    }
177
178    /// Returns `Err` if `token` fences out `self`, `Ok` otherwise.
179    fn maybe_fence(&mut self, token: FenceToken) -> Result<(), FenceError> {
180        match self {
181            FenceableToken::Initializing {
182                durable_token,
183                current_deploy_generation,
184                ..
185            } => {
186                match durable_token {
187                    Some(durable_token) => {
188                        *durable_token = max(durable_token.clone(), token.clone());
189                    }
190                    None => {
191                        *durable_token = Some(token.clone());
192                    }
193                }
194                if let Some(current_deploy_generation) = current_deploy_generation {
195                    if *current_deploy_generation < token.deploy_generation {
196                        *self = FenceableToken::Fenced {
197                            current_token: FenceToken {
198                                deploy_generation: *current_deploy_generation,
199                                epoch: token.epoch,
200                            },
201                            fence_token: token,
202                        };
203                        self.validate()?;
204                    }
205                }
206            }
207            FenceableToken::Unfenced { current_token } => {
208                if *current_token < token {
209                    *self = FenceableToken::Fenced {
210                        current_token: current_token.clone(),
211                        fence_token: token,
212                    };
213                    self.validate()?;
214                }
215            }
216            FenceableToken::Fenced { .. } => {
217                self.validate()?;
218            }
219        }
220
221        Ok(())
222    }
223
224    /// Returns a [`FenceableToken::Unfenced`] token and the updates to the catalog required to
225    /// transition to the `Unfenced` state if `self` is [`FenceableToken::Initializing`], otherwise
226    /// returns `None`.
227    fn generate_unfenced_token(
228        &self,
229        mode: Mode,
230    ) -> Result<Option<(Vec<(StateUpdateKind, Diff)>, FenceableToken)>, DurableCatalogError> {
231        let (durable_token, current_deploy_generation) = match self {
232            FenceableToken::Initializing {
233                durable_token,
234                current_deploy_generation,
235            } => (durable_token.clone(), current_deploy_generation.clone()),
236            FenceableToken::Unfenced { .. } | FenceableToken::Fenced { .. } => return Ok(None),
237        };
238
239        let mut fence_updates = Vec::with_capacity(2);
240
241        if let Some(durable_token) = &durable_token {
242            fence_updates.push((
243                StateUpdateKind::FenceToken(durable_token.clone()),
244                Diff::MINUS_ONE,
245            ));
246        }
247
248        let current_deploy_generation = current_deploy_generation
249            .or_else(|| durable_token.as_ref().map(|token| token.deploy_generation))
250            // We cannot initialize a catalog without a deploy generation.
251            .ok_or(DurableCatalogError::Uninitialized)?;
252        let mut current_epoch = durable_token
253            .map(|token| token.epoch)
254            .unwrap_or(MIN_EPOCH)
255            .get();
256        // Only writable catalogs attempt to increment the epoch.
257        if matches!(mode, Mode::Writable) {
258            current_epoch = current_epoch + 1;
259        }
260        let current_epoch = Epoch::new(current_epoch).expect("known to be non-zero");
261        let current_token = FenceToken {
262            deploy_generation: current_deploy_generation,
263            epoch: current_epoch,
264        };
265
266        fence_updates.push((
267            StateUpdateKind::FenceToken(current_token.clone()),
268            Diff::ONE,
269        ));
270
271        let current_fenceable_token = FenceableToken::Unfenced { current_token };
272
273        Ok(Some((fence_updates, current_fenceable_token)))
274    }
275}
276
277/// An error that can occur while executing [`PersistHandle::compare_and_append`].
278#[derive(Debug, thiserror::Error)]
279pub(crate) enum CompareAndAppendError {
280    #[error(transparent)]
281    Fence(#[from] FenceError),
282    /// Catalog encountered an upper mismatch when trying to write to the catalog: another
283    /// writer moved the upper between our snapshot of it and the write. Handled by the conflict
284    /// classification in the commit and advance paths (rebase over empty progress, surface
285    /// content conflicts as out-of-sync).
286    #[error(
287        "expected catalog upper {expected_upper:?} did not match actual catalog upper {actual_upper:?}"
288    )]
289    UpperMismatch {
290        expected_upper: Timestamp,
291        actual_upper: Timestamp,
292    },
293}
294
295impl CompareAndAppendError {
296    pub(crate) fn unwrap_fence_error(self) -> FenceError {
297        match self {
298            CompareAndAppendError::Fence(e) => e,
299            e @ CompareAndAppendError::UpperMismatch { .. } => {
300                panic!("unexpected upper mismatch: {e:?}")
301            }
302        }
303    }
304}
305
306impl From<UpperMismatch<Timestamp>> for CompareAndAppendError {
307    fn from(upper_mismatch: UpperMismatch<Timestamp>) -> Self {
308        Self::UpperMismatch {
309            expected_upper: antichain_to_timestamp(upper_mismatch.expected),
310            actual_upper: antichain_to_timestamp(upper_mismatch.current),
311        }
312    }
313}
314
315pub(crate) trait ApplyUpdate<T: IntoStateUpdateKindJson> {
316    /// Process and apply `update`.
317    ///
318    /// Returns `Some` if `update` should be cached in memory and `None` otherwise.
319    fn apply_update(
320        &mut self,
321        update: StateUpdate<T>,
322        current_fence_token: &mut FenceableToken,
323        metrics: &Arc<Metrics>,
324    ) -> Result<Option<StateUpdate<T>>, FenceError>;
325}
326
327/// A handle for interacting with the persist catalog shard.
328///
329/// The catalog shard is used in multiple different contexts, for example pre-open and post-open,
330/// but for all contexts the majority of the durable catalog's behavior is identical. This struct
331/// implements those behaviors that are identical while allowing the user to specify the different
332/// behaviors via generic parameters.
333///
334/// The behavior of the durable catalog can be different along one of two axes. The first is the
335/// format of each individual update, i.e. raw binary, the current protobuf version, previous
336/// protobuf versions, etc. The second axis is what to do with each individual update, for example
337/// before opening we cache all config updates but don't cache them after opening. These behaviors
338/// are customizable via the `T: TryIntoStateUpdateKind` and `U: ApplyUpdate<T>` generic parameters
339/// respectively.
340#[derive(Debug)]
341pub(crate) struct PersistHandle<T: TryIntoStateUpdateKind, U: ApplyUpdate<T>> {
342    /// The [`Mode`] that this catalog was opened in.
343    pub(crate) mode: Mode,
344    /// Since handle to control compaction.
345    since_handle: SinceHandle<SourceData, (), Timestamp, StorageDiff>,
346    /// Write handle to persist.
347    write_handle: WriteHandle<SourceData, (), Timestamp, StorageDiff>,
348    /// Listener to catalog changes.
349    listen: Listen<SourceData, (), Timestamp, StorageDiff>,
350    /// Handle for connecting to persist.
351    persist_client: PersistClient,
352    /// Catalog shard ID.
353    shard_id: ShardId,
354    /// Cache of the most recent catalog snapshot.
355    ///
356    /// We use a tuple instead of [`StateUpdate`] to make consolidation easier.
357    pub(crate) snapshot: Vec<(T, Timestamp, Diff)>,
358    /// Applies custom processing, filtering, and fencing for each individual update.
359    update_applier: U,
360    /// The current upper of the persist shard.
361    pub(crate) upper: Timestamp,
362    /// The fence token of the catalog, if one exists.
363    fenceable_token: FenceableToken,
364    /// The semantic version of the current binary.
365    catalog_content_version: semver::Version,
366    /// Flag to indicate if bootstrap is complete.
367    bootstrap_complete: bool,
368    /// Metrics for the persist catalog.
369    metrics: Arc<Metrics>,
370    /// Snapshot size at the last amortized consolidation, used by
371    /// [`Self::maybe_consolidate`] to decide when to consolidate. Initialized
372    /// lazily on the first call.
373    size_at_last_consolidation: Option<usize>,
374    /// Counts raw updates applied to this handle.
375    ///
376    /// This distinguishes empty upper progress from content. Memory updates are insufficient
377    /// because kinds such as ID allocators do not produce them.
378    updates_applied: u64,
379}
380
381impl<T: TryIntoStateUpdateKind, U: ApplyUpdate<T>> PersistHandle<T, U> {
382    /// Fetch the current upper of the catalog state.
383    #[mz_ore::instrument]
384    async fn current_upper(&mut self) -> Timestamp {
385        match self.mode {
386            Mode::Writable | Mode::Readonly => {
387                let upper = self.write_handle.fetch_recent_upper().await;
388                antichain_to_timestamp(upper.clone())
389            }
390            Mode::Savepoint => self.upper,
391        }
392    }
393
394    /// Appends `updates` iff the current global upper of the catalog is `self.upper`.
395    ///
396    /// Returns the next upper used to commit the transaction.
397    #[mz_ore::instrument]
398    pub(crate) async fn compare_and_append<S: IntoStateUpdateKindJson>(
399        &mut self,
400        updates: Vec<(S, Diff)>,
401        commit_ts: Timestamp,
402    ) -> Result<Timestamp, CompareAndAppendError> {
403        let updates = updates.into_iter().map(|(kind, diff)| {
404            let kind: StateUpdateKindJson = kind.into();
405            (
406                (Into::<SourceData>::into(kind), ()),
407                commit_ts,
408                diff.into_inner(),
409            )
410        });
411        let next_upper = commit_ts.step_forward();
412        // Upper mismatches are classified by the commit and advance callers.
413        self.compare_and_append_inner(updates, next_upper).await?;
414
415        self.sync(next_upper).await?;
416        Ok(next_upper)
417    }
418
419    /// Compare-and-append `updates` to the catalog shard, advancing the upper to `next_upper`.
420    ///
421    /// On success, updating `self.upper` is left to the caller. The caller can thus decide whether
422    /// or not it needs to sync the catalog.
423    ///
424    /// # Panics
425    ///
426    /// Panics if not in `Writable` mode.
427    /// Panics if `next_upper` is not greater than `self.upper`.
428    async fn compare_and_append_inner(
429        &mut self,
430        updates: impl IntoIterator<Item = ((SourceData, ()), Timestamp, StorageDiff)>,
431        next_upper: Timestamp,
432    ) -> Result<(), CompareAndAppendError> {
433        assert_eq!(self.mode, Mode::Writable);
434        assert!(
435            next_upper > self.upper,
436            "next_upper ({next_upper}) not greater than current upper ({})",
437            self.upper,
438        );
439
440        let res = self
441            .write_handle
442            .compare_and_append(
443                updates,
444                Antichain::from_elem(self.upper),
445                Antichain::from_elem(next_upper),
446            )
447            .await
448            .expect("invalid usage");
449
450        if let Err(e @ UpperMismatch { .. }) = res {
451            // Most likely we were fenced out.
452            // Sync to the current upper to detect that.
453            self.sync_to_current_upper().await?;
454            return Err(e.into());
455        }
456
457        // Lag the shard's upper by 1 to keep it readable.
458        let downgrade_to = Antichain::from_elem(next_upper.saturating_sub(1));
459
460        // The since handle gives us the ability to fence out other downgraders using an opaque token.
461        // (See the method documentation for details.)
462        // That's not needed here, so we use the since handle's opaque token to avoid any comparison
463        // failures.
464        let opaque = self.since_handle.opaque().clone();
465        let downgrade = self
466            .since_handle
467            .maybe_compare_and_downgrade_since(&opaque, (&opaque, &downgrade_to))
468            .await;
469        if let Some(Err(e)) = downgrade {
470            soft_panic_or_log!("found opaque value {e:?}, but expected {opaque:?}");
471        }
472
473        Ok(())
474    }
475
476    /// Accepts an upper mismatch caused only by empty progress.
477    ///
478    /// `updates_applied_before` must be captured before the compare-and-append, which synchronizes
479    /// this handle on mismatch.
480    fn classify_upper_mismatch(
481        &self,
482        updates_applied_before: u64,
483        actual_upper: Timestamp,
484    ) -> Result<(), DurableCatalogError> {
485        if self.updates_applied != updates_applied_before {
486            Err(DurableCatalogError::CatalogOutOfSync {
487                update_count: usize::cast_from(self.updates_applied - updates_applied_before),
488                upper: actual_upper,
489            })
490        } else {
491            Ok(())
492        }
493    }
494
495    /// Generates an iterator of [`StateUpdate`] that contain all unconsolidated updates to the
496    /// catalog state up to, and including, `as_of`.
497    #[mz_ore::instrument]
498    async fn snapshot_unconsolidated(&mut self) -> Vec<StateUpdate<StateUpdateKind>> {
499        let current_upper = self.current_upper().await;
500
501        let mut snapshot = Vec::new();
502        let mut read_handle = self.read_handle().await;
503        let as_of = as_of(&read_handle, current_upper);
504        let mut stream = Box::pin(
505            // We use `snapshot_and_stream` because it guarantees unconsolidated output.
506            read_handle
507                .snapshot_and_stream(Antichain::from_elem(as_of))
508                .await
509                .expect("we have advanced the restart_as_of by the since"),
510        );
511        while let Some(update) = stream.next().await {
512            snapshot.push(update)
513        }
514        read_handle.expire().await;
515        snapshot
516            .into_iter()
517            .map(Into::<StateUpdate<StateUpdateKindJson>>::into)
518            .map(|state_update| state_update.try_into().expect("kind decoding error"))
519            .collect()
520    }
521
522    /// Listen and apply all updates that are currently in persist.
523    ///
524    /// Returns an error if this instance has been fenced out.
525    #[mz_ore::instrument]
526    pub(crate) async fn sync_to_current_upper(&mut self) -> Result<(), FenceError> {
527        let upper = self.current_upper().await;
528        self.sync(upper).await
529    }
530
531    /// Listen and apply all updates up to `target_upper`.
532    ///
533    /// Returns an error if this instance has been fenced out.
534    #[mz_ore::instrument(level = "debug")]
535    pub(crate) async fn sync(&mut self, target_upper: Timestamp) -> Result<(), FenceError> {
536        self.metrics.syncs.inc();
537        let histogram = self.metrics.sync_latency_seconds.clone();
538        self.sync_inner(target_upper)
539            .wall_time()
540            .observe(histogram)
541            .await
542    }
543
544    #[mz_ore::instrument(level = "debug")]
545    async fn sync_inner(&mut self, target_upper: Timestamp) -> Result<(), FenceError> {
546        self.fenceable_token.validate()?;
547
548        // Savepoint catalogs do not yet know how to update themselves in response to concurrent
549        // writes from writer catalogs.
550        if self.mode == Mode::Savepoint {
551            self.upper = max(self.upper, target_upper);
552            return Ok(());
553        }
554
555        let mut updates: BTreeMap<_, Vec<_>> = BTreeMap::new();
556
557        // Reset the amortized consolidation tracker so it picks up the
558        // current snapshot size as its baseline.
559        self.size_at_last_consolidation = None;
560
561        while self.upper < target_upper {
562            let listen_events = self.listen.fetch_next().await;
563            for listen_event in listen_events {
564                match listen_event {
565                    ListenEvent::Progress(upper) => {
566                        debug!("synced up to {upper:?}");
567                        self.upper = antichain_to_timestamp(upper);
568                        // Attempt to apply updates in batches of a single timestamp. If another
569                        // catalog wrote a fence token at one timestamp and then updates in a new
570                        // format at a later timestamp, then we want to apply the fence token
571                        // before attempting to deserialize the new updates.
572                        while let Some((ts, updates)) = updates.pop_first() {
573                            assert!(ts < self.upper, "expected {} < {}", ts, self.upper);
574                            let updates = updates.into_iter().map(
575                                |update: StateUpdate<StateUpdateKindJson>| {
576                                    let kind =
577                                        T::try_from(update.kind).expect("kind decoding error");
578                                    StateUpdate {
579                                        kind,
580                                        ts: update.ts,
581                                        diff: update.diff,
582                                    }
583                                },
584                            );
585                            self.apply_updates(updates)?;
586                            self.maybe_consolidate();
587                        }
588                    }
589                    ListenEvent::Updates(batch_updates) => {
590                        for update in batch_updates {
591                            let update: StateUpdate<StateUpdateKindJson> = update.into();
592                            updates.entry(update.ts).or_default().push(update);
593                        }
594                    }
595                }
596            }
597        }
598        assert_eq!(updates, BTreeMap::new(), "all updates should be applied");
599        // Always consolidate at the end to ensure the snapshot is clean.
600        self.consolidate();
601        Ok(())
602    }
603
604    /// Apply a batch of updates and then consolidate the snapshot. This is the
605    /// typical entry point for callers that apply updates in a single batch.
606    ///
607    /// For hot loops that apply updates across many timestamps (e.g., `sync_inner`),
608    /// use `apply_updates` directly and call `consolidate()` periodically (e.g.,
609    /// on snapshot doubling) to bound memory while staying amortized O(N log N).
610    pub(crate) fn apply_updates_and_consolidate(
611        &mut self,
612        updates: impl IntoIterator<Item = StateUpdate<T>>,
613    ) -> Result<(), FenceError> {
614        self.apply_updates(updates)?;
615        self.consolidate();
616        Ok(())
617    }
618
619    /// Apply a batch of updates to the catalog state without consolidating.
620    ///
621    /// Does NOT consolidate the snapshot afterward. If you are calling this once,
622    /// prefer `apply_updates_and_consolidate`. This method exists for loops that
623    /// call it many times — consolidating per call would be O(K * N log N) instead
624    /// of O(N log N). Callers should consolidate periodically (e.g., on snapshot
625    /// doubling) to bound memory.
626    #[mz_ore::instrument(level = "debug")]
627    fn apply_updates(
628        &mut self,
629        updates: impl IntoIterator<Item = StateUpdate<T>>,
630    ) -> Result<(), FenceError> {
631        let mut updates: Vec<_> = updates
632            .into_iter()
633            .map(|StateUpdate { kind, ts, diff }| (kind, ts, diff))
634            .collect();
635
636        // This helps guarantee that for a single key, there is at most a single retraction and a
637        // single insertion per timestamp. Otherwise, we would need to match the retractions and
638        // insertions up by value and manually figure out what the end value should be.
639        differential_dataflow::consolidation::consolidate_updates(&mut updates);
640
641        // Updates must be applied in timestamp order. Within a timestamp retractions must be
642        // applied before insertions, or we might end up retracting the wrong value.
643        updates.sort_by(|(_, ts1, diff1), (_, ts2, diff2)| ts1.cmp(ts2).then(diff1.cmp(diff2)));
644
645        let mut errors = Vec::new();
646
647        for (kind, ts, diff) in updates {
648            if diff != Diff::ONE && diff != Diff::MINUS_ONE {
649                panic!("invalid update in consolidated trace: ({kind:?}, {ts:?}, {diff:?})");
650            }
651            self.updates_applied += 1;
652
653            match self.update_applier.apply_update(
654                StateUpdate { kind, ts, diff },
655                &mut self.fenceable_token,
656                &self.metrics,
657            ) {
658                Ok(Some(StateUpdate { kind, ts, diff })) => self.snapshot.push((kind, ts, diff)),
659                Ok(None) => {}
660                // Instead of returning immediately, we accumulate all the errors and return the one
661                // with the most information.
662                Err(err) => errors.push(err),
663            }
664        }
665
666        // Track the high-water mark of the unconsolidated snapshot size.
667        let len = i64::try_from(self.snapshot.len()).unwrap_or(i64::MAX);
668        if len > self.metrics.snapshot_max_entries.get() {
669            self.metrics.snapshot_max_entries.set(len);
670        }
671
672        errors.sort();
673        if let Some(err) = errors.into_iter().next() {
674            return Err(err);
675        }
676
677        Ok(())
678    }
679
680    /// Consolidate the snapshot if it has at least doubled in size since the
681    /// last consolidation. This amortizes the O(N log N) consolidation cost
682    /// over many small updates, keeping the total work O(N log N) rather than
683    /// O(K * N log N) for K timestamps.
684    fn maybe_consolidate(&mut self) {
685        let threshold = *self
686            .size_at_last_consolidation
687            // Use a minimum of 8 to avoid consolidating on every update when
688            // the snapshot is small or empty (since 0 * 2 = 0).
689            .get_or_insert_with(|| max(self.snapshot.len(), 8));
690        if self.snapshot.len() >= threshold * 2 {
691            self.consolidate();
692            self.size_at_last_consolidation = Some(self.snapshot.len());
693        }
694    }
695
696    #[mz_ore::instrument]
697    pub(crate) fn consolidate(&mut self) {
698        self.metrics.snapshot_consolidations.inc();
699        soft_assert_no_log!(
700            self.snapshot
701                .windows(2)
702                .all(|updates| updates[0].1 <= updates[1].1),
703            "snapshot should be sorted by timestamp, {:#?}",
704            self.snapshot
705        );
706
707        let new_ts = self
708            .snapshot
709            .last()
710            .map(|(_, ts, _)| *ts)
711            .unwrap_or_else(Timestamp::minimum);
712        for (_, ts, _) in &mut self.snapshot {
713            *ts = new_ts;
714        }
715        differential_dataflow::consolidation::consolidate_updates(&mut self.snapshot);
716    }
717
718    /// Execute and return the results of `f` on the current catalog trace.
719    ///
720    /// Will return an error if the catalog has been fenced out.
721    async fn with_trace<R>(
722        &mut self,
723        f: impl FnOnce(&Vec<(T, Timestamp, Diff)>) -> Result<R, CatalogError>,
724    ) -> Result<R, CatalogError> {
725        self.sync_to_current_upper().await?;
726        f(&self.snapshot)
727    }
728
729    /// Open a read handle to the catalog.
730    async fn read_handle(&self) -> ReadHandle<SourceData, (), Timestamp, StorageDiff> {
731        self.persist_client
732            .open_leased_reader(
733                self.shard_id,
734                Arc::new(persist_desc()),
735                Arc::new(UnitSchema::default()),
736                Diagnostics {
737                    shard_name: CATALOG_SHARD_NAME.to_string(),
738                    handle_purpose: "openable durable catalog state temporary reader".to_string(),
739                },
740                USE_CRITICAL_SINCE_CATALOG.get(self.persist_client.dyncfgs()),
741            )
742            .await
743            .expect("invalid usage")
744    }
745
746    /// Politely releases all external resources that can only be released in an async context.
747    async fn expire(self: Box<Self>) {
748        self.write_handle.expire().await;
749        self.listen.expire().await;
750    }
751}
752
753impl<U: ApplyUpdate<StateUpdateKind>> PersistHandle<StateUpdateKind, U> {
754    /// Execute and return the results of `f` on the current catalog snapshot.
755    ///
756    /// Will return an error if the catalog has been fenced out.
757    async fn with_snapshot<T>(
758        &mut self,
759        f: impl FnOnce(Snapshot) -> Result<T, CatalogError>,
760    ) -> Result<T, CatalogError> {
761        fn apply<K, V>(map: &mut BTreeMap<K, V>, key: &K, value: &V, diff: Diff)
762        where
763            K: Ord + Clone,
764            V: Ord + Clone + Debug,
765        {
766            let key = key.clone();
767            let value = value.clone();
768            if diff == Diff::ONE {
769                let prev = map.insert(key, value);
770                assert_eq!(
771                    prev, None,
772                    "values must be explicitly retracted before inserting a new value"
773                );
774            } else if diff == Diff::MINUS_ONE {
775                let prev = map.remove(&key);
776                assert_eq!(
777                    prev,
778                    Some(value),
779                    "retraction does not match existing value"
780                );
781            }
782        }
783
784        self.with_trace(|trace| {
785            let mut snapshot = Snapshot::empty();
786            for (kind, ts, diff) in trace {
787                let diff = *diff;
788                if diff != Diff::ONE && diff != Diff::MINUS_ONE {
789                    panic!("invalid update in consolidated trace: ({kind:?}, {ts:?}, {diff:?})");
790                }
791
792                match kind {
793                    StateUpdateKind::AuditLog(_key, ()) => {
794                        // Ignore for snapshots.
795                    }
796                    StateUpdateKind::Cluster(key, value) => {
797                        apply(&mut snapshot.clusters, key, value, diff);
798                    }
799                    StateUpdateKind::ClusterReplica(key, value) => {
800                        apply(&mut snapshot.cluster_replicas, key, value, diff);
801                    }
802                    StateUpdateKind::Comment(key, value) => {
803                        apply(&mut snapshot.comments, key, value, diff);
804                    }
805                    StateUpdateKind::Config(key, value) => {
806                        apply(&mut snapshot.configs, key, value, diff);
807                    }
808                    StateUpdateKind::Database(key, value) => {
809                        apply(&mut snapshot.databases, key, value, diff);
810                    }
811                    StateUpdateKind::DefaultPrivilege(key, value) => {
812                        apply(&mut snapshot.default_privileges, key, value, diff);
813                    }
814                    StateUpdateKind::FenceToken(_token) => {
815                        // Ignore for snapshots.
816                    }
817                    StateUpdateKind::IdAllocator(key, value) => {
818                        apply(&mut snapshot.id_allocator, key, value, diff);
819                    }
820                    StateUpdateKind::IntrospectionSourceIndex(key, value) => {
821                        apply(&mut snapshot.introspection_sources, key, value, diff);
822                    }
823                    StateUpdateKind::Item(key, value) => {
824                        apply(&mut snapshot.items, key, value, diff);
825                    }
826                    StateUpdateKind::NetworkPolicy(key, value) => {
827                        apply(&mut snapshot.network_policies, key, value, diff);
828                    }
829                    StateUpdateKind::Role(key, value) => {
830                        apply(&mut snapshot.roles, key, value, diff);
831                    }
832                    StateUpdateKind::Schema(key, value) => {
833                        apply(&mut snapshot.schemas, key, value, diff);
834                    }
835                    StateUpdateKind::Setting(key, value) => {
836                        apply(&mut snapshot.settings, key, value, diff);
837                    }
838                    StateUpdateKind::SourceReferences(key, value) => {
839                        apply(&mut snapshot.source_references, key, value, diff);
840                    }
841                    StateUpdateKind::SystemConfiguration(key, value) => {
842                        apply(&mut snapshot.system_configurations, key, value, diff);
843                    }
844                    StateUpdateKind::ClusterSystemConfiguration(key, value) => {
845                        apply(
846                            &mut snapshot.cluster_system_configurations,
847                            key,
848                            value,
849                            diff,
850                        );
851                    }
852                    StateUpdateKind::ReplicaSystemConfiguration(key, value) => {
853                        apply(
854                            &mut snapshot.replica_system_configurations,
855                            key,
856                            value,
857                            diff,
858                        );
859                    }
860                    StateUpdateKind::SystemObjectMapping(key, value) => {
861                        apply(&mut snapshot.system_object_mappings, key, value, diff);
862                    }
863                    StateUpdateKind::SystemPrivilege(key, value) => {
864                        apply(&mut snapshot.system_privileges, key, value, diff);
865                    }
866                    StateUpdateKind::StorageCollectionMetadata(key, value) => {
867                        apply(&mut snapshot.storage_collection_metadata, key, value, diff);
868                    }
869                    StateUpdateKind::UnfinalizedShard(key, ()) => {
870                        apply(&mut snapshot.unfinalized_shards, key, &(), diff);
871                    }
872                    StateUpdateKind::TxnWalShard((), value) => {
873                        apply(&mut snapshot.txn_wal_shard, &(), value, diff);
874                    }
875                    StateUpdateKind::RoleAuth(key, value) => {
876                        apply(&mut snapshot.role_auth, key, value, diff);
877                    }
878                }
879            }
880            f(snapshot)
881        })
882        .await
883    }
884
885    /// Generates an iterator of [`StateUpdate`] that contain all updates to the catalog
886    /// state.
887    ///
888    /// The output is fetched directly from persist instead of the in-memory cache.
889    ///
890    /// The output is consolidated and sorted by timestamp in ascending order.
891    #[mz_ore::instrument(level = "debug")]
892    async fn persist_snapshot(&self) -> impl Iterator<Item = StateUpdate> + DoubleEndedIterator {
893        let mut read_handle = self.read_handle().await;
894        let as_of = as_of(&read_handle, self.upper);
895        let snapshot = snapshot_binary(&mut read_handle, as_of, &self.metrics)
896            .await
897            .map(|update| update.try_into().expect("kind decoding error"));
898        read_handle.expire().await;
899        snapshot
900    }
901}
902
903/// Applies updates for an unopened catalog.
904#[derive(Debug)]
905pub(crate) struct UnopenedCatalogStateInner {
906    /// A cache of the config collection of the catalog.
907    configs: BTreeMap<String, u64>,
908    /// A cache of the settings collection of the catalog.
909    settings: BTreeMap<String, String>,
910}
911
912impl UnopenedCatalogStateInner {
913    fn new() -> UnopenedCatalogStateInner {
914        UnopenedCatalogStateInner {
915            configs: BTreeMap::new(),
916            settings: BTreeMap::new(),
917        }
918    }
919}
920
921impl ApplyUpdate<StateUpdateKindJson> for UnopenedCatalogStateInner {
922    fn apply_update(
923        &mut self,
924        update: StateUpdate<StateUpdateKindJson>,
925        current_fence_token: &mut FenceableToken,
926        _metrics: &Arc<Metrics>,
927    ) -> Result<Option<StateUpdate<StateUpdateKindJson>>, FenceError> {
928        if !update.kind.is_audit_log() && update.kind.is_always_deserializable() {
929            let kind = TryInto::try_into(&update.kind).expect("kind is known to be deserializable");
930            match (kind, update.diff) {
931                (StateUpdateKind::Config(key, value), Diff::ONE) => {
932                    let prev = self.configs.insert(key.key, value.value);
933                    assert_eq!(
934                        prev, None,
935                        "values must be explicitly retracted before inserting a new value"
936                    );
937                }
938                (StateUpdateKind::Config(key, value), Diff::MINUS_ONE) => {
939                    let prev = self.configs.remove(&key.key);
940                    assert_eq!(
941                        prev,
942                        Some(value.value),
943                        "retraction does not match existing value"
944                    );
945                }
946                (StateUpdateKind::Setting(key, value), Diff::ONE) => {
947                    let prev = self.settings.insert(key.name, value.value);
948                    assert_eq!(
949                        prev, None,
950                        "values must be explicitly retracted before inserting a new value"
951                    );
952                }
953                (StateUpdateKind::Setting(key, value), Diff::MINUS_ONE) => {
954                    let prev = self.settings.remove(&key.name);
955                    assert_eq!(
956                        prev,
957                        Some(value.value),
958                        "retraction does not match existing value"
959                    );
960                }
961                (StateUpdateKind::FenceToken(fence_token), Diff::ONE) => {
962                    current_fence_token.maybe_fence(fence_token)?;
963                }
964                _ => {}
965            }
966        }
967
968        Ok(Some(update))
969    }
970}
971
972/// A Handle to an unopened catalog stored in persist. The unopened catalog can serve `Config` data,
973/// `Setting` data, or the current epoch. All other catalog data may be un-migrated and should not
974/// be read until the catalog has been opened. The [`UnopenedPersistCatalogState`] is responsible
975/// for opening the catalog, see [`OpenableDurableCatalogState::open`] for more details.
976///
977/// Production users should call [`Self::expire`] before dropping an [`UnopenedPersistCatalogState`]
978/// so that it can expire its leases. If/when rust gets AsyncDrop, this will be done automatically.
979pub(crate) type UnopenedPersistCatalogState =
980    PersistHandle<StateUpdateKindJson, UnopenedCatalogStateInner>;
981
982impl UnopenedPersistCatalogState {
983    /// Create a new [`UnopenedPersistCatalogState`] to the catalog state associated with
984    /// `organization_id`.
985    ///
986    /// All usages of the persist catalog must go through this function. That includes the
987    /// catalog-debug tool, the adapter's catalog, etc.
988    #[mz_ore::instrument]
989    pub(crate) async fn new(
990        persist_client: PersistClient,
991        organization_id: Uuid,
992        version: semver::Version,
993        deploy_generation: Option<u64>,
994        metrics: Arc<Metrics>,
995    ) -> Result<UnopenedPersistCatalogState, DurableCatalogError> {
996        let catalog_shard_id = shard_id(organization_id, CATALOG_SEED);
997        debug!(?catalog_shard_id, "new persist backed catalog state");
998
999        // Check the catalog shard version to ensure that we are compatible with the persist
1000        // data format. This lets us return an error gracefully, rather than panicking later in
1001        // persist.
1002        let version_in_catalog_shard =
1003            fetch_catalog_shard_version(&persist_client, catalog_shard_id).await;
1004        if let Some(version_in_catalog_shard) = version_in_catalog_shard {
1005            if !mz_persist_client::cfg::code_can_write_data(&version, &version_in_catalog_shard) {
1006                return Err(DurableCatalogError::IncompatiblePersistVersion {
1007                    found_version: version_in_catalog_shard,
1008                    catalog_version: version,
1009                });
1010            }
1011        }
1012
1013        let open_handles_start = Instant::now();
1014        info!("startup: envd serve: catalog init: open handles beginning");
1015        let since_handle = persist_client
1016            .open_critical_since(
1017                catalog_shard_id,
1018                CATALOG_CRITICAL_SINCE.clone(),
1019                Opaque::encode(&i64::MIN),
1020                Diagnostics {
1021                    shard_name: CATALOG_SHARD_NAME.to_string(),
1022                    handle_purpose: "durable catalog state critical since".to_string(),
1023                },
1024            )
1025            .await
1026            .expect("invalid usage");
1027
1028        let (mut write_handle, mut read_handle) = persist_client
1029            .open(
1030                catalog_shard_id,
1031                Arc::new(persist_desc()),
1032                Arc::new(UnitSchema::default()),
1033                Diagnostics {
1034                    shard_name: CATALOG_SHARD_NAME.to_string(),
1035                    handle_purpose: "durable catalog state handles".to_string(),
1036                },
1037                USE_CRITICAL_SINCE_CATALOG.get(persist_client.dyncfgs()),
1038            )
1039            .await
1040            .expect("invalid usage");
1041        info!(
1042            "startup: envd serve: catalog init: open handles complete in {:?}",
1043            open_handles_start.elapsed()
1044        );
1045
1046        // Commit an empty write at the minimum timestamp so the catalog is always readable.
1047        let upper = {
1048            const EMPTY_UPDATES: &[((SourceData, ()), Timestamp, StorageDiff)] = &[];
1049            let upper = Antichain::from_elem(Timestamp::minimum());
1050            let next_upper = Timestamp::minimum().step_forward();
1051            match write_handle
1052                .compare_and_append(EMPTY_UPDATES, upper, Antichain::from_elem(next_upper))
1053                .await
1054                .expect("invalid usage")
1055            {
1056                Ok(()) => next_upper,
1057                Err(mismatch) => antichain_to_timestamp(mismatch.current),
1058            }
1059        };
1060
1061        let snapshot_start = Instant::now();
1062        info!("startup: envd serve: catalog init: snapshot beginning");
1063        let as_of = as_of(&read_handle, upper);
1064        let snapshot: Vec<_> = snapshot_binary(&mut read_handle, as_of, &metrics)
1065            .await
1066            .map(|StateUpdate { kind, ts, diff }| (kind, ts, diff))
1067            .collect();
1068        let listen = read_handle
1069            .listen(Antichain::from_elem(as_of))
1070            .await
1071            .expect("invalid usage");
1072        info!(
1073            "startup: envd serve: catalog init: snapshot complete in {:?}",
1074            snapshot_start.elapsed()
1075        );
1076
1077        let mut handle = UnopenedPersistCatalogState {
1078            // Unopened catalogs are always writeable until they're opened in an explicit mode.
1079            mode: Mode::Writable,
1080            since_handle,
1081            write_handle,
1082            listen,
1083            persist_client,
1084            shard_id: catalog_shard_id,
1085            // Initialize empty in-memory state.
1086            snapshot: Vec::new(),
1087            update_applier: UnopenedCatalogStateInner::new(),
1088            upper,
1089            fenceable_token: FenceableToken::new(deploy_generation),
1090            catalog_content_version: version,
1091            bootstrap_complete: false,
1092            metrics,
1093            size_at_last_consolidation: None,
1094            updates_applied: 0,
1095        };
1096        // If the snapshot is not consolidated, and we see multiple epoch values while applying the
1097        // updates, then we might accidentally fence ourselves out.
1098        soft_assert_no_log!(
1099            snapshot.iter().all(|(_, _, diff)| *diff == Diff::ONE),
1100            "snapshot should be consolidated: {snapshot:#?}"
1101        );
1102
1103        let apply_start = Instant::now();
1104        info!("startup: envd serve: catalog init: apply updates beginning");
1105        let updates = snapshot
1106            .into_iter()
1107            .map(|(kind, ts, diff)| StateUpdate { kind, ts, diff });
1108        handle.apply_updates_and_consolidate(updates)?;
1109        info!(
1110            "startup: envd serve: catalog init: apply updates complete in {:?}",
1111            apply_start.elapsed()
1112        );
1113
1114        // Validate that the binary version of the current process is not less than any binary
1115        // version that has written to the catalog.
1116        // This condition is only checked once, right here. If a new process comes along with a
1117        // higher version, it must fence this process out with one of the existing fencing
1118        // mechanisms.
1119        if let Some(found_version) = handle.get_catalog_content_version().await? {
1120            // Use cmp_precedence() to ignore build metadata per SemVer 2.0.0 spec
1121            if handle
1122                .catalog_content_version
1123                .cmp_precedence(&found_version)
1124                == std::cmp::Ordering::Less
1125            {
1126                return Err(DurableCatalogError::IncompatiblePersistVersion {
1127                    found_version,
1128                    catalog_version: handle.catalog_content_version,
1129                });
1130            }
1131        }
1132
1133        Ok(handle)
1134    }
1135
1136    #[mz_ore::instrument]
1137    async fn open_inner(
1138        mut self,
1139        mode: Mode,
1140        initial_ts: Timestamp,
1141        bootstrap_args: &BootstrapArgs,
1142    ) -> Result<Box<dyn DurableCatalogState>, CatalogError> {
1143        // It would be nice to use `initial_ts` here, but it comes from the system clock, not the
1144        // timestamp oracle.
1145        let mut commit_ts = self.upper;
1146        self.mode = mode;
1147
1148        // Validate the current deploy generation.
1149        match (&self.mode, &self.fenceable_token) {
1150            (_, FenceableToken::Unfenced { .. } | FenceableToken::Fenced { .. }) => {
1151                return Err(DurableCatalogError::Internal(
1152                    "catalog should not have fenced before opening".to_string(),
1153                )
1154                .into());
1155            }
1156            (
1157                Mode::Writable | Mode::Savepoint,
1158                FenceableToken::Initializing {
1159                    current_deploy_generation: None,
1160                    ..
1161                },
1162            ) => {
1163                return Err(DurableCatalogError::Internal(format!(
1164                    "cannot open in mode '{:?}' without a deploy generation",
1165                    self.mode,
1166                ))
1167                .into());
1168            }
1169            _ => {}
1170        }
1171
1172        let read_only = matches!(self.mode, Mode::Readonly);
1173
1174        // Fence out previous catalogs.
1175        loop {
1176            self.sync_to_current_upper().await?;
1177            commit_ts = max(commit_ts, self.upper);
1178            let (fence_updates, current_fenceable_token) = self
1179                .fenceable_token
1180                .generate_unfenced_token(self.mode)?
1181                .ok_or_else(|| {
1182                    DurableCatalogError::Internal(
1183                        "catalog should not have fenced before opening".to_string(),
1184                    )
1185                })?;
1186            debug!(
1187                ?self.upper,
1188                ?self.fenceable_token,
1189                ?current_fenceable_token,
1190                "fencing previous catalogs"
1191            );
1192            if matches!(self.mode, Mode::Writable) {
1193                match self
1194                    .compare_and_append(fence_updates.clone(), commit_ts)
1195                    .await
1196                {
1197                    Ok(upper) => {
1198                        commit_ts = upper;
1199                    }
1200                    Err(CompareAndAppendError::Fence(e)) => return Err(e.into()),
1201                    Err(e @ CompareAndAppendError::UpperMismatch { .. }) => {
1202                        warn!("catalog write failed due to upper mismatch, retrying: {e:?}");
1203                        continue;
1204                    }
1205                }
1206            }
1207            self.fenceable_token = current_fenceable_token;
1208            break;
1209        }
1210
1211        if matches!(self.mode, Mode::Writable) {
1212            // One-time migration: The catalog previously used `CONTROLLER_CRITICAL_SINCE` for its
1213            // since handle. Now it uses its own `CATALOG_CRITICAL_SINCE`, to free
1214            // `CONTROLLER_CRITICAL_SINCE` up for the storage controller. The catalog and
1215            // controller handles differ in the `Opaque` codec, so we need a migration.
1216            //
1217            // TODO: Remove this once we don't support upgrading from v26 anymore.
1218            let mut controller_handle = self
1219                .persist_client
1220                .open_critical_since::<SourceData, (), Timestamp, StorageDiff>(
1221                    self.shard_id,
1222                    PersistClient::CONTROLLER_CRITICAL_SINCE,
1223                    Opaque::encode(&i64::MIN),
1224                    Diagnostics {
1225                        shard_name: CATALOG_SHARD_NAME.to_string(),
1226                        handle_purpose: "durable catalog state critical since (migration)"
1227                            .to_string(),
1228                    },
1229                )
1230                .await
1231                .expect("invalid usage");
1232
1233            let since = controller_handle.since().clone();
1234            let res = controller_handle
1235                .compare_and_downgrade_since(
1236                    &Opaque::encode(&i64::MIN),
1237                    (&Opaque::encode(&PersistEpoch::default()), &since),
1238                )
1239                .await;
1240            match res {
1241                Ok(_) => info!("migrated Opaque of catalog since handle"),
1242                Err(_) => { /* critical since was already migrated */ }
1243            }
1244        }
1245
1246        let is_initialized = self.is_initialized_inner();
1247        if !matches!(self.mode, Mode::Writable) && !is_initialized {
1248            return Err(CatalogError::Durable(DurableCatalogError::NotWritable(
1249                format!(
1250                    "catalog tables do not exist; will not create in {:?} mode",
1251                    self.mode
1252                ),
1253            )));
1254        }
1255        soft_assert_ne_or_log!(self.upper, Timestamp::minimum());
1256
1257        // Audit log entries are served from `mz_internal.mz_catalog_raw` via
1258        // the `mz_audit_events` materialized view, so they do not need to live
1259        // in the in-memory catalog snapshot. Drop them here and only keep the
1260        // count for metrics.
1261        let (audit_logs, snapshot): (Vec<_>, Vec<_>) = self
1262            .snapshot
1263            .into_iter()
1264            .partition(|(update, _, _)| update.is_audit_log());
1265        self.snapshot = snapshot;
1266        let audit_log_count = audit_logs.iter().map(|(_, _, diff)| diff).sum::<Diff>();
1267        drop(audit_logs);
1268
1269        // Perform data migrations.
1270        if is_initialized && !read_only {
1271            commit_ts = upgrade(&mut self, commit_ts).await?;
1272        }
1273
1274        debug!(
1275            ?is_initialized,
1276            ?self.upper,
1277            "initializing catalog state"
1278        );
1279        let mut catalog = PersistCatalogState {
1280            mode: self.mode,
1281            since_handle: self.since_handle,
1282            write_handle: self.write_handle,
1283            listen: self.listen,
1284            persist_client: self.persist_client,
1285            shard_id: self.shard_id,
1286            upper: self.upper,
1287            fenceable_token: self.fenceable_token,
1288            // Initialize empty in-memory state.
1289            snapshot: Vec::new(),
1290            update_applier: CatalogStateInner::new(),
1291            catalog_content_version: self.catalog_content_version,
1292            bootstrap_complete: false,
1293            metrics: self.metrics,
1294            size_at_last_consolidation: None,
1295            updates_applied: 0,
1296        };
1297        catalog.metrics.collection_entries.reset();
1298        // Normally, `collection_entries` is updated in `apply_updates`. The audit log updates skip
1299        // over that function so we manually update it here.
1300        catalog
1301            .metrics
1302            .collection_entries
1303            .with_label_values(&[&CollectionType::AuditLog.to_string()])
1304            .add(audit_log_count.into_inner());
1305        let updates = self.snapshot.into_iter().map(|(kind, ts, diff)| {
1306            let kind = TryIntoStateUpdateKind::try_into(kind).expect("kind decoding error");
1307            StateUpdate { kind, ts, diff }
1308        });
1309        catalog.apply_updates_and_consolidate(updates)?;
1310
1311        let catalog_content_version = catalog.catalog_content_version.to_string();
1312        let txn = if is_initialized {
1313            let mut txn = catalog.transaction_unchecked().await?;
1314
1315            // Ad-hoc migration: Initialize the `migration_version` expected by adapter to be
1316            // present in existing catalogs.
1317            //
1318            // Note: Need to exclude read-only catalog mode here, because in that mode all
1319            // transactions are expected to be no-ops.
1320            // TODO: remove this once we only support upgrades from version >= 0.164
1321            if txn.get_setting("migration_version".into()).is_none() && mode != Mode::Readonly {
1322                let old_version = txn.get_catalog_content_version();
1323                txn.set_setting("migration_version".into(), old_version.map(Into::into))?;
1324            }
1325
1326            // Opening the catalog with write intent fences out every previous
1327            // catalog owner, so all sessions served by previous owners are
1328            // dead. Reclaim the temporary items they owned here, before
1329            // anything else reads the catalog.
1330            //
1331            // NOTE: This only reclaims on the fence, which covers a
1332            // single-writer world where every crash is followed by some
1333            // process's writable open. Once several serving envds run
1334            // concurrently, a peer crash triggers no fence here, so
1335            // reclaiming its sessions' items needs the durable
1336            // envd-heartbeat mechanism described in the durable temporary
1337            // objects design doc.
1338            if mode != Mode::Readonly {
1339                txn.remove_ephemeral_items();
1340            }
1341
1342            txn.set_catalog_content_version(catalog_content_version)?;
1343            txn
1344        } else {
1345            soft_assert_eq_no_log!(
1346                catalog
1347                    .snapshot
1348                    .iter()
1349                    .filter(|(kind, _, _)| !matches!(kind, StateUpdateKind::FenceToken(_)))
1350                    .count(),
1351                0,
1352                "trace should not contain any updates for an uninitialized catalog: {:#?}",
1353                catalog.snapshot
1354            );
1355
1356            let mut txn = catalog.transaction_unchecked().await?;
1357            initialize::initialize(
1358                &mut txn,
1359                bootstrap_args,
1360                initial_ts.into(),
1361                catalog_content_version,
1362            )
1363            .await?;
1364            txn
1365        };
1366
1367        if read_only {
1368            let (txn_batch, _) = txn.into_parts()?;
1369            // The upper here doesn't matter because we are only applying the updates in memory.
1370            let updates = StateUpdate::from_txn_batch_ts(txn_batch, catalog.upper);
1371            catalog.apply_updates_and_consolidate(updates)?;
1372        } else {
1373            txn.commit_internal(commit_ts).await?;
1374        }
1375
1376        if matches!(catalog.mode, Mode::Writable) {
1377            let write_handle = catalog
1378                .persist_client
1379                .open_writer::<SourceData, (), Timestamp, i64>(
1380                    catalog.write_handle.shard_id(),
1381                    Arc::new(persist_desc()),
1382                    Arc::new(UnitSchema::default()),
1383                    Diagnostics {
1384                        shard_name: CATALOG_SHARD_NAME.to_string(),
1385                        handle_purpose: "compact catalog".to_string(),
1386                    },
1387                )
1388                .await
1389                .expect("invalid usage");
1390            let fuel = CATALOG_FORCE_COMPACTION_FUEL.handle(catalog.persist_client.dyncfgs());
1391            let wait = CATALOG_FORCE_COMPACTION_WAIT.handle(catalog.persist_client.dyncfgs());
1392            // We're going to gradually turn this on via dyncfgs. Run it in a task so that it
1393            // doesn't block startup.
1394            let _task = mz_ore::task::spawn(|| "catalog::force_shard_compaction", async move {
1395                let () =
1396                    mz_persist_client::cli::admin::dangerous_force_compaction_and_break_pushdown(
1397                        &write_handle,
1398                        || fuel.get(),
1399                        || wait.get(),
1400                    )
1401                    .await;
1402            });
1403        }
1404
1405        Ok(Box::new(catalog))
1406    }
1407
1408    /// Reports if the catalog state has been initialized.
1409    ///
1410    /// NOTE: This is the answer as of the last call to [`PersistHandle::sync`] or [`PersistHandle::sync_to_current_upper`],
1411    /// not necessarily what is currently in persist.
1412    #[mz_ore::instrument]
1413    fn is_initialized_inner(&self) -> bool {
1414        !self.update_applier.configs.is_empty()
1415    }
1416
1417    /// Get the current value of config `key`.
1418    ///
1419    /// Some configs need to be read before the catalog is opened for bootstrapping.
1420    #[mz_ore::instrument]
1421    async fn get_current_config(&mut self, key: &str) -> Result<Option<u64>, DurableCatalogError> {
1422        self.sync_to_current_upper().await?;
1423        Ok(self.update_applier.configs.get(key).cloned())
1424    }
1425
1426    /// Get the user version of this instance.
1427    ///
1428    /// The user version is used to determine if a migration is needed.
1429    #[mz_ore::instrument]
1430    pub(crate) async fn get_user_version(&mut self) -> Result<Option<u64>, DurableCatalogError> {
1431        self.get_current_config(USER_VERSION_KEY).await
1432    }
1433
1434    /// Get the current value of setting `name`.
1435    ///
1436    /// Some settings need to be read before the catalog is opened for bootstrapping.
1437    #[mz_ore::instrument]
1438    async fn get_current_setting(
1439        &mut self,
1440        name: &str,
1441    ) -> Result<Option<String>, DurableCatalogError> {
1442        self.sync_to_current_upper().await?;
1443        Ok(self.update_applier.settings.get(name).cloned())
1444    }
1445
1446    /// Get the catalog content version.
1447    ///
1448    /// The catalog content version is the semantic version of the most recent binary that wrote to
1449    /// the catalog.
1450    #[mz_ore::instrument]
1451    async fn get_catalog_content_version(
1452        &mut self,
1453    ) -> Result<Option<semver::Version>, DurableCatalogError> {
1454        let version = self
1455            .get_current_setting(CATALOG_CONTENT_VERSION_KEY)
1456            .await?;
1457        let version = version.map(|version| version.parse().expect("invalid version persisted"));
1458        Ok(version)
1459    }
1460}
1461
1462#[async_trait]
1463impl OpenableDurableCatalogState for UnopenedPersistCatalogState {
1464    #[mz_ore::instrument]
1465    async fn open_savepoint(
1466        mut self: Box<Self>,
1467        initial_ts: Timestamp,
1468        bootstrap_args: &BootstrapArgs,
1469    ) -> Result<Box<dyn DurableCatalogState>, CatalogError> {
1470        self.open_inner(Mode::Savepoint, initial_ts, bootstrap_args)
1471            .boxed()
1472            .await
1473    }
1474
1475    #[mz_ore::instrument]
1476    async fn open_read_only(
1477        mut self: Box<Self>,
1478        bootstrap_args: &BootstrapArgs,
1479    ) -> Result<Box<dyn DurableCatalogState>, CatalogError> {
1480        self.open_inner(Mode::Readonly, EpochMillis::MIN.into(), bootstrap_args)
1481            .boxed()
1482            .await
1483    }
1484
1485    #[mz_ore::instrument]
1486    async fn open(
1487        mut self: Box<Self>,
1488        initial_ts: Timestamp,
1489        bootstrap_args: &BootstrapArgs,
1490    ) -> Result<Box<dyn DurableCatalogState>, CatalogError> {
1491        self.open_inner(Mode::Writable, initial_ts, bootstrap_args)
1492            .boxed()
1493            .await
1494    }
1495
1496    #[mz_ore::instrument(level = "debug")]
1497    async fn open_debug(mut self: Box<Self>) -> Result<DebugCatalogState, CatalogError> {
1498        Ok(DebugCatalogState(*self))
1499    }
1500
1501    #[mz_ore::instrument]
1502    async fn is_initialized(&mut self) -> Result<bool, CatalogError> {
1503        self.sync_to_current_upper().await?;
1504        Ok(self.is_initialized_inner())
1505    }
1506
1507    #[mz_ore::instrument]
1508    async fn epoch(&mut self) -> Result<Epoch, CatalogError> {
1509        self.sync_to_current_upper().await?;
1510        self.fenceable_token
1511            .validate()?
1512            .map(|token| token.epoch)
1513            .ok_or(CatalogError::Durable(DurableCatalogError::Uninitialized))
1514    }
1515
1516    #[mz_ore::instrument]
1517    async fn get_deployment_generation(&mut self) -> Result<u64, CatalogError> {
1518        self.sync_to_current_upper().await?;
1519        self.fenceable_token
1520            .token()
1521            .map(|token| token.deploy_generation)
1522            .ok_or(CatalogError::Durable(DurableCatalogError::Uninitialized))
1523    }
1524
1525    #[mz_ore::instrument(level = "debug")]
1526    async fn get_0dt_deployment_max_wait(&mut self) -> Result<Option<Duration>, CatalogError> {
1527        let value = self
1528            .get_current_config(WITH_0DT_DEPLOYMENT_MAX_WAIT)
1529            .await?;
1530        match value {
1531            None => Ok(None),
1532            Some(millis) => Ok(Some(Duration::from_millis(millis))),
1533        }
1534    }
1535
1536    #[mz_ore::instrument(level = "debug")]
1537    async fn get_0dt_deployment_ddl_check_interval(
1538        &mut self,
1539    ) -> Result<Option<Duration>, CatalogError> {
1540        let value = self
1541            .get_current_config(WITH_0DT_DEPLOYMENT_DDL_CHECK_INTERVAL)
1542            .await?;
1543        match value {
1544            None => Ok(None),
1545            Some(millis) => Ok(Some(Duration::from_millis(millis))),
1546        }
1547    }
1548
1549    #[mz_ore::instrument(level = "debug")]
1550    async fn get_enable_0dt_deployment_panic_after_timeout(
1551        &mut self,
1552    ) -> Result<Option<bool>, CatalogError> {
1553        let value = self
1554            .get_current_config(ENABLE_0DT_DEPLOYMENT_PANIC_AFTER_TIMEOUT)
1555            .await?;
1556        match value {
1557            None => Ok(None),
1558            Some(0) => Ok(Some(false)),
1559            Some(1) => Ok(Some(true)),
1560            Some(v) => Err(
1561                DurableCatalogError::from(TryFromProtoError::UnknownEnumVariant(format!(
1562                    "{v} is not a valid boolean value"
1563                )))
1564                .into(),
1565            ),
1566        }
1567    }
1568
1569    #[mz_ore::instrument]
1570    async fn has_system_config_synced_once(&mut self) -> Result<bool, DurableCatalogError> {
1571        self.get_current_config(SYSTEM_CONFIG_SYNCED_KEY)
1572            .await
1573            .map(|value| value.map(|value| value > 0).unwrap_or(false))
1574    }
1575
1576    #[mz_ore::instrument]
1577    async fn trace_unconsolidated(&mut self) -> Result<Trace, CatalogError> {
1578        self.sync_to_current_upper().await?;
1579        if self.is_initialized_inner() {
1580            let snapshot = self.snapshot_unconsolidated().await;
1581            Ok(Trace::from_snapshot(snapshot))
1582        } else {
1583            Err(CatalogError::Durable(DurableCatalogError::Uninitialized))
1584        }
1585    }
1586
1587    #[mz_ore::instrument]
1588    async fn trace_consolidated(&mut self) -> Result<Trace, CatalogError> {
1589        self.sync_to_current_upper().await?;
1590        if self.is_initialized_inner() {
1591            let snapshot = self.current_snapshot().await?;
1592            Ok(Trace::from_snapshot(snapshot))
1593        } else {
1594            Err(CatalogError::Durable(DurableCatalogError::Uninitialized))
1595        }
1596    }
1597
1598    #[mz_ore::instrument(level = "debug")]
1599    async fn expire(self: Box<Self>) {
1600        self.expire().await
1601    }
1602}
1603
1604/// Applies updates for an opened catalog.
1605#[derive(Debug)]
1606struct CatalogStateInner {
1607    /// A trace of all catalog updates that can be consumed by some higher layer.
1608    updates: VecDeque<memory::objects::StateUpdate>,
1609}
1610
1611impl CatalogStateInner {
1612    fn new() -> CatalogStateInner {
1613        CatalogStateInner {
1614            updates: VecDeque::new(),
1615        }
1616    }
1617}
1618
1619impl ApplyUpdate<StateUpdateKind> for CatalogStateInner {
1620    fn apply_update(
1621        &mut self,
1622        update: StateUpdate<StateUpdateKind>,
1623        current_fence_token: &mut FenceableToken,
1624        metrics: &Arc<Metrics>,
1625    ) -> Result<Option<StateUpdate<StateUpdateKind>>, FenceError> {
1626        if let Some(collection_type) = update.kind.collection_type() {
1627            metrics
1628                .collection_entries
1629                .with_label_values(&[&collection_type.to_string()])
1630                .add(update.diff.into_inner());
1631        }
1632
1633        {
1634            let update: Option<memory::objects::StateUpdate> = (&update)
1635                .try_into()
1636                .expect("invalid persisted update: {update:#?}");
1637            if let Some(update) = update {
1638                self.updates.push_back(update);
1639            }
1640        }
1641
1642        match (update.kind, update.diff) {
1643            (StateUpdateKind::AuditLog(_, ()), _) => Ok(None),
1644            // Nothing to due for fence token retractions but wait for the next insertion.
1645            (StateUpdateKind::FenceToken(_), Diff::MINUS_ONE) => Ok(None),
1646            (StateUpdateKind::FenceToken(token), Diff::ONE) => {
1647                current_fence_token.maybe_fence(token)?;
1648                Ok(None)
1649            }
1650            (kind, diff) => Ok(Some(StateUpdate {
1651                kind,
1652                ts: update.ts,
1653                diff,
1654            })),
1655        }
1656    }
1657}
1658
1659/// A durable store of the catalog state using Persist as an implementation. The durable store can
1660/// serve any catalog data and transactionally modify catalog data.
1661///
1662/// Production users should call [`Self::expire`] before dropping a [`PersistCatalogState`]
1663/// so that it can expire its leases. If/when rust gets AsyncDrop, this will be done automatically.
1664type PersistCatalogState = PersistHandle<StateUpdateKind, CatalogStateInner>;
1665
1666impl PersistHandle<StateUpdateKind, CatalogStateInner> {
1667    /// Creates a transaction without validating pending catalog updates.
1668    async fn transaction_unchecked(&mut self) -> Result<Transaction<'_>, CatalogError> {
1669        self.metrics.transactions_started.inc();
1670        let snapshot = self.snapshot().await?;
1671        let commit_ts = self.upper;
1672        Transaction::new(self, snapshot, commit_ts)
1673    }
1674}
1675
1676#[async_trait]
1677impl ReadOnlyDurableCatalogState for PersistCatalogState {
1678    fn epoch(&self) -> Epoch {
1679        self.fenceable_token
1680            .token()
1681            .expect("opened catalog state must have an epoch")
1682            .epoch
1683    }
1684
1685    fn metrics(&self) -> &Metrics {
1686        &self.metrics
1687    }
1688
1689    #[mz_ore::instrument(level = "debug")]
1690    async fn expire(self: Box<Self>) {
1691        self.expire().await
1692    }
1693
1694    fn is_bootstrap_complete(&self) -> bool {
1695        self.bootstrap_complete
1696    }
1697
1698    async fn get_audit_logs(&mut self) -> Result<Vec<VersionedEvent>, CatalogError> {
1699        self.sync_to_current_upper().await?;
1700        let audit_logs: Vec<_> = self
1701            .persist_snapshot()
1702            .await
1703            .filter_map(
1704                |StateUpdate {
1705                     kind,
1706                     ts: _,
1707                     diff: _,
1708                 }| match kind {
1709                    StateUpdateKind::AuditLog(key, ()) => Some(key),
1710                    _ => None,
1711                },
1712            )
1713            .collect();
1714        let mut audit_logs: Vec<_> = audit_logs
1715            .into_iter()
1716            .map(RustType::from_proto)
1717            .map_ok(|key: AuditLogKey| key.event)
1718            .collect::<Result<_, _>>()?;
1719        audit_logs.sort_by(|a, b| a.sortable_id().cmp(&b.sortable_id()));
1720        Ok(audit_logs)
1721    }
1722
1723    #[mz_ore::instrument(level = "debug")]
1724    async fn get_next_id(&mut self, id_type: &str) -> Result<u64, CatalogError> {
1725        self.with_trace(|trace| {
1726            Ok(trace
1727                .into_iter()
1728                .rev()
1729                .filter_map(|(kind, _, _)| match kind {
1730                    StateUpdateKind::IdAllocator(key, value) if key.name == id_type => {
1731                        Some(value.next_id)
1732                    }
1733                    _ => None,
1734                })
1735                .next()
1736                .expect("must exist"))
1737        })
1738        .await
1739    }
1740
1741    #[mz_ore::instrument(level = "debug")]
1742    async fn get_deployment_generation(&mut self) -> Result<u64, CatalogError> {
1743        self.sync_to_current_upper().await?;
1744        Ok(self
1745            .fenceable_token
1746            .token()
1747            .expect("opened catalogs must have a token")
1748            .deploy_generation)
1749    }
1750
1751    #[mz_ore::instrument(level = "debug")]
1752    async fn snapshot(&mut self) -> Result<Snapshot, CatalogError> {
1753        self.with_snapshot(Ok).await
1754    }
1755
1756    #[mz_ore::instrument(level = "debug")]
1757    async fn sync_to_current_updates(
1758        &mut self,
1759    ) -> Result<Vec<memory::objects::StateUpdate>, CatalogError> {
1760        let upper = self.current_upper().await;
1761        self.sync_updates(upper).await
1762    }
1763
1764    #[mz_ore::instrument(level = "debug")]
1765    async fn sync_updates(
1766        &mut self,
1767        target_upper: mz_repr::Timestamp,
1768    ) -> Result<Vec<memory::objects::StateUpdate>, CatalogError> {
1769        self.sync(target_upper).await?;
1770        let mut updates = Vec::new();
1771        while let Some(update) = self.update_applier.updates.front() {
1772            if update.ts >= target_upper {
1773                break;
1774            }
1775
1776            let update = self
1777                .update_applier
1778                .updates
1779                .pop_front()
1780                .expect("peeked above");
1781            updates.push(update);
1782        }
1783        Ok(updates)
1784    }
1785
1786    #[mz_ore::instrument(level = "debug")]
1787    async fn ensure_not_out_of_sync(
1788        &mut self,
1789        target_upper: Timestamp,
1790    ) -> Result<(), CatalogError> {
1791        self.sync(target_upper).await?;
1792        let update_count = self
1793            .update_applier
1794            .updates
1795            .iter()
1796            .take_while(|update| update.ts < target_upper)
1797            .count();
1798        if update_count == 0 {
1799            Ok(())
1800        } else {
1801            Err(DurableCatalogError::CatalogOutOfSync {
1802                update_count,
1803                upper: target_upper,
1804            }
1805            .into())
1806        }
1807    }
1808
1809    async fn current_upper(&mut self) -> Timestamp {
1810        self.current_upper().await
1811    }
1812}
1813
1814#[async_trait]
1815#[allow(mismatched_lifetime_syntaxes)]
1816impl DurableCatalogState for PersistCatalogState {
1817    fn is_read_only(&self) -> bool {
1818        matches!(self.mode, Mode::Readonly)
1819    }
1820
1821    fn is_savepoint(&self) -> bool {
1822        matches!(self.mode, Mode::Savepoint)
1823    }
1824
1825    async fn mark_bootstrap_complete(&mut self) {
1826        self.bootstrap_complete = true;
1827        if matches!(self.mode, Mode::Writable) {
1828            self.since_handle
1829                .upgrade_version()
1830                .await
1831                .expect("invalid usage")
1832        }
1833    }
1834
1835    #[mz_ore::instrument(level = "debug")]
1836    async fn transaction(&mut self) -> Result<Transaction, CatalogError> {
1837        let mut txn = self.transaction_unchecked().await?;
1838        txn.ensure_not_out_of_sync().await?;
1839        Ok(txn)
1840    }
1841
1842    fn transaction_from_snapshot(
1843        &mut self,
1844        snapshot: Snapshot,
1845    ) -> Result<DryRunTransaction, CatalogError> {
1846        let commit_ts = self.upper;
1847        Transaction::new(self, snapshot, commit_ts).map(DryRunTransaction::new)
1848    }
1849
1850    #[mz_ore::instrument(level = "debug")]
1851    async fn allocate_id(
1852        &mut self,
1853        id_type: &str,
1854        amount: u64,
1855        commit_ts: Timestamp,
1856    ) -> Result<Vec<u64>, CatalogError> {
1857        let start = Instant::now();
1858        if amount == 0 {
1859            return Ok(Vec::new());
1860        }
1861        let mut txn = self.transaction_unchecked().await?;
1862        let ids = txn.get_and_increment_id_by(id_type.to_string(), amount)?;
1863        txn.commit_internal(commit_ts).await?;
1864        self.metrics
1865            .allocate_id_seconds
1866            .observe(start.elapsed().as_secs_f64());
1867        Ok(ids)
1868    }
1869
1870    #[mz_ore::instrument(level = "debug")]
1871    async fn commit_transaction(
1872        &mut self,
1873        txn_batch: TransactionBatch,
1874        commit_ts: Timestamp,
1875    ) -> Result<Timestamp, CatalogError> {
1876        async fn commit_transaction_inner(
1877            catalog: &mut PersistCatalogState,
1878            txn_batch: TransactionBatch,
1879            commit_ts: Timestamp,
1880        ) -> Result<Timestamp, CatalogError> {
1881            // If the transaction is empty then we don't error, even in read-only mode.
1882            // This is mostly for legacy reasons (i.e. with enough elbow grease this
1883            // behavior can be changed without breaking any fundamental assumptions).
1884            if catalog.mode == Mode::Readonly {
1885                let updates: Vec<_> = StateUpdate::from_txn_batch(txn_batch).collect();
1886                if !updates.is_empty() {
1887                    let collection_types: Vec<_> = updates
1888                        .iter()
1889                        .filter_map(|u| u.0.collection_type())
1890                        .collect();
1891                    return Err(DurableCatalogError::NotWritable(format!(
1892                        "cannot commit a transaction in a read-only catalog: \
1893                         {} updates across collections: {collection_types:?}",
1894                        updates.len(),
1895                    ))
1896                    .into());
1897                }
1898                return Ok(catalog.upper);
1899            }
1900
1901            // The handle is mutex-protected from transaction creation through commit, so only a
1902            // local programming error can change its cached upper here.
1903            assert_eq!(
1904                catalog.upper, txn_batch.upper,
1905                "the handle was mutated mid-transaction"
1906            );
1907
1908            let updates: Vec<_> = StateUpdate::from_txn_batch(txn_batch).collect();
1909            debug!("committing updates: {updates:?}");
1910
1911            // Empty upper progress does not invalidate the transaction.
1912            let mut commit_ts = max(commit_ts, catalog.upper);
1913
1914            let next_upper = match catalog.mode {
1915                Mode::Writable => loop {
1916                    let updates_applied_before = catalog.updates_applied;
1917                    match catalog.compare_and_append(updates.clone(), commit_ts).await {
1918                        Ok(next_upper) => break next_upper,
1919                        Err(CompareAndAppendError::Fence(e)) => return Err(e.into()),
1920                        Err(CompareAndAppendError::UpperMismatch { actual_upper, .. }) => {
1921                            // The mismatch synchronized the handle. Retry only if it applied no
1922                            // content.
1923                            catalog
1924                                .classify_upper_mismatch(updates_applied_before, actual_upper)?;
1925                            commit_ts = max(commit_ts, catalog.upper);
1926                        }
1927                    }
1928                },
1929                Mode::Savepoint => {
1930                    let updates = updates.into_iter().map(|(kind, diff)| StateUpdate {
1931                        kind,
1932                        ts: commit_ts,
1933                        diff,
1934                    });
1935                    catalog.apply_updates_and_consolidate(updates)?;
1936                    catalog.upper = commit_ts.step_forward();
1937                    catalog.upper
1938                }
1939                Mode::Readonly => unreachable!("handled above"),
1940            };
1941
1942            Ok(next_upper)
1943        }
1944        self.metrics.transaction_commits.inc();
1945        let histogram = self.metrics.transaction_commit_latency_seconds.clone();
1946        commit_transaction_inner(self, txn_batch, commit_ts)
1947            .wall_time()
1948            .observe(histogram)
1949            .await
1950    }
1951
1952    #[mz_ore::instrument(level = "debug")]
1953    async fn advance_upper(&mut self, new_upper: Timestamp) -> Result<(), CatalogError> {
1954        loop {
1955            if self.upper >= new_upper {
1956                // This does not consult Persist. It only revalidates a fence already cached by
1957                // this handle, since a sync can advance `upper` while recording the fence.
1958                self.fenceable_token.validate()?;
1959                return Ok(());
1960            }
1961
1962            match self.mode {
1963                Mode::Writable => {}
1964                Mode::Savepoint => {
1965                    self.upper = new_upper;
1966                    return Ok(());
1967                }
1968                Mode::Readonly => {
1969                    return Err(DurableCatalogError::NotWritable(
1970                        "cannot advance upper of a read-only catalog".into(),
1971                    )
1972                    .into());
1973                }
1974            }
1975
1976            let updates_applied_before = self.updates_applied;
1977            match self.compare_and_append_inner([], new_upper).await {
1978                Ok(()) => {
1979                    self.upper = new_upper;
1980                    // No sync needed since no data was written.
1981                    return Ok(());
1982                }
1983                Err(CompareAndAppendError::Fence(e)) => return Err(e.into()),
1984                Err(CompareAndAppendError::UpperMismatch { actual_upper, .. }) => {
1985                    // The mismatch synchronized the handle. Retry only if it applied no content.
1986                    self.classify_upper_mismatch(updates_applied_before, actual_upper)?;
1987                }
1988            }
1989        }
1990    }
1991
1992    fn shard_id(&self) -> ShardId {
1993        self.shard_id
1994    }
1995}
1996
1997/// Deterministically generate a shard ID for the given `organization_id` and `seed`.
1998pub fn shard_id(organization_id: Uuid, seed: usize) -> ShardId {
1999    let hash = sha2::Sha256::digest(format!("{organization_id}{seed}")).to_vec();
2000    soft_assert_eq_or_log!(hash.len(), 32, "SHA256 returns 32 bytes (256 bits)");
2001    let uuid = Uuid::from_slice(&hash[0..16]).expect("from_slice accepts exactly 16 bytes");
2002    ShardId::from_str(&format!("s{uuid}")).expect("known to be valid")
2003}
2004
2005/// Generates a timestamp for reading from `read_handle` that is as fresh as possible, given
2006/// `upper`.
2007fn as_of(
2008    read_handle: &ReadHandle<SourceData, (), Timestamp, StorageDiff>,
2009    upper: Timestamp,
2010) -> Timestamp {
2011    let since = read_handle.since().clone();
2012    let mut as_of = upper.checked_sub(1).unwrap_or_else(|| {
2013        panic!("catalog persist shard should be initialize, found upper: {upper:?}")
2014    });
2015    // We only downgrade the since after writing, and we always set the since to one less than the
2016    // upper.
2017    soft_assert_or_log!(
2018        since.less_equal(&as_of),
2019        "since={since:?}, as_of={as_of:?}; since must be less than or equal to as_of"
2020    );
2021    // This should be a no-op if the assert above passes, however if it doesn't then we'd like to
2022    // continue with a correct timestamp instead of entering a panic loop.
2023    as_of.advance_by(since.borrow());
2024    as_of
2025}
2026
2027/// Fetch the persist version of the catalog shard, if one exists. A version will not
2028/// exist if we are creating a brand-new environment.
2029async fn fetch_catalog_shard_version(
2030    persist_client: &PersistClient,
2031    catalog_shard_id: ShardId,
2032) -> Option<semver::Version> {
2033    let shard_state = persist_client
2034        .inspect_shard::<Timestamp>(&catalog_shard_id)
2035        .await
2036        .ok()?;
2037    let json_state = serde_json::to_value(shard_state).expect("state serialization error");
2038    let json_version = json_state
2039        .get("applier_version")
2040        .cloned()
2041        .expect("missing applier_version");
2042    let version = serde_json::from_value(json_version).expect("version deserialization error");
2043    Some(version)
2044}
2045
2046/// Generates an iterator of [`StateUpdate`] that contain all updates to the catalog
2047/// state up to, and including, `as_of`.
2048///
2049/// The output is consolidated and sorted by timestamp in ascending order.
2050#[mz_ore::instrument(level = "debug")]
2051async fn snapshot_binary(
2052    read_handle: &mut ReadHandle<SourceData, (), Timestamp, StorageDiff>,
2053    as_of: Timestamp,
2054    metrics: &Arc<Metrics>,
2055) -> impl Iterator<Item = StateUpdate<StateUpdateKindJson>> + DoubleEndedIterator + use<> {
2056    metrics.snapshots_taken.inc();
2057    let histogram = metrics.snapshot_latency_seconds.clone();
2058    snapshot_binary_inner(read_handle, as_of)
2059        .wall_time()
2060        .observe(histogram)
2061        .await
2062}
2063
2064/// Generates an iterator of [`StateUpdate`] that contain all updates to the catalog
2065/// state up to, and including, `as_of`.
2066///
2067/// The output is consolidated and sorted by timestamp in ascending order.
2068#[mz_ore::instrument(level = "debug")]
2069async fn snapshot_binary_inner(
2070    read_handle: &mut ReadHandle<SourceData, (), Timestamp, StorageDiff>,
2071    as_of: Timestamp,
2072) -> impl Iterator<Item = StateUpdate<StateUpdateKindJson>> + DoubleEndedIterator + use<> {
2073    let snapshot = read_handle
2074        .snapshot_and_fetch(Antichain::from_elem(as_of))
2075        .await
2076        .expect("we have advanced the restart_as_of by the since");
2077    soft_assert_no_log!(
2078        snapshot.iter().all(|(_, _, diff)| *diff == 1),
2079        "snapshot_and_fetch guarantees a consolidated result: {snapshot:#?}"
2080    );
2081    snapshot
2082        .into_iter()
2083        .map(Into::<StateUpdate<StateUpdateKindJson>>::into)
2084        .sorted_by(|a, b| Ord::cmp(&b.ts, &a.ts))
2085}
2086
2087/// Convert an [`Antichain<Timestamp>`] to a [`Timestamp`].
2088///
2089/// The correctness of this function relies on [`Timestamp`] being totally ordered and never
2090/// finalizing the catalog shard.
2091pub(crate) fn antichain_to_timestamp(antichain: Antichain<Timestamp>) -> Timestamp {
2092    antichain
2093        .into_option()
2094        .expect("we use a totally ordered time and never finalize the shard")
2095}
2096
2097// Debug methods used by the catalog-debug tool.
2098
2099impl Trace {
2100    /// Generates a [`Trace`] from snapshot.
2101    fn from_snapshot(snapshot: impl IntoIterator<Item = StateUpdate>) -> Trace {
2102        let mut trace = Trace::new();
2103        for StateUpdate { kind, ts, diff } in snapshot {
2104            match kind {
2105                StateUpdateKind::AuditLog(k, v) => trace.audit_log.values.push(((k, v), ts, diff)),
2106                StateUpdateKind::Cluster(k, v) => trace.clusters.values.push(((k, v), ts, diff)),
2107                StateUpdateKind::ClusterReplica(k, v) => {
2108                    trace.cluster_replicas.values.push(((k, v), ts, diff))
2109                }
2110                StateUpdateKind::Comment(k, v) => trace.comments.values.push(((k, v), ts, diff)),
2111                StateUpdateKind::Config(k, v) => trace.configs.values.push(((k, v), ts, diff)),
2112                StateUpdateKind::Database(k, v) => trace.databases.values.push(((k, v), ts, diff)),
2113                StateUpdateKind::DefaultPrivilege(k, v) => {
2114                    trace.default_privileges.values.push(((k, v), ts, diff))
2115                }
2116                StateUpdateKind::FenceToken(_) => {
2117                    // Fence token not included in trace.
2118                }
2119                StateUpdateKind::IdAllocator(k, v) => {
2120                    trace.id_allocator.values.push(((k, v), ts, diff))
2121                }
2122                StateUpdateKind::IntrospectionSourceIndex(k, v) => {
2123                    trace.introspection_sources.values.push(((k, v), ts, diff))
2124                }
2125                StateUpdateKind::Item(k, v) => trace.items.values.push(((k, v), ts, diff)),
2126                StateUpdateKind::NetworkPolicy(k, v) => {
2127                    trace.network_policies.values.push(((k, v), ts, diff))
2128                }
2129                StateUpdateKind::Role(k, v) => trace.roles.values.push(((k, v), ts, diff)),
2130                StateUpdateKind::Schema(k, v) => trace.schemas.values.push(((k, v), ts, diff)),
2131                StateUpdateKind::Setting(k, v) => trace.settings.values.push(((k, v), ts, diff)),
2132                StateUpdateKind::SourceReferences(k, v) => {
2133                    trace.source_references.values.push(((k, v), ts, diff))
2134                }
2135                StateUpdateKind::SystemConfiguration(k, v) => {
2136                    trace.system_configurations.values.push(((k, v), ts, diff))
2137                }
2138                StateUpdateKind::ClusterSystemConfiguration(k, v) => trace
2139                    .cluster_system_configurations
2140                    .values
2141                    .push(((k, v), ts, diff)),
2142                StateUpdateKind::ReplicaSystemConfiguration(k, v) => trace
2143                    .replica_system_configurations
2144                    .values
2145                    .push(((k, v), ts, diff)),
2146                StateUpdateKind::SystemObjectMapping(k, v) => {
2147                    trace.system_object_mappings.values.push(((k, v), ts, diff))
2148                }
2149                StateUpdateKind::SystemPrivilege(k, v) => {
2150                    trace.system_privileges.values.push(((k, v), ts, diff))
2151                }
2152                StateUpdateKind::StorageCollectionMetadata(k, v) => trace
2153                    .storage_collection_metadata
2154                    .values
2155                    .push(((k, v), ts, diff)),
2156                StateUpdateKind::UnfinalizedShard(k, ()) => {
2157                    trace.unfinalized_shards.values.push(((k, ()), ts, diff))
2158                }
2159                StateUpdateKind::TxnWalShard((), v) => {
2160                    trace.txn_wal_shard.values.push((((), v), ts, diff))
2161                }
2162                StateUpdateKind::RoleAuth(k, v) => trace.role_auth.values.push(((k, v), ts, diff)),
2163            }
2164        }
2165        trace
2166    }
2167}
2168
2169impl UnopenedPersistCatalogState {
2170    /// Manually update value of `key` in collection `T` to `value`.
2171    #[mz_ore::instrument]
2172    pub(crate) async fn debug_edit<T: Collection>(
2173        &mut self,
2174        key: T::Key,
2175        value: T::Value,
2176    ) -> Result<Option<T::Value>, CatalogError>
2177    where
2178        T::Key: PartialEq + Eq + Debug + Clone,
2179        T::Value: Debug + Clone,
2180    {
2181        let prev_value = loop {
2182            let key = key.clone();
2183            let value = value.clone();
2184            let snapshot = self.current_snapshot().await?;
2185            let trace = Trace::from_snapshot(snapshot);
2186            let collection_trace = T::collection_trace(trace);
2187            let prev_values: Vec<_> = collection_trace
2188                .values
2189                .into_iter()
2190                .filter(|((k, _), _, diff)| {
2191                    soft_assert_eq_or_log!(*diff, Diff::ONE, "trace is consolidated");
2192                    &key == k
2193                })
2194                .collect();
2195
2196            let prev_value = match &prev_values[..] {
2197                [] => None,
2198                [((_, v), _, _)] => Some(v.clone()),
2199                prev_values => panic!("multiple values found for key {key:?}: {prev_values:?}"),
2200            };
2201
2202            let mut updates: Vec<_> = prev_values
2203                .into_iter()
2204                .map(|((k, v), _, _)| (T::update(k, v), Diff::MINUS_ONE))
2205                .collect();
2206            updates.push((T::update(key, value), Diff::ONE));
2207            // We must fence out all other catalogs, if we haven't already, since we are writing.
2208            match self.fenceable_token.generate_unfenced_token(self.mode)? {
2209                Some((fence_updates, current_fenceable_token)) => {
2210                    updates.extend(fence_updates.clone());
2211                    match self.compare_and_append(updates, self.upper).await {
2212                        Ok(_) => {
2213                            self.fenceable_token = current_fenceable_token;
2214                            break prev_value;
2215                        }
2216                        Err(CompareAndAppendError::Fence(e)) => return Err(e.into()),
2217                        Err(e @ CompareAndAppendError::UpperMismatch { .. }) => {
2218                            warn!("catalog write failed due to upper mismatch, retrying: {e:?}");
2219                            continue;
2220                        }
2221                    }
2222                }
2223                None => {
2224                    self.compare_and_append(updates, self.upper)
2225                        .await
2226                        .map_err(|e| e.unwrap_fence_error())?;
2227                    break prev_value;
2228                }
2229            }
2230        };
2231        Ok(prev_value)
2232    }
2233
2234    /// Manually delete `key` from collection `T`.
2235    #[mz_ore::instrument]
2236    pub(crate) async fn debug_delete<T: Collection>(
2237        &mut self,
2238        key: T::Key,
2239    ) -> Result<(), CatalogError>
2240    where
2241        T::Key: PartialEq + Eq + Debug + Clone,
2242        T::Value: Debug,
2243    {
2244        loop {
2245            let key = key.clone();
2246            let snapshot = self.current_snapshot().await?;
2247            let trace = Trace::from_snapshot(snapshot);
2248            let collection_trace = T::collection_trace(trace);
2249            let mut retractions: Vec<_> = collection_trace
2250                .values
2251                .into_iter()
2252                .filter(|((k, _), _, diff)| {
2253                    soft_assert_eq_or_log!(*diff, Diff::ONE, "trace is consolidated");
2254                    &key == k
2255                })
2256                .map(|((k, v), _, _)| (T::update(k, v), Diff::MINUS_ONE))
2257                .collect();
2258
2259            // We must fence out all other catalogs, if we haven't already, since we are writing.
2260            match self.fenceable_token.generate_unfenced_token(self.mode)? {
2261                Some((fence_updates, current_fenceable_token)) => {
2262                    retractions.extend(fence_updates.clone());
2263                    match self.compare_and_append(retractions, self.upper).await {
2264                        Ok(_) => {
2265                            self.fenceable_token = current_fenceable_token;
2266                            break;
2267                        }
2268                        Err(CompareAndAppendError::Fence(e)) => return Err(e.into()),
2269                        Err(e @ CompareAndAppendError::UpperMismatch { .. }) => {
2270                            warn!("catalog write failed due to upper mismatch, retrying: {e:?}");
2271                            continue;
2272                        }
2273                    }
2274                }
2275                None => {
2276                    self.compare_and_append(retractions, self.upper)
2277                        .await
2278                        .map_err(|e| e.unwrap_fence_error())?;
2279                    break;
2280                }
2281            }
2282        }
2283        Ok(())
2284    }
2285
2286    /// Generates a [`Vec<StateUpdate>`] that contain all updates to the catalog
2287    /// state.
2288    ///
2289    /// The output is consolidated and sorted by timestamp in ascending order and the current upper.
2290    async fn current_snapshot(
2291        &mut self,
2292    ) -> Result<impl IntoIterator<Item = StateUpdate> + '_, CatalogError> {
2293        self.sync_to_current_upper().await?;
2294        self.consolidate();
2295        Ok(self.snapshot.iter().cloned().map(|(kind, ts, diff)| {
2296            let kind = TryIntoStateUpdateKind::try_into(kind).expect("kind decoding error");
2297            StateUpdate { kind, ts, diff }
2298        }))
2299    }
2300}