rt: move budget state to context thread-local

tokio/src/runtime/context.rs

@@ -1,3 +1,8 @@
+use crate::runtime::coop;
+
+use std::cell::Cell;
+
+#[cfg(any(feature = "rt", feature = "macros"))]
 use crate::util::rand::{FastRand, RngSeed};
 
 cfg_rt! {
@@ -9,15 +14,23 @@ struct Context {
     /// Handle to the runtime scheduler running on the current thread.
     #[cfg(feature = "rt")]
     scheduler: RefCell<Option<scheduler::Handle>>,
+
+    #[cfg(any(feature = "rt", feature = "macros"))]
     rng: FastRand,
+    /// Tracks the amount of "work" a task may still do before yielding back to
+    /// the sheduler
+    budget: Cell<coop::Budget>,
 }
 
 tokio_thread_local! {
     static CONTEXT: Context = {
         Context {
             #[cfg(feature = "rt")]
             scheduler: RefCell::new(None),
+
+            #[cfg(any(feature = "rt", feature = "macros"))]
             rng: FastRand::new(RngSeed::new()),
+            budget: Cell::new(coop::Budget::unconstrained()),
         }
     }
 }
@@ -27,6 +40,10 @@ pub(crate) fn thread_rng_n(n: u32) -> u32 {
     CONTEXT.with(|ctx| ctx.rng.fastrand_n(n))
 }
 
+pub(super) fn budget<R>(f: impl FnOnce(&Cell<coop::Budget>) -> R) -> R {
+    CONTEXT.with(|ctx| f(&ctx.budget))
+}
+
 cfg_rt! {
     use crate::runtime::TryCurrentError;
 

tokio/src/runtime/coop.rs

@@ -29,11 +29,9 @@
 // other futures. By doing this, you avoid double-counting each iteration of
 // the outer future against the cooperating budget.
 
-use std::cell::Cell;
+use crate::runtime::context;
 
-tokio_thread_local! {
-    static CURRENT: Cell<Budget> = const { Cell::new(Budget::unconstrained()) };
-}
+use std::cell::Cell;
 
 /// Opaque type tracking the amount of "work" a task may still do before
 /// yielding back to the scheduler.
@@ -56,7 +54,7 @@ impl Budget {
     }
 
     /// Returns an unconstrained budget. Operations will not be limited.
-    const fn unconstrained() -> Budget {
+    pub(super) const fn unconstrained() -> Budget {
         Budget(None)
     }
 
@@ -92,7 +90,7 @@ fn with_budget<R>(budget: Budget, f: impl FnOnce() -> R) -> R {
         }
     }
 
-    CURRENT.with(move |cell| {
+    context::budget(|cell| {
         let prev = cell.get();
 
         cell.set(budget);
@@ -105,13 +103,13 @@ fn with_budget<R>(budget: Budget, f: impl FnOnce() -> R) -> R {
 
 #[inline(always)]
 pub(crate) fn has_budget_remaining() -> bool {
-    CURRENT.with(|cell| cell.get().has_remaining())
+    context::budget(|cell| cell.get().has_remaining())
 }
 
 cfg_rt_multi_thread! {
     /// Sets the current task's budget.
     pub(crate) fn set(budget: Budget) {
-        CURRENT.with(|cell| cell.set(budget))
+        context::budget(|cell| cell.set(budget))
     }
 }
 
@@ -120,7 +118,7 @@ cfg_rt! {
     ///
     /// Returns the remaining budget
     pub(crate) fn stop() -> Budget {
-        CURRENT.with(|cell| {
+        context::budget(|cell| {
             let prev = cell.get();
             cell.set(Budget::unconstrained());
             prev
@@ -146,7 +144,7 @@ cfg_coop! {
             // They are both represented as the remembered budget being unconstrained.
             let budget = self.0.get();
             if !budget.is_unconstrained() {
-                CURRENT.with(|cell| {
+                context::budget(|cell| {
                     cell.set(budget);
                 });
             }
@@ -162,12 +160,12 @@ cfg_coop! {
     /// that the budget empties appropriately.
     ///
     /// Note that `RestoreOnPending` restores the budget **as it was before `poll_proceed`**.
-    /// Therefore, if the budget is _further_ adjusted between when `poll_proceed` returns and
+    /// Therefore, if the budget is _fCURRENT.withurther_ adjusted between when `poll_proceed` returns and
     /// `RestRestoreOnPending` is dropped, those adjustments are erased unless the caller indicates
     /// that progress was made.
     #[inline]
     pub(crate) fn poll_proceed(cx: &mut Context<'_>) -> Poll<RestoreOnPending> {
-        CURRENT.with(|cell| {
+        context::budget(|cell| {
             let mut budget = cell.get();
 
             if budget.decrement() {
@@ -211,7 +209,7 @@ mod test {
     use wasm_bindgen_test::wasm_bindgen_test as test;
 
     fn get() -> Budget {
-        CURRENT.with(|cell| cell.get())
+        context::budget(|cell| cell.get())
     }
 
     #[test]

tokio/src/runtime/mod.rs

@@ -177,7 +177,6 @@
 #[macro_use]
 mod tests;
 
-#[cfg(any(feature = "rt", feature = "macros"))]
 pub(crate) mod context;
 
 pub(crate) mod coop;