adapter: key migration shard by deploy generation too

This commit adds the deploy generation to the `Key` used in the
migration shard. This is necessary because some of our tests simulate
0dt upgrades by bumping the deploy generation without actually bumping
the version, and those tests get confused when the deploy generation
isn't in the key.

Note that this requires changing the serialization format of the
migration shard key. We choose JSON, to make the format easier to modify
in the future. We still need to be able to parse the older format so we
can clean up old entries.

Author: teskje

src/adapter/src/catalog.rs

@@ -684,6 +684,7 @@ impl Catalog {
                 unsafe_mode: true,
                 all_features: false,
                 build_info,
+                deploy_generation: 0,
                 environment_id: environment_id.unwrap_or_else(EnvironmentId::for_tests),
                 read_only,
                 now,

src/adapter/src/catalog/open.rs

@@ -181,6 +181,8 @@ impl Catalog {
             license_key: config.license_key,
         };
 
+        let deploy_generation = storage.get_deployment_generation().await?;
+
         let mut updates: Vec<_> = storage.sync_to_current_updates().await?;
         assert!(!updates.is_empty(), "initial catalog snapshot is missing");
         let mut txn = storage.transaction().await?;
@@ -481,6 +483,7 @@ impl Catalog {
         // Migrate builtin items.
         let schema_migration_result = builtin_schema_migration::run(
             config.build_info,
+            deploy_generation,
             &mut txn,
             config.builtin_item_migration_config,
         )

src/adapter/src/catalog/open/builtin_schema_migration.rs

@@ -190,6 +190,7 @@ impl Default for MigrationResult {
 /// required migrations, respectively.
 pub(super) async fn run(
     build_info: &BuildInfo,
+    deploy_generation: u64,
     txn: &mut Transaction<'_>,
     config: BuiltinItemMigrationConfig,
 ) -> Result<MigrationResult, Error> {
@@ -223,6 +224,7 @@ pub(super) async fn run(
     let migration = Migration::new(
         durable_version.clone(),
         build_version.clone(),
+        deploy_generation,
         txn,
         builtins,
         config,
@@ -243,6 +245,7 @@ pub(super) async fn run(
 struct Migration<'a, 'b> {
     source_version: Version,
     target_version: Version,
+    deploy_generation: u64,
     txn: &'a mut Transaction<'b>,
     builtins: BTreeMap<SystemObjectDescription, &'static Builtin<NameReference>>,
     object_ids: BTreeMap<SystemObjectDescription, SystemObjectUniqueIdentifier>,
@@ -253,6 +256,7 @@ impl<'a, 'b> Migration<'a, 'b> {
     fn new(
         source_version: Version,
         target_version: Version,
+        deploy_generation: u64,
         txn: &'a mut Transaction<'b>,
         builtins: BTreeMap<SystemObjectDescription, &'static Builtin<NameReference>>,
         config: BuiltinItemMigrationConfig,
@@ -265,6 +269,7 @@ impl<'a, 'b> Migration<'a, 'b> {
         Self {
             source_version,
             target_version,
+            deploy_generation,
             txn,
             builtins,
             object_ids,
@@ -600,7 +605,7 @@ impl<'a, 'b> Migration<'a, 'b> {
     }
 
     /// Try to get or insert replacement shards for the given IDs into the migration shard, at
-    /// `target_version`.
+    /// `target_version` and `deploy_generation`.
     ///
     /// This method looks for existing entries in the migration shards and returns those if they
     /// are present. Otherwise it generates new shard IDs and tries to insert them.
@@ -617,13 +622,13 @@ impl<'a, 'b> Migration<'a, 'b> {
         let upper = persist_write.fetch_recent_upper().await;
         let write_ts = *upper.as_option().expect("migration shard not sealed");
 
-        // Another process might already have done a shard replacement at our version, in which
-        // case we can directly reuse the replacement shards.
-        //
-        // TODO: We should check for deploy generation as well as build version. Storing the deploy
-        // generation as well requires changing the key schema of the migration shard.
+        // Another process might already have done a shard replacement at our version and
+        // generation, in which case we can directly reuse the replacement shards.
         if let Some(read_ts) = write_ts.step_back() {
-            let pred = |key: &migration_shard::Key| key.build_version == self.target_version;
+            let pred = |key: &migration_shard::Key| {
+                key.build_version == self.target_version
+                    && key.deploy_generation == Some(self.deploy_generation)
+            };
             if let Some(entries) = read_migration_shard(persist_read, read_ts, pred).await {
                 let replaced_shards: BTreeMap<_, _> = entries
                     .into_iter()
@@ -659,6 +664,7 @@ impl<'a, 'b> Migration<'a, 'b> {
             let key = migration_shard::Key {
                 global_id,
                 build_version: self.target_version.clone(),
+                deploy_generation: Some(self.deploy_generation),
             };
             updates.push(((key, shard_id), write_ts, 1));
         }
@@ -939,23 +945,41 @@ mod migration_shard {
     use mz_persist_types::columnar::Schema;
     use mz_persist_types::stats::NoneStats;
     use semver::Version;
+    use serde::{Deserialize, Serialize};
 
-    #[derive(Debug, Clone, Eq, Ord, PartialEq, PartialOrd)]
+    #[derive(Debug, Clone, Eq, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
     pub(super) struct Key {
         pub(super) global_id: u64,
         pub(super) build_version: Version,
+        // Versions < 26.0 didn't include the deploy generation. As long as we still might
+        // encounter migration shard entries that don't have it, we need to keep this an `Option`
+        // and keep supporting both key formats.
+        pub(super) deploy_generation: Option<u64>,
     }
 
     impl fmt::Display for Key {
         fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-            write!(f, "{}-{}", self.global_id, self.build_version)
+            if self.deploy_generation.is_some() {
+                // current format
+                let s = serde_json::to_string(self).expect("JSON serializable");
+                f.write_str(&s)
+            } else {
+                // pre-26.0 format
+                write!(f, "{}-{}", self.global_id, self.build_version)
+            }
         }
     }
 
     impl FromStr for Key {
         type Err = String;
 
         fn from_str(s: &str) -> Result<Self, String> {
+            // current format
+            if let Ok(key) = serde_json::from_str(s) {
+                return Ok(key);
+            };
+
+            // pre-26.0 format
             let parts: Vec<_> = s.splitn(2, '-').collect();
             let &[global_id, build_version] = parts.as_slice() else {
                 return Err(format!("invalid Key '{s}'"));
@@ -967,6 +991,7 @@ mod migration_shard {
             Ok(Key {
                 global_id,
                 build_version,
+                deploy_generation: None,
             })
         }
     }
@@ -976,6 +1001,7 @@ mod migration_shard {
             Self {
                 global_id: Default::default(),
                 build_version: Version::new(0, 0, 0),
+                deploy_generation: Some(0),
             }
         }
     }

src/adapter/src/coord.rs

@@ -4140,6 +4140,7 @@ pub fn serve(
                 unsafe_mode,
                 all_features,
                 build_info,
+                deploy_generation: controller_config.deploy_generation,
                 environment_id: environment_id.clone(),
                 read_only: read_only_controllers,
                 now: now.clone(),

src/catalog-debug/src/main.rs

@@ -605,6 +605,7 @@ async fn upgrade_check(
             unsafe_mode: true,
             all_features: false,
             build_info: &BUILD_INFO,
+            deploy_generation: args.deploy_generation.unwrap_or(0),
             environment_id: args.environment_id.clone(),
             read_only,
             now,

src/catalog/src/config.rs

@@ -50,6 +50,8 @@ pub struct StateConfig {
     pub all_features: bool,
     /// Information about this build of Materialize.
     pub build_info: &'static BuildInfo,
+    /// The deploy generation with which the process was started.
+    pub deploy_generation: u64,
     /// A persistent ID associated with the environment.
     pub environment_id: EnvironmentId,
     /// Whether to start Materialize in read-only mode.