diff --git a/tokio-util/src/sync/cancellation_token.rs b/tokio-util/src/sync/cancellation_token.rs index c04d8dc19e4..d785ed424de 100644 --- a/tokio-util/src/sync/cancellation_token.rs +++ b/tokio-util/src/sync/cancellation_token.rs @@ -268,6 +268,12 @@ impl CancellationToken { /// unless the [`CancellationToken`] is cancelled. In that case the function returns /// `None` and the future gets dropped. /// + /// # 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. + /// /// # Cancellation safety /// /// This method is only cancel safe if `fut` is cancel safe. @@ -303,11 +309,15 @@ impl CancellationToken { } } - RunUntilCancelledFuture { - cancellation: self.cancelled(), - future: fut, + if self.is_cancelled() { + None + } else { + RunUntilCancelledFuture { + cancellation: self.cancelled(), + future: fut, + } + .await } - .await } /// Runs a future to completion and returns its result wrapped inside of an `Option` @@ -316,6 +326,12 @@ impl CancellationToken { /// /// The function takes self by value and returns a future that owns the token. /// + /// # 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. + /// /// # Cancellation safety /// /// This method is only cancel safe if `fut` is cancel safe. diff --git a/tokio-util/tests/sync_cancellation_token.rs b/tokio-util/tests/sync_cancellation_token.rs index 9332a8f9d02..995890171b7 100644 --- a/tokio-util/tests/sync_cancellation_token.rs +++ b/tokio-util/tests/sync_cancellation_token.rs @@ -492,6 +492,21 @@ fn run_until_cancelled_test() { fut.as_mut().poll(&mut Context::from_waker(&waker)) ); } + + // Do not poll the future when token is already cancelled. + { + let token = CancellationToken::new(); + + let fut = token.run_until_cancelled(async { panic!("fut polled after cancellation") }); + pin!(fut); + + token.cancel(); + + assert_eq!( + Poll::Ready(None), + fut.as_mut().poll(&mut Context::from_waker(&waker)) + ); + } } #[test]