persist: register schema at write time, not writer open time

This commit changes persist clients to defer calling `register_schema`
on a shard from write handle creation time to the time of the first
append operation. The rationale is that we don't need to enforce a
schema if we are not attempting to write to a shard.

The plan is for 0dt upgrades to make good use of the new, more lenient
behavior. Read-only environments can open write handles with evolved
schemas without having to durably write down the new schemas. This will
allow us to back out of version upgrades without the risk of permanently
poisoning the persist state for lower versions.

Author: teskje

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

@@ -254,6 +254,22 @@ where
             })
     }
 
+    /// See [crate::PersistClient::find_schema].
+    pub fn find_schema(&self, key_schema: &K::Schema, val_schema: &V::Schema) -> Option<SchemaId> {
+        self.state
+            .read_lock(&self.metrics.locks.applier_read_cacheable, |state| {
+                // The common case is that the requested schema is a recent one, so as a minor
+                // optimization, do this search in reverse order.
+                let mut schemas = state.collections.schemas.iter().rev();
+                schemas
+                    .find(|(_, x)| {
+                        K::decode_schema(&x.key) == *key_schema
+                            && V::decode_schema(&x.val) == *val_schema
+                    })
+                    .map(|(id, _)| *id)
+            })
+    }
+
     /// Returns whether the current's state `since` and `upper` are both empty.
     ///
     /// Due to sharing state with other handles, successive reads to this fn or any other may

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

@@ -65,8 +65,6 @@ use crate::{PersistConfig, ShardId, WriterId, cfg};
 /// A key and value `Schema` of data written to a batch or shard.
 #[derive(Debug)]
 pub struct Schemas<K: Codec, V: Codec> {
-    // TODO: Remove the Option once this finishes rolling out and all shards
-    // have a registered schema.
     /// Id under which this schema is registered in the shard's schema registry,
     /// if any.
     pub id: Option<SchemaId>,

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

@@ -700,6 +700,11 @@ where
         self.applier.latest_schema()
     }
 
+    /// See [crate::PersistClient::find_schema].
+    pub fn find_schema(&self, key_schema: &K::Schema, val_schema: &V::Schema) -> Option<SchemaId> {
+        self.applier.find_schema(key_schema, val_schema)
+    }
+
     /// See [crate::PersistClient::compare_and_evolve_schema].
     ///
     /// TODO: Unify this with [Self::register_schema]?

src/persist-client/src/lib.rs

@@ -24,7 +24,7 @@ use differential_dataflow::lattice::Lattice;
 use itertools::Itertools;
 use mz_build_info::{BuildInfo, build_info};
 use mz_dyncfg::ConfigSet;
-use mz_ore::{instrument, soft_assert_or_log};
+use mz_ore::instrument;
 use mz_persist::location::{Blob, Consensus, ExternalError};
 use mz_persist_types::schema::SchemaId;
 use mz_persist_types::{Codec, Codec64, Opaque};
