sync: fix CancellationToken failing to cancel the ready futures (#7462)

This patch fixes an issue where the `CancellationToken::run_until_cancelled` never cancels the `Future` that returns `Ready` at the first `poll`.

---------

Co-authored-by: Luca BRUNO <[email protected]>

Author: stepantubanov

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.

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]