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 DerefMut for MergeTree { #[derive(Debug)] pub enum Pending { Writing(JoinHandle), - Blocking, Finished(T), } @@ -176,25 +175,38 @@ impl Pending { 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); + } }