tokio/runtime/
context.rs

1use crate::loom::thread::AccessError;
2use crate::task::coop;
3
4use std::cell::Cell;
5
6#[cfg(any(feature = "rt", feature = "macros"))]
7use crate::util::rand::FastRand;
8
9cfg_rt! {
10    mod blocking;
11    pub(crate) use blocking::{disallow_block_in_place, try_enter_blocking_region, BlockingRegionGuard};
12
13    mod current;
14    pub(crate) use current::{with_current, try_set_current, SetCurrentGuard};
15
16    mod runtime;
17    pub(crate) use runtime::{EnterRuntime, enter_runtime};
18
19    mod scoped;
20    use scoped::Scoped;
21
22    use crate::runtime::{scheduler, task::Id};
23
24    use std::task::Waker;
25
26    cfg_taskdump! {
27        use crate::runtime::task::trace;
28    }
29}
30
31cfg_rt_multi_thread! {
32    mod runtime_mt;
33    pub(crate) use runtime_mt::{current_enter_context, exit_runtime};
34}
35
36struct Context {
37    /// Uniquely identifies the current thread
38    #[cfg(feature = "rt")]
39    thread_id: Cell<Option<ThreadId>>,
40
41    /// Handle to the runtime scheduler running on the current thread.
42    #[cfg(feature = "rt")]
43    current: current::HandleCell,
44
45    /// Handle to the scheduler's internal "context"
46    #[cfg(feature = "rt")]
47    scheduler: Scoped<scheduler::Context>,
48
49    #[cfg(feature = "rt")]
50    current_task_id: Cell<Option<Id>>,
51
52    /// Tracks if the current thread is currently driving a runtime.
53    /// Note, that if this is set to "entered", the current scheduler
54    /// handle may not reference the runtime currently executing. This
55    /// is because other runtime handles may be set to current from
56    /// within a runtime.
57    #[cfg(feature = "rt")]
58    runtime: Cell<EnterRuntime>,
59
60    #[cfg(any(feature = "rt", feature = "macros"))]
61    rng: Cell<Option<FastRand>>,
62
63    /// Tracks the amount of "work" a task may still do before yielding back to
64    /// the scheduler
65    budget: Cell<coop::Budget>,
66
67    #[cfg(all(
68        tokio_unstable,
69        feature = "taskdump",
70        feature = "rt",
71        target_os = "linux",
72        any(
73            target_arch = "aarch64",
74            target_arch = "x86",
75            target_arch = "x86_64",
76            target_arch = "s390x"
77        )
78    ))]
79    trace: trace::Context,
80}
81
82tokio_thread_local! {
83    static CONTEXT: Context = const {
84        Context {
85            #[cfg(feature = "rt")]
86            thread_id: Cell::new(None),
87
88            // Tracks the current runtime handle to use when spawning,
89            // accessing drivers, etc...
90            #[cfg(feature = "rt")]
91            current: current::HandleCell::new(),
92
93            // Tracks the current scheduler internal context
94            #[cfg(feature = "rt")]
95            scheduler: Scoped::new(),
96
97            #[cfg(feature = "rt")]
98            current_task_id: Cell::new(None),
99
100            // Tracks if the current thread is currently driving a runtime.
101            // Note, that if this is set to "entered", the current scheduler
102            // handle may not reference the runtime currently executing. This
103            // is because other runtime handles may be set to current from
104            // within a runtime.
105            #[cfg(feature = "rt")]
106            runtime: Cell::new(EnterRuntime::NotEntered),
107
108            #[cfg(any(feature = "rt", feature = "macros"))]
109            rng: Cell::new(None),
110
111            budget: Cell::new(coop::Budget::unconstrained()),
112
113            #[cfg(all(
114                tokio_unstable,
115                feature = "taskdump",
116                feature = "rt",
117                target_os = "linux",
118                any(
119                    target_arch = "aarch64",
120                    target_arch = "x86",
121                    target_arch = "x86_64",
122                    target_arch = "s390x"
123                )
124            ))]
125            trace: trace::Context::new(),
126        }
127    }
128}
129
130#[cfg(any(feature = "macros", feature = "rt"))]
131pub(crate) fn thread_rng_n(n: u32) -> u32 {
132    CONTEXT.with(|ctx| {
133        let mut rng = ctx.rng.get().unwrap_or_else(FastRand::new);
134        let ret = rng.fastrand_n(n);
135        ctx.rng.set(Some(rng));
136        ret
137    })
138}
139
140pub(crate) fn budget<R>(f: impl FnOnce(&Cell<coop::Budget>) -> R) -> Result<R, AccessError> {
141    CONTEXT.try_with(|ctx| f(&ctx.budget))
142}
143
144cfg_rt! {
145    use crate::runtime::ThreadId;
146
147    pub(crate) fn thread_id() -> Result<ThreadId, AccessError> {
148        CONTEXT.try_with(|ctx| {
149            match ctx.thread_id.get() {
150                Some(id) => id,
151                None => {
152                    let id = ThreadId::next();
153                    ctx.thread_id.set(Some(id));
154                    id
155                }
156            }
157        })
158    }
159
160    pub(crate) fn set_current_task_id(id: Option<Id>) -> Option<Id> {
161        CONTEXT.try_with(|ctx| ctx.current_task_id.replace(id)).unwrap_or_default()
162    }
163
164    pub(crate) fn current_task_id() -> Option<Id> {
165        CONTEXT.try_with(|ctx| ctx.current_task_id.get()).unwrap_or_default()
166    }
167
168    pub(crate) fn worker_index() -> Option<usize> {
169        with_scheduler(|ctx| ctx.and_then(|c| c.worker_index()))
170    }
171
172    #[track_caller]
173    pub(crate) fn defer(waker: &Waker) {
174        with_scheduler(|maybe_scheduler| {
175            if let Some(scheduler) = maybe_scheduler {
176                scheduler.defer(waker);
177            } else {
178                // Called from outside of the runtime, immediately wake the
179                // task.
180                waker.wake_by_ref();
181            }
182        });
183    }
184
185    pub(super) fn set_scheduler<R>(v: &scheduler::Context, f: impl FnOnce() -> R) -> R {
186        CONTEXT.with(|c| c.scheduler.set(v, f))
187    }
188
189    #[cfg(all(target_os = "emscripten", not(target_feature = "atomics")))]
190    pub(crate) struct ClearSchedulerGuard(*const scheduler::Context);
191
192    #[cfg(all(target_os = "emscripten", not(target_feature = "atomics")))]
193    impl Drop for ClearSchedulerGuard {
194        fn drop(&mut self) {
195            CONTEXT.with(|c| c.scheduler.inner.set(self.0));
196        }
197    }
198
199    /// Unsets the scheduler context until the guard drops, on unwind too.
200    #[cfg(all(target_os = "emscripten", not(target_feature = "atomics")))]
201    pub(crate) fn clear_scheduler() -> ClearSchedulerGuard {
202        CONTEXT.with(|c| ClearSchedulerGuard(c.scheduler.inner.replace(std::ptr::null())))
203    }
204
205    #[track_caller]
206    pub(super) fn with_scheduler<R>(f: impl FnOnce(Option<&scheduler::Context>) -> R) -> R {
207        let mut f = Some(f);
208        CONTEXT.try_with(|c| {
209            let f = f.take().unwrap();
210            if matches!(c.runtime.get(), EnterRuntime::Entered { .. }) {
211                c.scheduler.with(f)
212            } else {
213                f(None)
214            }
215        })
216            .unwrap_or_else(|_| (f.take().unwrap())(None))
217    }
218
219    cfg_taskdump! {
220        /// SAFETY: Callers of this function must ensure that trace frames always
221        /// form a valid linked list.
222        pub(crate) unsafe fn with_trace<R>(f: impl FnOnce(&trace::Context) -> R) -> Option<R> {
223            CONTEXT.try_with(|c| f(&c.trace)).ok()
224        }
225    }
226}
tokio/runtime/context.rs

tokio/runtime/
context.rs
