From 202899f8e3d4a4c1c38b3329f16f6779033dee24 Mon Sep 17 00:00:00 2001
From: Moritz Hoffmann <antiguru@gmail.com>
Date: Mon, 31 Aug 2026 09:46:33 +0200
Subject: [PATCH] persist: make Pending cancel safe

`Pending::block_until_ready` replaced the state with a `Blocking` placeholder
before awaiting the write handle. Dropping the future at that await left the
placeholder in place permanently, and every later `into_result` on that
`Pending` panicked with "block_until_ready cancelled?".

`BatchBuilder` reaches that await in its ordinary write path, in the loop that
bounds outstanding part writes, and `finish` reaches `into_result` through the
run-completion paths. A caller who drove `BatchBuilder::add` inside a
cancellable future and then called `finish`, for example to obtain a `Batch` it
could `delete()`, hit an unconditional panic. Since `install_enhanced_handler`
aborts the process for any panic outside a catch, that took down the whole
process rather than losing the batch.

Await the handle through the borrow instead. The state is never taken out, so
there is nothing to leave behind and the `Blocking` variant disappears along
with its panic. The spawned task keeps running across the cancellation, so a
later `block_until_ready` or `into_result` still yields the value.

Adds `test_pending_survives_cancellation`, which cancels a `block_until_ready`
at its await and then reads the value.

Closes: PER-70

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 src/persist-client/src/internal/merge.rs | 45 +++++++++++++++++++++---
 1 file changed, 40 insertions(+), 5 deletions(-)

diff --git a/src/persist-client/src/internal/merge.rs b/src/persist-client/src/internal/merge.rs
index 88f1d072e9004..0ed52cf8eb799 100644
--- a/src/persist-client/src/internal/merge.rs
+++ b/src/persist-client/src/internal/merge.rs
@@ -163,7 +163,6 @@ impl<T> DerefMut for MergeTree<T> {
 #[derive(Debug)]
 pub enum Pending<T> {
     Writing(JoinHandle<T>),
-    Blocking,
     Finished(T),
 }
 
@@ -176,25 +175,38 @@ impl<T: Send + 'static> Pending<T> {
         matches!(self, Self::Finished(_))
     }
 
+    /// Wait for the value, consuming `self`.
     pub async fn into_result(self) -> T {
         match self {
             Pending::Writing(h) => h.await,
-            Pending::Blocking => panic!("block_until_ready cancelled?"),
             Pending::Finished(t) => t,
         }
     }
 
+    /// Wait for the value and store it, so that later calls resolve without waiting.
+    ///
+    /// Cancel safe: the spawned task keeps running and the handle is retained, so a `Pending`
+    /// whose `block_until_ready` was dropped mid-await still yields its value from a later
+    /// `block_until_ready` or `into_result`.
     pub async fn block_until_ready(&mut self) {
-        let pending = mem::replace(self, Self::Blocking);
-        let value = pending.into_result().await;
+        let value = match self {
+            // Await through the borrow rather than taking the handle out: taking it would have to
+            // leave a placeholder behind, and cancellation at this await point would make that
+            // placeholder permanent.
+            Pending::Writing(handle) => handle.await,
+            Pending::Finished(_) => return,
+        };
         *self = Pending::Finished(value);
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use super::*;
+    use futures::poll;
     use mz_ore::cast::CastLossy;
+    use tokio::sync::oneshot;
+
+    use super::*;
 
     #[mz_ore::test]
     #[cfg_attr(miri, ignore)] // too slow
@@ -260,4 +272,27 @@ mod tests {
             }
         }
     }
+
+    /// A `Pending` whose `block_until_ready` is cancelled mid-await still resolves later. Callers
+    /// reach that await through ordinary cancellable futures, and the value must survive it.
+    #[mz_ore::test(tokio::test)]
+    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `epoll_create1`
+    async fn test_pending_survives_cancellation() {
+        let (tx, rx) = oneshot::channel();
+        let mut pending = Pending::new(mz_ore::task::spawn(|| "pending-test", async move {
+            rx.await.expect("sender not dropped")
+        }));
+
+        // The task cannot complete before the send below, so this poll is guaranteed to suspend at
+        // the await inside `block_until_ready`. Dropping the future there is the cancellation.
+        {
+            let mut blocked = Box::pin(pending.block_until_ready());
+            assert!(poll!(&mut blocked).is_pending());
+        }
+
+        tx.send(42).expect("receiver not dropped");
+        pending.block_until_ready().await;
+        assert!(pending.is_finished());
+        assert_eq!(pending.into_result().await, 42);
+    }
 }