sync: add CancellationToken::run_until_cancelled (#6618)

tokio-util/src/sync/cancellation_token.rs

@@ -241,6 +241,52 @@ impl CancellationToken {
     pub fn drop_guard(self) -> DropGuard {
         DropGuard { inner: Some(self) }
     }
+
+    /// Runs a future to completion and returns its result wrapped inside of an `Option`
+    /// unless the `CancellationToken` is cancelled. In that case the function returns
+    /// `None` and the future gets dropped.
+    ///
+    /// # Cancel safety
+    ///
+    /// This method is only cancel safe if `fut` is cancel safe.
+    pub async fn run_until_cancelled<F>(&self, fut: F) -> Option<F::Output>
+    where
+        F: Future,
+    {
+        pin_project! {
+            /// A Future that is resolved once the corresponding [`CancellationToken`]
+            /// is cancelled or a given Future gets resolved. It is biased towards the
+            /// Future completion.
+            #[must_use = "futures do nothing unless polled"]
+            struct RunUntilCancelledFuture<'a, F: Future> {
+                #[pin]
+                cancellation: WaitForCancellationFuture<'a>,
+                #[pin]
+                future: F,
+            }
+        }
+
+        impl<'a, F: Future> Future for RunUntilCancelledFuture<'a, F> {
+            type Output = Option<F::Output>;
+
+            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+                let this = self.project();
+                if let Poll::Ready(res) = this.future.poll(cx) {
+                    Poll::Ready(Some(res))
+                } else if this.cancellation.poll(cx).is_ready() {
+                    Poll::Ready(None)
+                } else {
+                    Poll::Pending
+                }
+            }
+        }
+
+        RunUntilCancelledFuture {
+            cancellation: self.cancelled(),
+            future: fut,
+        }
+        .await
+    }
 }
 
 // ===== impl WaitForCancellationFuture =====