future: add adapters of CancellationToken for FutureExt (#7475)

tokio-util/src/future/with_cancellation_token.rs

@@ -0,0 +1,79 @@
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+use pin_project_lite::pin_project;
+
+use crate::sync::{CancellationToken, RunUntilCancelledFuture, RunUntilCancelledFutureOwned};
+
+pin_project! {
+    /// A [`Future`] that is resolved once the corresponding [`CancellationToken`]
+    /// is cancelled or a given [`Future`] gets resolved.
+    ///
+    /// This future is immediately resolved if the corresponding [`CancellationToken`]
+    /// is already cancelled, otherwise, in case of concurrent completion and
+    /// cancellation, this is biased towards the future completion.
+    #[must_use = "futures do nothing unless polled"]
+    pub struct WithCancellationTokenFuture<'a, F: Future> {
+        #[pin]
+        run_until_cancelled: Option<RunUntilCancelledFuture<'a, F>>
+    }
+}
+
+impl<'a, F: Future> WithCancellationTokenFuture<'a, F> {
+    pub(crate) fn new(cancellation_token: &'a CancellationToken, future: F) -> Self {
+        Self {
+            run_until_cancelled: (!cancellation_token.is_cancelled())
+                .then(|| RunUntilCancelledFuture::new(cancellation_token, future)),
+        }
+    }
+}
+
+impl<'a, F: Future> Future for WithCancellationTokenFuture<'a, F> {
+    type Output = Option<F::Output>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        match this.run_until_cancelled.as_pin_mut() {
+            Some(fut) => fut.poll(cx),
+            None => Poll::Ready(None),
+        }
+    }
+}
+
+pin_project! {
+    /// A [`Future`] that is resolved once the corresponding [`CancellationToken`]
+    /// is cancelled or a given [`Future`] gets resolved.
+    ///
+    /// This future is immediately resolved if the corresponding [`CancellationToken`]
+    /// is already cancelled, otherwise, in case of concurrent completion and
+    /// cancellation, this is biased towards the future completion.
+    #[must_use = "futures do nothing unless polled"]
+    pub struct WithCancellationTokenFutureOwned<F: Future> {
+        #[pin]
+        run_until_cancelled: Option<RunUntilCancelledFutureOwned<F>>
+    }
+}
+
+impl<F: Future> WithCancellationTokenFutureOwned<F> {
+    pub(crate) fn new(cancellation_token: CancellationToken, future: F) -> Self {
+        Self {
+            run_until_cancelled: (!cancellation_token.is_cancelled())
+                .then(|| RunUntilCancelledFutureOwned::new(cancellation_token, future)),
+        }
+    }
+}
+
+impl<F: Future> Future for WithCancellationTokenFutureOwned<F> {
+    type Output = Option<F::Output>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        match this.run_until_cancelled.as_pin_mut() {
+            Some(fut) => fut.poll(cx),
+            None => Poll::Ready(None),
+        }
+    }
+}