@@ -490,10 +490,6 @@ impl PersistClient {
     ///
     /// Use this to save latency and a bit of persist traffic if you're just
     /// going to immediately drop or expire the [ReadHandle].
-    ///
-    /// The `_schema` parameter is currently unused, but should be an object
-    /// that represents the schema of the data in the shard. This will be required
-    /// in the future.
     #[instrument(level = "debug", fields(shard = %shard_id))]
     pub async fn open_writer<K, V, T, D>(
         &self,
@@ -511,23 +507,11 @@ impl PersistClient {
         let machine = self.make_machine(shard_id, diagnostics.clone()).await?;
         let gc = GarbageCollector::new(machine.clone(), Arc::clone(&self.isolated_runtime));
 
-        // TODO: Because schemas are ordered, as part of the persist schema
-        // changes work, we probably want to build some way to allow persist
-        // users to control the order. For example, maybe a
-        // `PersistClient::compare_and_append_schema(current_schema_id,
-        // next_schema)`. Presumably this would then be passed in to open_writer
-        // instead of us implicitly registering it here.
-        // NB: The overwhelming common case is that this schema is already
-        // registered. In this case, the cmd breaks early and nothing is
-        // written to (or read from) CRDB.
-        let (schema_id, maintenance) = machine.register_schema(&*key_schema, &*val_schema).await;
-        maintenance.start_performing(&machine, &gc);
-        soft_assert_or_log!(
-            schema_id.is_some(),
-            "unable to register schemas {:?} {:?}",
-            key_schema,
-            val_schema,
-        );
+        // We defer registering the schema until write time, to allow opening
+        // write handles in a "read-only" mode where they don't implicitly
+        // modify persist state. But it might already be registered, in which
+        // case we can fetch its ID.
+        let schema_id = machine.find_schema(&*key_schema, &*val_schema);
 
         let writer_id = WriterId::new();
         let schemas = Schemas {
@@ -718,6 +702,26 @@ impl PersistClient {
         Ok(machine.latest_schema())
     }
 
+    /// Returns the ID of the given schema, if known at the current state.
+    pub async fn find_schema<K, V, T, D>(
+        &self,
+        shard_id: ShardId,
+        key_schema: &K::Schema,
+        val_schema: &V::Schema,
+        diagnostics: Diagnostics,
+    ) -> Result<Option<SchemaId>, InvalidUsage<T>>
+    where
+        K: Debug + Codec,
+        V: Debug + Codec,
+        T: Timestamp + Lattice + Codec64 + Sync,
+        D: Semigroup + Codec64 + Send + Sync,
+    {
+        let machine = self
+            .make_machine::<K, V, T, D>(shard_id, diagnostics)
+            .await?;
+        Ok(machine.find_schema(key_schema, val_schema))
+    }
+
     /// Registers a new latest schema for the given shard.
     ///
     /// This new schema must be [backward_compatible] with all previous schemas

src/persist-client/src/schema.rs

@@ -605,7 +605,7 @@ mod tests {
         let schema0 = StringsSchema(vec![false]);
         let schema1 = StringsSchema(vec![false, true]);
 
-        let write0 = client
+        let mut write0 = client
             .open_writer::<Strings, (), u64, i64>(
                 shard_id,
                 Arc::new(schema0.clone()),
@@ -614,6 +614,8 @@ mod tests {
             )
             .await
             .unwrap();
+
+        write0.ensure_schema_registered().await;
         assert_eq!(write0.write_schemas.id.unwrap(), SchemaId(0));
 
         // Not backward compatible (yet... we don't support dropping a column at

src/persist-client/src/write.rs

@@ -221,6 +221,31 @@ where
         self.write_schemas.id
     }
 
+    /// Registers the write schema, if it isn't already registered.
+    ///
+    /// # Panics
+    ///
+    /// This method expects that either the shard doesn't yet have any schema registered, or one of
+    /// the registered schemas is the same as the write schema. If all registered schemas are
+    /// different from the write schema, it panics.
+    pub async fn ensure_schema_registered(&mut self) -> SchemaId {
+        let Schemas { id, key, val } = &self.write_schemas;
+
+        if let Some(id) = id {
+            return *id;
+        }
+
+        let (schema_id, maintenance) = self.machine.register_schema(key, val).await;
+        maintenance.start_performing(&self.machine, &self.gc);
+
+        let Some(schema_id) = schema_id else {
+            panic!("unable to register schemas: {key:?} {val:?}");
+        };
+
+        self.write_schemas.id = Some(schema_id);
+        schema_id
+    }
+
     /// A cached version of the shard-global `upper` frontier.
     ///
     /// This is the most recent upper discovered by this handle. It is
@@ -507,6 +532,16 @@ where
             }
         }
 
+        // If we are writing any data, we require a registered write schema.
+        //
+        // We allow appending empty batches without a registered schema, because
+        // some clients expect to be able to advance the write frontier without
+        // knowing the schema. For example, shard finalization may not always
+        // know the schemas of shards to be finalized.
+        if batches.iter().any(|b| b.batch.len > 0) && self.schema_id().is_none() {
+            self.ensure_schema_registered().await;
+        }
+
         let lower = expected_upper.clone();
         let upper = new_upper;
         let since = Antichain::from_elem(T::minimum());

src/txn-wal/src/txns.rs

@@ -214,7 +214,15 @@ where
     ) -> Result<Tidy, T> {
         let op = &Arc::clone(&self.metrics).register;
         op.run(async {
-            let data_writes = data_writes.into_iter().collect::<Vec<_>>();
+            let mut data_writes = data_writes.into_iter().collect::<Vec<_>>();
+
+            // The txns system requires that all participating data shards have a
+            // schema registered. Importantly, we must register a data shard's
+            // schema _before_ we publish it to the txns shard.
+            for data_write in &mut data_writes {
+                data_write.ensure_schema_registered().await;
+            }
+
             let updates = data_writes
                 .iter()
                 .map(|data_write| {
@@ -285,45 +293,47 @@ where
                 }
             }
             for data_write in data_writes {
-                let new_schema_id = data_write.schema_id();
-
                 // If we already have a write handle for a newer version of a table, don't replace
                 // it! Currently we only support adding columns to tables with a default value, so
                 // the latest/newest schema will always be the most complete.
                 //
-                // TODO(alter_table): Revist when we support dropping columns.
+                // TODO(alter_table): Revisit when we support dropping columns.
                 match self.datas.data_write_for_commit.get(&data_write.shard_id()) {
                     None => {
                         self.datas
                             .data_write_for_commit
                             .insert(data_write.shard_id(), DataWriteCommit(data_write));
                     }
                     Some(previous) => {
-                        match (previous.schema_id(), new_schema_id) {
-                            (Some(previous_id), None) => {
-                                mz_ore::soft_panic_or_log!(
-                                    "tried registering a WriteHandle replacing one with a SchemaId prev_schema_id: {:?} shard_id: {:?}",
-                                    previous_id,
-                                    previous.shard_id(),
-                                );
-                            },
-                            (Some(previous_id), Some(new_id)) if previous_id > new_id => {
-                                mz_ore::soft_panic_or_log!(
-                                    "tried registering a WriteHandle with an older SchemaId prev_schema_id: {:?} new_schema_id: {:?} shard_id: {:?}",
-                                    previous_id,
-                                    new_id,
-                                    previous.shard_id(),
-                                );
-                            },
-                            (previous_schema_id, new_schema_id) => {
-                                if previous_schema_id.is_none() && new_schema_id.is_none() {
-                                    tracing::warn!("replacing WriteHandle without any SchemaIds to reason about");
-                                } else {
-                                    tracing::info!(?previous_schema_id, ?new_schema_id, shard_id = ?previous.shard_id(), "replacing WriteHandle");
-                                }
-                                self.datas.data_write_for_commit.insert(data_write.shard_id(), DataWriteCommit(data_write));
-                            }
+                        let new_schema_id = data_write.schema_id().expect("ensured above");
+
+                        if let Some(prev_schema_id) = previous.schema_id()
+                            && prev_schema_id > new_schema_id
+                        {
+                            mz_ore::soft_panic_or_log!(
+                                "tried registering a WriteHandle with an older SchemaId; \
+                                 prev_schema_id: {} new_schema_id: {} shard_id: {}",
+                                prev_schema_id,
+                                new_schema_id,
+                                previous.shard_id(),
+                            );
+                            continue;
+                        } else if previous.schema_id().is_none() {
+                            mz_ore::soft_panic_or_log!(
+                                "encountered data shard without a schema; shard_id: {}",
+                                previous.shard_id(),
+                            );
                         }
+
+                        tracing::info!(
+                            prev_schema_id = ?previous.schema_id(),
+                            ?new_schema_id,
+                            shard_id = %previous.shard_id(),
+                            "replacing WriteHandle"
+                        );
+                        self.datas
+                            .data_write_for_commit
+                            .insert(data_write.shard_id(), DataWriteCommit(data_write));
                     }
                 }
             }
@@ -756,12 +766,8 @@ where
             .expect("codecs have not changed");
         let (key_schema, val_schema) = match schemas {
             Some((_, key_schema, val_schema)) => (Arc::new(key_schema), Arc::new(val_schema)),
-            // - For new shards we will always have at least one schema
-            //   registered by the time we reach this point, because that
-            //   happens at txn-registration time.
-            // - For pre-existing shards, every txns shard will have had
-            //   open_writer called on it at least once in the previous release,
-            //   so the schema should exist.
+            // We will always have at least one schema registered by the time we reach this point,
+            // because that is ensured at txn-registration time.
             None => unreachable!("data shard {} should have a schema", data_id),
         };
         let wrapped = self