macros: make select! budget-aware

tokio/src/macros/select.rs

@@ -660,6 +660,14 @@ doc! {macro_rules! select {
             let mut futures = &mut futures;
 
             $crate::macros::support::poll_fn(|cx| {
+                // Return `Pending` when the task budget is depleted since budget-aware futures
+                // are going to yield anyway and other futures will not cooperate.
+                if !$crate::macros::support::has_budget_remaining() {
+                    cx.waker().wake_by_ref();
+
+                    return ::std::task::Poll::Pending;
+                }
+
                 // Track if any branch returns pending. If no branch completes
                 // **or** returns pending, this implies that all branches are
                 // disabled.

tokio/src/macros/support.rs

@@ -7,6 +7,14 @@ cfg_macros! {
     pub fn thread_rng_n(n: u32) -> u32 {
         crate::runtime::context::thread_rng_n(n)
     }
+
+    #[doc(hidden)]
+    pub fn has_budget_remaining() -> bool {
+        #[cfg(feature = "rt")]
+        { crate::task::coop::has_budget_remaining() }
+        #[cfg(not(feature = "rt"))]
+        { true }
+    }
 }
 
 pub use std::future::{Future, IntoFuture};

tokio/tests/macros_select.rs

@@ -716,3 +716,27 @@ async fn temporary_lifetime_extension() {
         () = &mut std::future::ready(()) => {},
     }
 }
+
+#[tokio::test]
+async fn select_is_budget_aware() {
+    const BUDGET: usize = 128;
+
+    let task = || {
+        Box::pin(async move {
+            tokio::select! {
+                biased;
+
+                () = tokio::task::coop::consume_budget() => {},
+                () = std::future::ready(()) => {}
+            }
+        })
+    };
+
+    for _ in 0..BUDGET {
+        let poll = futures::poll!(&mut task());
+        assert!(poll.is_ready());
+    }
+
+    let poll = futures::poll!(&mut task());
+    assert!(poll.is_pending());
+}