From 19392ef637551087d2af1ae968aedcaeff521878 Mon Sep 17 00:00:00 2001 From: Carl Lerche Date: Thu, 4 Jun 2026 12:13:02 -0700 Subject: [PATCH] coop: time-based yielding, proof of concept --- tokio/src/runtime/builder.rs | 61 +++++++ tokio/src/runtime/config.rs | 6 + .../runtime/scheduler/current_thread/mod.rs | 10 +- .../runtime/scheduler/multi_thread/worker.rs | 3 +- tokio/src/task/coop/mod.rs | 162 ++++++++++++++---- tokio/tests/coop_time_budget.rs | 75 ++++++++ 6 files changed, 279 insertions(+), 38 deletions(-) create mode 100644 tokio/tests/coop_time_budget.rs diff --git a/tokio/src/runtime/builder.rs b/tokio/src/runtime/builder.rs index a37d47f18..8cf0cedcf 100644 --- a/tokio/src/runtime/builder.rs +++ b/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, + /// 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, diff --git a/tokio/src/runtime/config.rs b/tokio/src/runtime/config.rs index e97482fa9..dd905d8db 100644 --- a/tokio/src/runtime/config.rs +++ b/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, + /// Callback for a worker parking itself pub(crate) before_park: Option, diff --git a/tokio/src/runtime/scheduler/current_thread/mod.rs b/tokio/src/runtime/scheduler/current_thread/mod.rs index f0b072d57..3749e5057 100644 --- a/tokio/src/runtime/scheduler/current_thread/mod.rs +++ b/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(&self, mut core: Box, f: impl FnOnce() -> R) -> (Box, 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; diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index e222edbf9..1e6a8d30c 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/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)] diff --git a/tokio/src/task/coop/mod.rs b/tokio/src/task/coop/mod.rs index 4a520e6f5..6f50d26db 100644 --- a/tokio/src/task/coop/mod.rs +++ b/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); +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(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(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( + time_budget: Option, + 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 { + 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))); + }); + } } diff --git a/tokio/tests/coop_time_budget.rs b/tokio/tests/coop_time_budget.rs new file mode 100644 index 000000000..866187720 --- /dev/null +++ b/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(); + }); +}