persist: provide WriteHandle::advance_upper

... and use it to enable shard finalization without a known shard
schema.

Author: teskje

src/persist-client/src/cli/admin.rs

@@ -231,7 +231,6 @@ pub async fn run(command: AdminArgs) -> Result<(), anyhow::Error> {
                     Arc::clone(&pubsub_sender),
                 ));
 
-                // We need a PersistClient to open a write handle so we can append an empty batch.
                 let persist_client = PersistClient::new(
                     cfg,
                     blob,
@@ -259,21 +258,7 @@ pub async fn run(command: AdminArgs) -> Result<(), anyhow::Error> {
                         diagnostics,
                     )
                     .await?;
-
-                if !write_handle.upper().is_empty() {
-                    let empty_batch: Vec<(
-                        (crate::cli::inspect::K, crate::cli::inspect::V),
-                        u64,
-                        i64,
-                    )> = vec![];
-                    let lower = write_handle.upper().clone();
-                    let upper = Antichain::new();
-
-                    let result = write_handle.append(empty_batch, lower, upper).await?;
-                    if let Err(err) = result {
-                        anyhow::bail!("failed to force downgrade upper, {err:?}");
-                    }
-                }
+                write_handle.advance_upper(&Antichain::new()).await;
             }
 
             if force_downgrade_since {

src/persist-client/src/lib.rs

@@ -1996,7 +1996,6 @@ mod tests {
     #[mz_persist_proc::test(tokio::test)]
     #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
     async fn finalize_empty_shard(dyncfgs: ConfigUpdates) {
-        const EMPTY: &[(((), ()), u64, i64)] = &[];
         let persist_client = new_test_client(&dyncfgs).await;
 
         let shard_id = ShardId::new();
@@ -2010,11 +2009,7 @@ mod tests {
         // Advance since and upper to empty, which is a pre-requisite for
         // finalization/tombstoning.
         let () = read.downgrade_since(&Antichain::new()).await;
-        let () = write
-            .compare_and_append(EMPTY, Antichain::from_elem(0), Antichain::new())
-            .await
-            .expect("usage should be valid")
-            .expect("upper should match");
+        let () = write.advance_upper(&Antichain::new()).await;
 
         let mut since_handle: SinceHandle<(), (), u64, i64, u64> = persist_client
             .open_critical_since(shard_id, CRITICAL_SINCE, Diagnostics::for_tests())
@@ -2051,7 +2046,6 @@ mod tests {
     #[mz_persist_proc::test(tokio::test)]
     #[cfg_attr(miri, ignore)] // unsupported operation: returning ready events from epoll_wait is not yet implemented
     async fn finalize_shard(dyncfgs: ConfigUpdates) {
-        const EMPTY: &[(((), ()), u64, i64)] = &[];
         const DATA: &[(((), ()), u64, i64)] = &[(((), ()), 0, 1)];
         let persist_client = new_test_client(&dyncfgs).await;
 
@@ -2073,11 +2067,7 @@ mod tests {
         // Advance since and upper to empty, which is a pre-requisite for
         // finalization/tombstoning.
         let () = read.downgrade_since(&Antichain::new()).await;
-        let () = write
-            .compare_and_append(EMPTY, Antichain::from_elem(1), Antichain::new())
-            .await
-            .expect("usage should be valid")
-            .expect("upper should match");
+        let () = write.advance_upper(&Antichain::new()).await;
 
         let mut since_handle: SinceHandle<(), (), u64, i64, u64> = persist_client
             .open_critical_since(shard_id, CRITICAL_SINCE, Diagnostics::for_tests())

src/persist-client/src/write.rs

@@ -281,6 +281,50 @@ where
         &self.upper
     }
 
+    /// Advance the shard's upper by the given frontier.
+    ///
+    /// If the provided `target` is less than or equal to the shard's upper, this is a no-op.
+    ///
+    /// In contrast to the various compare-and-append methods, this method does not require the
+    /// handle's write schema to be registered with the shard. That is, it is fine to use a dummy
+    /// schema when creating a writer just to advance a shard upper.
+    pub async fn advance_upper(&mut self, target: &Antichain<T>) {
+        // We avoid `fetch_recent_upper` here, to avoid a consensus roundtrip if the known upper is
+        // already beyond the target.
+        let mut lower = self.shared_upper().clone();
+
+        while !PartialOrder::less_equal(target, &lower) {
+            let since = Antichain::from_elem(T::minimum());
+            let desc = Description::new(lower.clone(), target.clone(), since);
+            let batch = HollowBatch::empty(desc);
+
+            let heartbeat_timestamp = (self.cfg.now)();
+            let res = self
+                .machine
+                .compare_and_append(
+                    &batch,
+                    &self.writer_id,
+                    &self.debug_state,
+                    heartbeat_timestamp,
+                )
+                .await;
+
+            use CompareAndAppendRes::*;
+            let new_upper = match res {
+                Success(_seq_no, maintenance) => {
+                    maintenance.start_performing(&self.machine, &self.gc, self.compact.as_ref());
+                    batch.desc.upper().clone()
+                }
+                UpperMismatch(_seq_no, actual_upper) => actual_upper,
+                InvalidUsage(_invalid_usage) => unreachable!("batch bounds checked above"),
+                InlineBackpressure => unreachable!("batch was empty"),
+            };
+
+            self.upper.clone_from(&new_upper);
+            lower = new_upper;
+        }
+    }
+
     /// Applies `updates` to this shard and downgrades this handle's upper to
     /// `upper`.
     ///
@@ -532,15 +576,8 @@ 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;
-        }
+        // Before we append any data, we require a registered write schema.
+        self.ensure_schema_registered().await;
 
         let lower = expected_upper.clone();
         let upper = new_upper;

src/storage-client/src/storage_collections.rs

@@ -3248,40 +3248,23 @@ async fn finalize_shards_task<T>(
                     Some(shard_id)
                 } else {
                     debug!(%shard_id, "finalizing shard");
-                        let finalize = || async move {
+                    let finalize = || async move {
                         // TODO: thread the global ID into the shard finalization WAL
                         let diagnostics = Diagnostics::from_purpose("finalizing shards");
 
-                        let schemas = persist_client.latest_schema::<SourceData, (), T, StorageDiff>(shard_id, diagnostics.clone()).await.expect("codecs have not changed");
-                        let (key_schema, val_schema) = match schemas {
-                            Some((_, key_schema, val_schema)) => (key_schema, val_schema),
-                            None => (RelationDesc::empty(), UnitSchema),
-                        };
-
-                        let empty_batch: Vec<((SourceData, ()), T, StorageDiff)> = vec![];
+                        // We only use the writer to advance the upper, so using a dummy schema is
+                        // fine.
                         let mut write_handle: WriteHandle<SourceData, (), T, StorageDiff> =
                             persist_client
                                 .open_writer(
                                     shard_id,
-                                    Arc::new(key_schema),
-                                    Arc::new(val_schema),
+                                    Arc::new(RelationDesc::empty()),
+                                    Arc::new(UnitSchema),
                                     diagnostics,
                                 )
                                 .await
                                 .expect("invalid persist usage");
-
-                        let upper = write_handle.upper();
-
-                        if !upper.is_empty() {
-                            let append = write_handle
-                                .append(empty_batch, upper.clone(), Antichain::new())
-                                .await?;
-
-                            if let Err(e) = append {
-                                warn!(%shard_id, "tried to finalize a shard with an advancing upper: {e:?}");
-                                return Ok(());
-                            }
-                        }
+                        write_handle.advance_upper(&Antichain::new()).await;
                         write_handle.expire().await;
 
                         if force_downgrade_since {
@@ -3317,9 +3300,7 @@ async fn finalize_shards_task<T>(
                                 .compare_and_downgrade_since(&epoch, (&epoch, &new_since))
                                 .await;
                             if let Err(e) = downgrade {
-                                warn!(
-                                    "tried to finalize a shard with an advancing epoch: {e:?}"
-                                );
+                                warn!("tried to finalize a shard with an advancing epoch: {e:?}");
                                 return Ok(());
                             }
                             // Not available now, so finalization is broken.