coop: time-based yielding, proof of concept

tokio/src/runtime/builder.rs

@@ -124,6 +124,12 @@ pub struct Builder {
     /// How many ticks before yielding to the driver for timer and I/O events?
     pub(super) event_interval: u32,
 
+    /// Optional per-poll wall-clock budget used for cooperative scheduling.
+    ///
+    /// When `Some`, the scheduler uses an elapsed-time strategy to determine
+    /// when a task should yield, rather than the default tick-based strategy.
+    pub(super) coop_time_budget: Option<Duration>,
+
     /// When true, the multi-threade scheduler LIFO slot should not be used.
     ///
     /// This option should only be exposed as unstable.
@@ -326,6 +332,9 @@ impl Builder {
             global_queue_interval: None,
             event_interval,
 
+            // Default: tick-based cooperative budget.
+            coop_time_budget: None,
+
             seed_generator: RngSeedGenerator::new(RngSeed::new()),
 
             #[cfg(tokio_unstable)]
@@ -1222,6 +1231,56 @@ impl Builder {
         self
     }
 
+    /// Switch the cooperative scheduling budget to a time-based strategy.
+    ///
+    /// By default, Tokio uses a tick-based budget: each task may consume up
+    /// to a fixed number of cooperative "ticks" (yield points in
+    /// Tokio-aware leaf futures) per poll before being forced to yield back
+    /// to the scheduler. This works well as long as each tick corresponds to
+    /// a roughly similar amount of work, but it can lead to long polls when
+    /// individual ticks are expensive — for example, when CPU-bound work
+    /// sits between yield points.
+    ///
+    /// Calling this method opts the runtime into an alternative,
+    /// time-based budget: instead of counting ticks, the runtime records the
+    /// time when each task poll begins and considers the budget exhausted
+    /// once the configured wall-clock `duration` has elapsed.
+    ///
+    /// This causes Tokio to call `Instant::now()` at every cooperative yield
+    /// point, which is more expensive than decrementing a counter. Tasks
+    /// that perform many cheap yield points may therefore see throughput
+    /// regressions; the time-based mode is most appropriate when tasks
+    /// perform a significant amount of work between yield points and the
+    /// goal is to bound poll latency.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `duration` is zero.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #[cfg(not(target_family = "wasm"))]
+    /// # {
+    /// use std::time::Duration;
+    /// use tokio::runtime;
+    ///
+    /// let rt = runtime::Builder::new_multi_thread()
+    ///     .coop_time_budget(Duration::from_micros(100))
+    ///     .build()
+    ///     .unwrap();
+    /// # }
+    /// ```
+    #[track_caller]
+    pub fn coop_time_budget(&mut self, duration: Duration) -> &mut Self {
+        assert!(
+            !duration.is_zero(),
+            "coop_time_budget must be greater than zero"
+        );
+        self.coop_time_budget = Some(duration);
+        self
+    }
+
     cfg_unstable! {
         /// Configure how the runtime responds to an unhandled panic on a
         /// spawned task.
@@ -1700,6 +1759,7 @@ impl Builder {
                 after_termination: self.after_termination.clone(),
                 global_queue_interval: self.global_queue_interval,
                 event_interval: self.event_interval,
+                coop_time_budget: self.coop_time_budget,
                 #[cfg(tokio_unstable)]
                 unhandled_panic: self.unhandled_panic.clone(),
                 disable_lifo_slot: self.disable_lifo_slot,
@@ -1886,6 +1946,7 @@ cfg_rt_multi_thread! {
                     after_termination: self.after_termination.clone(),
                     global_queue_interval: self.global_queue_interval,
                     event_interval: self.event_interval,
+                    coop_time_budget: self.coop_time_budget,
                     #[cfg(tokio_unstable)]
                     unhandled_panic: self.unhandled_panic.clone(),
                     disable_lifo_slot: self.disable_lifo_slot,

tokio/src/runtime/config.rs

@@ -12,6 +12,12 @@ pub(crate) struct Config {
     /// How many ticks before yielding to the driver for timer and I/O events?
     pub(crate) event_interval: u32,
 
+    /// When `Some`, the cooperative scheduling budget for each task poll is
+    /// time-based: the task may continue making cooperative yield points until
+    /// the configured wall-clock duration has elapsed, instead of using the
+    /// default tick-based budget.
+    pub(crate) coop_time_budget: Option<std::time::Duration>,
+
     /// Callback for a worker parking itself
     pub(crate) before_park: Option<Callback>,
 

tokio/src/runtime/scheduler/current_thread/mod.rs

@@ -369,8 +369,11 @@ impl Context {
     /// Execute the closure with the given scheduler core stored in the
     /// thread-local context.
     fn run_task<R>(&self, mut core: Box<Core>, f: impl FnOnce() -> R) -> (Box<Core>, R) {
+        let coop_time_budget = self.handle.shared.config.coop_time_budget;
         core.metrics.start_poll();
-        let mut ret = self.enter(core, || crate::task::coop::budget(f));
+        let mut ret = self.enter(core, || {
+            crate::task::coop::run_with_coop_budget(coop_time_budget, f)
+        });
         ret.0.metrics.end_poll();
         ret
     }
@@ -774,8 +777,11 @@ impl CoreGuard<'_> {
                 let handle = &context.handle;
 
                 if handle.reset_woken() {
+                    let coop_time_budget = handle.shared.config.coop_time_budget;
                     let (c, res) = context.enter(core, || {
-                        crate::task::coop::budget(|| future.as_mut().poll(&mut cx))
+                        crate::task::coop::run_with_coop_budget(coop_time_budget, || {
+                            future.as_mut().poll(&mut cx)
+                        })
                     });
 
                     core = c;

tokio/src/runtime/scheduler/multi_thread/worker.rs

@@ -672,7 +672,8 @@ impl Context {
         *self.core.borrow_mut() = Some(core);
 
         // Run the task
-        coop::budget(|| {
+        let coop_time_budget = self.worker.handle.shared.config.coop_time_budget;
+        coop::run_with_coop_budget(coop_time_budget, || {
             // Unlike the poll time above, poll start callback is attached to the task id,
             // so it is tightly associated with the actual poll invocation.
             #[cfg(tokio_unstable)]

tokio/src/task/coop/mod.rs

@@ -91,10 +91,25 @@ cfg_rt! {
 
 use crate::runtime::context;
 
+use std::time::{Duration, Instant};
+
 /// Opaque type tracking the amount of "work" a task may still do before
 /// yielding back to the scheduler.
 #[derive(Debug, Copy, Clone)]
-pub(crate) struct Budget(Option<u8>);
+pub(crate) struct Budget(BudgetInner);
+
+#[derive(Debug, Copy, Clone)]
+enum BudgetInner {
+    /// No budget constraint. Operations will not be limited.
+    Unconstrained,
+    /// Tick-based budget. Each call to [`Budget::decrement`] consumes one tick;
+    /// the budget is exhausted when the counter reaches zero.
+    Ticks(u8),
+    /// Time-based budget. The contained `Instant` is the deadline at which
+    /// the budget is considered exhausted. Each call to [`Budget::decrement`]
+    /// reads `Instant::now()` and compares it against the deadline.
+    Time(Instant),
+}
 
 pub(crate) struct BudgetDecrement {
     success: bool,
@@ -102,7 +117,7 @@ pub(crate) struct BudgetDecrement {
 }
 
 impl Budget {
-    /// Budget assigned to a task on each poll.
+    /// Default tick budget assigned to a task on each poll.
     ///
     /// The value itself is chosen somewhat arbitrarily. It needs to be high
     /// enough to amortize wakeup and scheduling costs, but low enough that we
@@ -112,17 +127,27 @@ impl Budget {
     ///
     /// Note that as more yield points are added in the ecosystem, this value
     /// will probably also have to be raised.
-    const fn initial() -> Budget {
-        Budget(Some(128))
+    const fn initial_ticks() -> Budget {
+        Budget(BudgetInner::Ticks(128))
+    }
+
+    /// Returns a time-based budget that is considered exhausted once
+    /// `Instant::now()` reaches or exceeds the given deadline.
+    fn initial_time(duration: Duration) -> Budget {
+        Budget(BudgetInner::Time(Instant::now() + duration))
     }
 
     /// Returns an unconstrained budget. Operations will not be limited.
     pub(crate) const fn unconstrained() -> Budget {
-        Budget(None)
+        Budget(BudgetInner::Unconstrained)
     }
 
     fn has_remaining(self) -> bool {
-        self.0.map_or(true, |budget| budget > 0)
+        match self.0 {
+            BudgetInner::Unconstrained => true,
+            BudgetInner::Ticks(budget) => budget > 0,
+            BudgetInner::Time(deadline) => Instant::now() < deadline,
+        }
     }
 }
 
@@ -130,7 +155,31 @@ impl Budget {
 /// returns, the budget is reset to the value prior to calling the function.
 #[inline(always)]
 pub(crate) fn budget<R>(f: impl FnOnce() -> R) -> R {
-    with_budget(Budget::initial(), f)
+    with_budget(Budget::initial_ticks(), f)
+}
+
+/// Runs the given closure with a time-based cooperative task budget. The task
+/// is considered out of budget once the elapsed wall-clock duration since
+/// entering this function exceeds `duration`. When the function returns, the
+/// budget is reset to the value prior to calling the function.
+#[inline(always)]
+pub(crate) fn budget_time<R>(duration: Duration, f: impl FnOnce() -> R) -> R {
+    with_budget(Budget::initial_time(duration), f)
+}
+
+/// Runs the given closure with a cooperative task budget chosen by the runtime
+/// configuration. If `time_budget` is `Some(d)`, the closure runs with a
+/// time-based budget of duration `d`; otherwise, it runs with the default
+/// tick-based budget.
+#[inline(always)]
+pub(crate) fn run_with_coop_budget<R>(
+    time_budget: Option<Duration>,
+    f: impl FnOnce() -> R,
+) -> R {
+    match time_budget {
+        Some(d) => budget_time(d, f),
+        None => budget(f),
+    }
 }
 
 /// Runs the given closure with an unconstrained task budget. When the function returns, the budget
@@ -411,23 +460,33 @@ cfg_coop! {
         /// Decrements the budget. Returns `true` if successful. Decrementing fails
         /// when there is not enough remaining budget.
         fn decrement(&mut self) -> BudgetDecrement {
-            if let Some(num) = &mut self.0 {
-                if *num > 0 {
-                    *num -= 1;
-
-                    let hit_zero = *num == 0;
-
-                    BudgetDecrement { success: true, hit_zero }
-                } else {
-                    BudgetDecrement { success: false, hit_zero: false }
+            match &mut self.0 {
+                BudgetInner::Unconstrained => {
+                    BudgetDecrement { success: true, hit_zero: false }
+                }
+                BudgetInner::Ticks(num) => {
+                    if *num > 0 {
+                        *num -= 1;
+
+                        let hit_zero = *num == 0;
+
+                        BudgetDecrement { success: true, hit_zero }
+                    } else {
+                        BudgetDecrement { success: false, hit_zero: false }
+                    }
+                }
+                BudgetInner::Time(deadline) => {
+                    if Instant::now() < *deadline {
+                        BudgetDecrement { success: true, hit_zero: false }
+                    } else {
+                        BudgetDecrement { success: false, hit_zero: false }
+                    }
                 }
-            } else {
-                BudgetDecrement { success: true, hit_zero: false }
             }
         }
 
         fn is_unconstrained(self) -> bool {
-            self.0.is_none()
+            matches!(self.0, BudgetInner::Unconstrained)
         }
     }
 
@@ -503,58 +562,69 @@ mod test {
         context::budget(|cell| cell.get()).unwrap_or(Budget::unconstrained())
     }
 
+    fn ticks(b: Budget) -> Option<u8> {
+        match b.0 {
+            BudgetInner::Ticks(n) => Some(n),
+            _ => None,
+        }
+    }
+
+    fn initial_ticks() -> u8 {
+        ticks(Budget::initial_ticks()).unwrap()
+    }
+
     #[test]
     fn budgeting() {
         use std::future::poll_fn;
         use tokio_test::*;
 
-        assert!(get().0.is_none());
+        assert!(get().is_unconstrained());
 
         let coop = assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
 
-        assert!(get().0.is_none());
+        assert!(get().is_unconstrained());
         drop(coop);
-        assert!(get().0.is_none());
+        assert!(get().is_unconstrained());
 
         budget(|| {
-            assert_eq!(get().0, Budget::initial().0);
+            assert_eq!(ticks(get()), Some(initial_ticks()));
 
             let coop = assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
-            assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 1);
+            assert_eq!(ticks(get()).unwrap(), initial_ticks() - 1);
             drop(coop);
             // we didn't make progress
-            assert_eq!(get().0, Budget::initial().0);
+            assert_eq!(ticks(get()), Some(initial_ticks()));
 
             let coop = assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
-            assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 1);
+            assert_eq!(ticks(get()).unwrap(), initial_ticks() - 1);
             coop.made_progress();
             drop(coop);
             // we _did_ make progress
-            assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 1);
+            assert_eq!(ticks(get()).unwrap(), initial_ticks() - 1);
 
             let coop = assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
-            assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 2);
+            assert_eq!(ticks(get()).unwrap(), initial_ticks() - 2);
             coop.made_progress();
             drop(coop);
-            assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 2);
+            assert_eq!(ticks(get()).unwrap(), initial_ticks() - 2);
 
             budget(|| {
-                assert_eq!(get().0, Budget::initial().0);
+                assert_eq!(ticks(get()), Some(initial_ticks()));
 
                 let coop = assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
-                assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 1);
+                assert_eq!(ticks(get()).unwrap(), initial_ticks() - 1);
                 coop.made_progress();
                 drop(coop);
-                assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 1);
+                assert_eq!(ticks(get()).unwrap(), initial_ticks() - 1);
             });
 
-            assert_eq!(get().0.unwrap(), Budget::initial().0.unwrap() - 2);
+            assert_eq!(ticks(get()).unwrap(), initial_ticks() - 2);
         });
 
-        assert!(get().0.is_none());
+        assert!(get().is_unconstrained());
 
         budget(|| {
-            let n = get().0.unwrap();
+            let n = ticks(get()).unwrap();
 
             for _ in 0..n {
                 let coop = assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
@@ -570,4 +640,26 @@ mod test {
             assert_pending!(task.poll());
         });
     }
+
+    #[test]
+    fn time_budget_exhausted() {
+        use std::thread;
+        use tokio_test::*;
+
+        // A small but non-trivial duration; we'll sleep past it to force exhaustion.
+        let dur = Duration::from_millis(10);
+
+        budget_time(dur, || {
+            // First poll_proceed should succeed (time has not yet elapsed).
+            let coop = assert_ready!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
+            coop.made_progress();
+
+            // Sleep past the deadline.
+            thread::sleep(dur + Duration::from_millis(5));
+
+            // Now the budget should be exhausted.
+            assert!(!has_budget_remaining());
+            let _ = assert_pending!(task::spawn(()).enter(|cx, _| poll_proceed(cx)));
+        });
+    }
 }

tokio/tests/coop_time_budget.rs

@@ -0,0 +1,75 @@
+#![warn(rust_2018_idioms)]
+#![cfg(feature = "full")]
+
+use std::time::Duration;
+use tokio::runtime::Builder;
+use tokio::task::coop::{consume_budget, has_budget_remaining};
+
+/// With time-based coop budgeting, a task should be permitted to consume many
+/// more cooperative yield points per poll than the default tick budget (128),
+/// provided each yield point is cheap and the wall-clock budget has not yet
+/// elapsed.
+#[test]
+fn time_budget_allows_more_than_tick_budget() {
+    let rt = Builder::new_current_thread()
+        .enable_all()
+        .coop_time_budget(Duration::from_secs(1))
+        .build()
+        .unwrap();
+
+    rt.block_on(async {
+        // The tick-based budget would exhaust after 128 calls. With a 1s
+        // time-based budget the loop below should easily exceed that without
+        // ever observing exhaustion.
+        for _ in 0..1_000 {
+            assert!(has_budget_remaining());
+            consume_budget().await;
+        }
+    });
+}
+
+/// With time-based coop budgeting, once the configured duration elapses, the
+/// budget should be reported as exhausted.
+#[test]
+fn time_budget_exhausts_after_duration() {
+    let rt = Builder::new_current_thread()
+        .enable_all()
+        .coop_time_budget(Duration::from_millis(10))
+        .build()
+        .unwrap();
+
+    rt.block_on(async {
+        assert!(has_budget_remaining());
+
+        // Burn through the time budget without yielding back to the runtime.
+        std::thread::sleep(Duration::from_millis(20));
+
+        assert!(!has_budget_remaining());
+    });
+}
+
+#[test]
+#[should_panic]
+fn time_budget_zero_panics() {
+    let _ = Builder::new_current_thread().coop_time_budget(Duration::ZERO);
+}
+
+#[test]
+fn time_budget_works_on_multi_thread() {
+    let rt = Builder::new_multi_thread()
+        .worker_threads(1)
+        .enable_all()
+        .coop_time_budget(Duration::from_secs(1))
+        .build()
+        .unwrap();
+
+    rt.block_on(async {
+        let handle = tokio::spawn(async {
+            for _ in 0..1_000 {
+                assert!(has_budget_remaining());
+                consume_budget().await;
+            }
+        });
+        handle.await.unwrap();
+    });
+}