From 0e5c5d64f5aa4387e7d0fb3c1a69561b170fbfcf Mon Sep 17 00:00:00 2001 From: yanyuxing Date: Tue, 29 Jul 2025 18:09:13 +0800 Subject: [PATCH] future: add adapters of `CancellationToken` for `FutureExt` (#7475) --- tokio-util/Cargo.toml | 4 +- tokio-util/src/future.rs | 138 +++++++++ .../src/future/with_cancellation_token.rs | 79 ++++++ tokio-util/src/lib.rs | 2 + tokio-util/src/sync/cancellation_token.rs | 102 +++++-- tokio-util/src/sync/mod.rs | 1 + tokio-util/src/time/mod.rs | 34 +-- tokio-util/tests/future.rs | 268 ++++++++++++++++++ 8 files changed, 568 insertions(+), 60 deletions(-) create mode 100644 tokio-util/src/future.rs create mode 100644 tokio-util/src/future/with_cancellation_token.rs create mode 100644 tokio-util/tests/future.rs diff --git a/tokio-util/Cargo.toml b/tokio-util/Cargo.toml index 084d123fa..e8e004c00 100644 --- a/tokio-util/Cargo.toml +++ b/tokio-util/Cargo.toml @@ -24,9 +24,9 @@ default = [] full = ["codec", "compat", "io-util", "time", "net", "rt"] net = ["tokio/net"] -compat = ["futures-io",] +compat = ["futures-io"] codec = [] -time = ["tokio/time","slab"] +time = ["tokio/time", "slab"] io = [] io-util = ["io", "tokio/rt", "tokio/io-util"] rt = ["tokio/rt", "tokio/sync", "futures-util", "hashbrown"] diff --git a/tokio-util/src/future.rs b/tokio-util/src/future.rs new file mode 100644 index 000000000..ba4779908 --- /dev/null +++ b/tokio-util/src/future.rs @@ -0,0 +1,138 @@ +//! An extension trait for Futures that provides a variety of convenient adapters. + +mod with_cancellation_token; +use with_cancellation_token::{WithCancellationTokenFuture, WithCancellationTokenFutureOwned}; + +use std::future::Future; + +use crate::sync::CancellationToken; + +/// A trait which contains a variety of convenient adapters and utilities for `Future`s. +pub trait FutureExt: Future { + cfg_time! { + /// A wrapper around [`tokio::time::timeout`], with the advantage that it is easier to write + /// fluent call chains. + /// + /// # Examples + /// + /// ```rust + /// use tokio::{sync::oneshot, time::Duration}; + /// use tokio_util::future::FutureExt; + /// + /// # async fn dox() { + /// let (_tx, rx) = oneshot::channel::<()>(); + /// + /// let res = rx.timeout(Duration::from_millis(10)).await; + /// assert!(res.is_err()); + /// # } + /// ``` + fn timeout(self, timeout: std::time::Duration) -> tokio::time::Timeout + where + Self: Sized, + { + tokio::time::timeout(timeout, self) + } + + /// A wrapper around [`tokio::time::timeout_at`], with the advantage that it is easier to write + /// fluent call chains. + /// + /// # Examples + /// + /// ```rust + /// use tokio::{sync::oneshot, time::{Duration, Instant}}; + /// use tokio_util::future::FutureExt; + /// + /// # async fn dox() { + /// let (_tx, rx) = oneshot::channel::<()>(); + /// let deadline = Instant::now() + Duration::from_millis(10); + /// + /// let res = rx.timeout_at(deadline).await; + /// assert!(res.is_err()); + /// # } + /// ``` + fn timeout_at(self, deadline: tokio::time::Instant) -> tokio::time::Timeout + where + Self: Sized, + { + tokio::time::timeout_at(deadline, self) + } + } + + /// Similar to [`CancellationToken::run_until_cancelled`], + /// but with the advantage that it is easier to write fluent call chains, + /// and biased towards waiting for [`CancellationToken`] to complete. + /// + /// # Fairness + /// + /// Calling this on an already-cancelled token directly returns `None`. + /// For all subsequent polls, in case of concurrent completion and + /// cancellation, this is biased towards the future completion. + /// + /// # Examples + /// + /// ```rust + /// use tokio::sync::oneshot; + /// use tokio_util::future::FutureExt; + /// use tokio_util::sync::CancellationToken; + /// + /// # async fn dox() { + /// let (_tx, rx) = oneshot::channel::<()>(); + /// let token = CancellationToken::new(); + /// let token_clone = token.clone(); + /// tokio::spawn(async move { + /// tokio::time::sleep(std::time::Duration::from_millis(10)).await; + /// token.cancel(); + /// }); + /// assert!(rx.with_cancellation_token(&token_clone).await.is_none()) + /// # } + /// ``` + fn with_cancellation_token( + self, + cancellation_token: &CancellationToken, + ) -> WithCancellationTokenFuture<'_, Self> + where + Self: Sized, + { + WithCancellationTokenFuture::new(cancellation_token, self) + } + + /// Similar to [`CancellationToken::run_until_cancelled_owned`], + /// but with the advantage that it is easier to write fluent call chains, + /// and biased towards waiting for [`CancellationToken`] to complete. + /// + /// # Fairness + /// + /// Calling this on an already-cancelled token directly returns `None`. + /// For all subsequent polls, in case of concurrent completion and + /// cancellation, this is biased towards the future completion. + /// + /// # Examples + /// + /// ```rust + /// use tokio::sync::oneshot; + /// use tokio_util::future::FutureExt; + /// use tokio_util::sync::CancellationToken; + /// + /// # async fn dox() { + /// let (_tx, rx) = oneshot::channel::<()>(); + /// let token = CancellationToken::new(); + /// let token_clone = token.clone(); + /// tokio::spawn(async move { + /// tokio::time::sleep(std::time::Duration::from_millis(10)).await; + /// token.cancel(); + /// }); + /// assert!(rx.with_cancellation_token_owned(token_clone).await.is_none()) + /// # } + /// ``` + fn with_cancellation_token_owned( + self, + cancellation_token: CancellationToken, + ) -> WithCancellationTokenFutureOwned + where + Self: Sized, + { + WithCancellationTokenFutureOwned::new(cancellation_token, self) + } +} + +impl FutureExt for T {} diff --git a/tokio-util/src/future/with_cancellation_token.rs b/tokio-util/src/future/with_cancellation_token.rs new file mode 100644 index 000000000..7b1d0f0bf --- /dev/null +++ b/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> + } +} + +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; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + 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 { + #[pin] + run_until_cancelled: Option> + } +} + +impl WithCancellationTokenFutureOwned { + 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 Future for WithCancellationTokenFutureOwned { + type Output = Option; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.project(); + match this.run_until_cancelled.as_pin_mut() { + Some(fut) => fut.poll(cx), + None => Poll::Ready(None), + } + } +} diff --git a/tokio-util/src/lib.rs b/tokio-util/src/lib.rs index a68a6ba38..677027d6e 100644 --- a/tokio-util/src/lib.rs +++ b/tokio-util/src/lib.rs @@ -59,3 +59,5 @@ pub mod either; pub use bytes; mod util; + +pub mod future; diff --git a/tokio-util/src/sync/cancellation_token.rs b/tokio-util/src/sync/cancellation_token.rs index d785ed424..1b397c2bb 100644 --- a/tokio-util/src/sync/cancellation_token.rs +++ b/tokio-util/src/sync/cancellation_token.rs @@ -281,34 +281,6 @@ impl CancellationToken { 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; - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - 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 - } - } - } - if self.is_cancelled() { None } else { @@ -439,3 +411,77 @@ impl Future for WaitForCancellationFutureOwned { } } } + +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"] + pub(crate) struct RunUntilCancelledFuture<'a, F: Future> { + #[pin] + cancellation: WaitForCancellationFuture<'a>, + #[pin] + future: F, + } +} + +impl<'a, F: Future> RunUntilCancelledFuture<'a, F> { + pub(crate) fn new(cancellation_token: &'a CancellationToken, future: F) -> Self { + Self { + cancellation: cancellation_token.cancelled(), + future, + } + } +} + +impl<'a, F: Future> Future for RunUntilCancelledFuture<'a, F> { + type Output = Option; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + 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 + } + } +} + +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"] + pub(crate) struct RunUntilCancelledFutureOwned { + #[pin] + cancellation: WaitForCancellationFutureOwned, + #[pin] + future: F, + } +} + +impl Future for RunUntilCancelledFutureOwned { + type Output = Option; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + 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 + } + } +} + +impl RunUntilCancelledFutureOwned { + pub(crate) fn new(cancellation_token: CancellationToken, future: F) -> Self { + Self { + cancellation: cancellation_token.cancelled_owned(), + future, + } + } +} diff --git a/tokio-util/src/sync/mod.rs b/tokio-util/src/sync/mod.rs index 28f2fbd2d..c6c97471e 100644 --- a/tokio-util/src/sync/mod.rs +++ b/tokio-util/src/sync/mod.rs @@ -5,6 +5,7 @@ pub use cancellation_token::{ guard::DropGuard, guard_ref::DropGuardRef, CancellationToken, WaitForCancellationFuture, WaitForCancellationFutureOwned, }; +pub(crate) use cancellation_token::{RunUntilCancelledFuture, RunUntilCancelledFutureOwned}; mod mpsc; pub use mpsc::{PollSendError, PollSender}; diff --git a/tokio-util/src/time/mod.rs b/tokio-util/src/time/mod.rs index b4c5a772b..076f4832e 100644 --- a/tokio-util/src/time/mod.rs +++ b/tokio-util/src/time/mod.rs @@ -8,45 +8,19 @@ //! //! This type must be used from within the context of the `Runtime`. -use std::future::Future; use std::time::Duration; -use tokio::time::Timeout; mod wheel; pub mod delay_queue; +// re-export `FutureExt` to avoid breaking change +#[doc(inline)] +pub use crate::future::FutureExt; + #[doc(inline)] pub use delay_queue::DelayQueue; -/// A trait which contains a variety of convenient adapters and utilities for `Future`s. -pub trait FutureExt: Future { - /// A wrapper around [`tokio::time::timeout`], with the advantage that it is easier to write - /// fluent call chains. - /// - /// # Examples - /// - /// ```rust - /// use tokio::{sync::oneshot, time::Duration}; - /// use tokio_util::time::FutureExt; - /// - /// # async fn dox() { - /// let (tx, rx) = oneshot::channel::<()>(); - /// - /// let res = rx.timeout(Duration::from_millis(10)).await; - /// assert!(res.is_err()); - /// # } - /// ``` - fn timeout(self, timeout: Duration) -> Timeout - where - Self: Sized, - { - tokio::time::timeout(timeout, self) - } -} - -impl FutureExt for T {} - // ===== Internal utils ===== enum Round { diff --git a/tokio-util/tests/future.rs b/tokio-util/tests/future.rs new file mode 100644 index 000000000..974a934cd --- /dev/null +++ b/tokio-util/tests/future.rs @@ -0,0 +1,268 @@ +use std::{ + future::{pending, ready, Future}, + task::{Context, Poll}, +}; + +use futures_test::task::new_count_waker; +use tokio::pin; +use tokio_test::{assert_pending, assert_ready_eq}; +use tokio_util::{future::FutureExt, sync::CancellationToken}; + +#[derive(Default)] +struct ReadyOnTheSecondPollFuture { + polled: bool, +} + +impl Future for ReadyOnTheSecondPollFuture { + type Output = (); + + fn poll(mut self: std::pin::Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { + if !self.polled { + self.polled = true; + return Poll::Pending; + } + Poll::Ready(()) + } +} + +#[test] +fn ready_fut_with_cancellation_token_test() { + let (waker, _) = new_count_waker(); + let token = CancellationToken::new(); + + let ready_fut = ready(()); + + let ready_with_token_fut = ready_fut.with_cancellation_token(&token); + + pin!(ready_with_token_fut); + + let res = ready_with_token_fut + .as_mut() + .poll(&mut Context::from_waker(&waker)); + + assert_ready_eq!(res, Some(())); +} + +#[test] +fn pending_fut_with_cancellation_token_test() { + let (waker, _) = new_count_waker(); + let token = CancellationToken::new(); + + let pending_fut = pending::<()>(); + + let pending_with_token_fut = pending_fut.with_cancellation_token(&token); + + pin!(pending_with_token_fut); + + let res = pending_with_token_fut + .as_mut() + .poll(&mut Context::from_waker(&waker)); + + assert_pending!(res); +} + +#[test] +fn ready_fut_with_already_cancelled_token_test() { + let (waker, _) = new_count_waker(); + let token = CancellationToken::new(); + token.cancel(); + + let ready_fut = ready(()); + + let ready_fut_with_token_fut = ready_fut.with_cancellation_token(&token); + + pin!(ready_fut_with_token_fut); + + let res = ready_fut_with_token_fut + .as_mut() + .poll(&mut Context::from_waker(&waker)); + + assert_ready_eq!(res, None); +} + +#[test] +fn pending_fut_with_already_cancelled_token_test() { + let (waker, wake_count) = new_count_waker(); + let token = CancellationToken::new(); + token.cancel(); + + let pending_fut = pending::<()>(); + + let pending_with_token_fut = pending_fut.with_cancellation_token(&token); + + pin!(pending_with_token_fut); + + let res = pending_with_token_fut + .as_mut() + .poll(&mut Context::from_waker(&waker)); + + assert_ready_eq!(res, None); + assert_eq!(wake_count, 0); +} + +#[test] +fn pending_fut_with_token_cancelled_test() { + let (waker, wake_count) = new_count_waker(); + let token = CancellationToken::new(); + + let pending_fut = pending::<()>(); + + let pending_with_token_fut = pending_fut.with_cancellation_token(&token); + + pin!(pending_with_token_fut); + + let res = pending_with_token_fut + .as_mut() + .poll(&mut Context::from_waker(&waker)); + assert_pending!(res); + + token.cancel(); + + let res = pending_with_token_fut + .as_mut() + .poll(&mut Context::from_waker(&waker)); + assert_ready_eq!(res, None); + assert_eq!(wake_count, 1); +} + +#[test] +fn pending_only_on_first_poll_with_cancellation_token_test() { + let (waker, wake_count) = new_count_waker(); + let token = CancellationToken::new(); + let fut = ReadyOnTheSecondPollFuture::default().with_cancellation_token(&token); + pin!(fut); + + // first poll, ReadyOnTheSecondPollFuture returned Pending + let res = fut.as_mut().poll(&mut Context::from_waker(&waker)); + assert_pending!(res); + + token.cancel(); + assert_eq!(wake_count, 1); + + // due to the polling fairness (biased behavior) of `WithCancellationToken` Future, + // subsequent polls are biased toward polling ReadyOnTheSecondPollFuture, + // which results in always returning Ready. + let res = fut.as_mut().poll(&mut Context::from_waker(&waker)); + assert_ready_eq!(res, Some(())); +} + +#[test] +fn ready_fut_with_cancellation_owned_token_test() { + let (waker, _) = new_count_waker(); + let token = CancellationToken::new(); + + let ready_fut = ready(()); + + let ready_with_token_fut = ready_fut.with_cancellation_token_owned(token); + + pin!(ready_with_token_fut); + + let res = ready_with_token_fut + .as_mut() + .poll(&mut Context::from_waker(&waker)); + + assert_ready_eq!(res, Some(())); +} + +#[test] +fn pending_fut_with_cancellation_token_owned_test() { + let (waker, _) = new_count_waker(); + let token = CancellationToken::new(); + + let pending_fut = pending::<()>(); + + let pending_with_token_fut = pending_fut.with_cancellation_token_owned(token); + + pin!(pending_with_token_fut); + + let res = pending_with_token_fut + .as_mut() + .poll(&mut Context::from_waker(&waker)); + + assert_pending!(res); +} + +#[test] +fn ready_fut_with_already_cancelled_token_owned_test() { + let (waker, _) = new_count_waker(); + let token = CancellationToken::new(); + token.cancel(); + + let ready_fut = ready(()); + + let ready_fut_with_token_fut = ready_fut.with_cancellation_token_owned(token); + + pin!(ready_fut_with_token_fut); + + let res = ready_fut_with_token_fut + .as_mut() + .poll(&mut Context::from_waker(&waker)); + + assert_ready_eq!(res, None); +} + +#[test] +fn pending_fut_with_already_cancelled_token_owned_test() { + let (waker, wake_count) = new_count_waker(); + let token = CancellationToken::new(); + token.cancel(); + + let pending_fut = pending::<()>(); + + let pending_with_token_fut = pending_fut.with_cancellation_token_owned(token); + + pin!(pending_with_token_fut); + + let res = pending_with_token_fut + .as_mut() + .poll(&mut Context::from_waker(&waker)); + + assert_ready_eq!(res, None); + assert_eq!(wake_count, 0); +} + +#[test] +fn pending_fut_with_owned_token_cancelled_test() { + let (waker, wake_count) = new_count_waker(); + let token = CancellationToken::new(); + + let pending_fut = pending::<()>(); + + let pending_with_token_fut = pending_fut.with_cancellation_token_owned(token.clone()); + + pin!(pending_with_token_fut); + + let res = pending_with_token_fut + .as_mut() + .poll(&mut Context::from_waker(&waker)); + assert_pending!(res); + + token.cancel(); + + let res = pending_with_token_fut + .as_mut() + .poll(&mut Context::from_waker(&waker)); + assert_ready_eq!(res, None); + assert_eq!(wake_count, 1); +} + +#[test] +fn pending_only_on_first_poll_with_cancellation_token_owned_test() { + let (waker, wake_count) = new_count_waker(); + let token = CancellationToken::new(); + let fut = ReadyOnTheSecondPollFuture::default().with_cancellation_token(&token); + pin!(fut); + + // first poll, ReadyOnTheSecondPollFuture returned Pending + let res = fut.as_mut().poll(&mut Context::from_waker(&waker)); + assert_pending!(res); + + token.cancel(); + assert_eq!(wake_count, 1); + + // due to the polling fairness (biased behavior) of `WithCancellationToken` Future, + // subsequent polls are biased toward polling ReadyOnTheSecondPollFuture, + // which results in always returning Ready. + let res = fut.as_mut().poll(&mut Context::from_waker(&waker)); + assert_ready_eq!(res, Some(())); +}