Explicit upgrades for special shards

Author: bkirwi

src/adapter/src/catalog/open.rs

@@ -589,7 +589,7 @@ impl Catalog {
                 );
             }
 
-            catalog.storage().await.mark_bootstrap_complete();
+            catalog.storage().await.mark_bootstrap_complete().await;
 
             Ok(OpenCatalogResult {
                 catalog,

src/adapter/src/catalog/open/builtin_item_migration.rs

@@ -357,7 +357,7 @@ async fn migrate_builtin_collections_incompatible(
             shard_id,
             Arc::new(TableKeySchema),
             Arc::new(ShardIdSchema),
-            diagnostics,
+            diagnostics.clone(),
             USE_CRITICAL_SINCE_CATALOG.get(persist_client.dyncfgs()),
         )
         .await
@@ -550,6 +550,10 @@ async fn migrate_builtin_collections_incompatible(
                     error!("Unable to remove old entries from migration shard: {e:?}");
                 }
             }
+            persist_client
+                .upgrade_version::<TableKey, ShardId, Timestamp, StorageDiff>(shard_id, diagnostics)
+                .await
+                .expect("valid usage");
         }
     }
     .boxed();

src/catalog/src/durable.rs

@@ -290,7 +290,7 @@ pub trait DurableCatalogState: ReadOnlyDurableCatalogState {
     fn is_savepoint(&self) -> bool;
 
     /// Marks the bootstrap process as complete.
-    fn mark_bootstrap_complete(&mut self);
+    async fn mark_bootstrap_complete(&mut self);
 
     /// Creates a new durable catalog state transaction.
     async fn transaction(&mut self) -> Result<Transaction, CatalogError>;

src/catalog/src/durable/persist.rs

@@ -1299,6 +1299,7 @@ impl UnopenedPersistCatalogState {
             catalog
                 .increment_catalog_upgrade_shard_version(self.update_applier.organization_id)
                 .await;
+
             let write_handle = catalog
                 .persist_client
                 .open_writer::<SourceData, (), Timestamp, i64>(
@@ -1711,8 +1712,14 @@ impl DurableCatalogState for PersistCatalogState {
         matches!(self.mode, Mode::Savepoint)
     }
 
-    fn mark_bootstrap_complete(&mut self) {
+    async fn mark_bootstrap_complete(&mut self) {
         self.bootstrap_complete = true;
+        if matches!(self.mode, Mode::Writable) {
+            self.since_handle
+                .upgrade_version()
+                .await
+                .expect("invalid usage")
+        }
     }
 
     #[mz_ore::instrument(level = "debug")]

src/catalog/src/expr_cache.rs

@@ -252,6 +252,11 @@ impl ExpressionCache {
             .collect();
         self.durable_cache.try_set_many(&keys_to_remove).await?;
 
+        if remove_prior_versions {
+            // We've purged old versions from the cache; upgrade the backing Persist version as well.
+            self.durable_cache.upgrade_version().await;
+        }
+
         if compact_shard {
             let fuel = EXPRESSION_CACHE_FORCE_COMPACTION_FUEL.handle(dyncfgs);
             let wait = EXPRESSION_CACHE_FORCE_COMPACTION_WAIT.handle(dyncfgs);

src/durable-cache/src/lib.rs

@@ -390,6 +390,15 @@ impl<C: DurableCacheCodec> DurableCache<C> {
         )
         .await
     }
+
+    /// Upgrade the version associated with the backing shard. This should only be done once
+    /// we've durably upgraded to the new version.
+    pub async fn upgrade_version(&self) {
+        self.since_handle
+            .upgrade_version()
+            .await
+            .expect("invalid usage")
+    }
 }
 
 #[cfg(test)]

src/ore/src/process.rs

@@ -41,7 +41,7 @@
 #[macro_export]
 macro_rules! halt {
     ($($arg:expr),* $(,)?) => {{
-        $crate::__private::tracing::warn!("halting process: {}", format!($($arg),*));
+        $crate::__private::tracing::warn!("halting process: {}\n{:?}", format!($($arg),*), ::std::backtrace::Backtrace::capture());
         $crate::process::exit_thread_safe(166);
     }}
 }

src/persist-client/src/critical.rs

@@ -23,6 +23,7 @@ use serde::{Deserialize, Serialize};
 use timely::progress::{Antichain, Timestamp};
 use uuid::Uuid;
 
+use crate::error::InvalidUsage;
 use crate::internal::machine::Machine;
 use crate::internal::state::Since;
 use crate::stats::SnapshotStats;
@@ -324,6 +325,19 @@ where
         }
     }
 
+    /// Upgrade the version associated with this shard, applying any state migrations. This
+    /// is irrevocable and will fence out old versions: it should only be run only once Materialize
+    /// has fully committed to the new version.
+    pub async fn upgrade_version(&self) -> Result<(), InvalidUsage<T>> {
+        match self.machine.upgrade_version().await {
+            Ok(maintenance) => {
+                let () = maintenance.perform(&self.machine, &self.gc).await;
+                Ok(())
+            }
+            Err(version) => Err(InvalidUsage::IncompatibleVersion { version }),
+        }
+    }
+
     // Expiry temporarily removed.
     // If you'd like to stop this handle from holding back the since of the shard,
     // downgrade it to [].

src/persist-client/src/internal/encoding.rs

@@ -206,11 +206,7 @@ pub(crate) fn assert_code_can_read_data(code_version: &Version, data_version: &V
     if !cfg::code_can_read_data(code_version, data_version) {
         // We can't catch halts, so panic in test, so we can get unit test
         // coverage.
-        if cfg!(test) {
-            panic!("code at version {code_version} cannot read data with version {data_version}");
-        } else {
-            halt!("code at version {code_version} cannot read data with version {data_version}");
-        }
+        panic!("code at version {code_version} cannot read data with version {data_version}");
     }
 }