tokio/runtime/context/
runtime_mt.rs

1use super::{EnterRuntime, CONTEXT};
2
3/// Returns true if in a runtime context.
4pub(crate) fn current_enter_context() -> EnterRuntime {
5    CONTEXT.with(|c| c.runtime.get())
6}
7
8/// Forces the current "entered" state to be cleared while the closure
9/// is executed.
10pub(crate) fn exit_runtime<F: FnOnce() -> R, R>(f: F) -> R {
11    // Reset in case the closure panics
12    struct Reset(EnterRuntime);
13
14    impl Drop for Reset {
15        fn drop(&mut self) {
16            CONTEXT.with(|c| {
17                assert!(
18                    !c.runtime.get().is_entered(),
19                    "closure claimed permanent executor"
20                );
21                c.runtime.set(self.0);
22            });
23        }
24    }
25
26    let was = CONTEXT.with(|c| {
27        let e = c.runtime.get();
28        assert!(e.is_entered(), "asked to exit when not entered");
29        c.runtime.set(EnterRuntime::NotEntered);
30        e
31    });
32
33    let _reset = Reset(was);
34    // dropping _reset after f() will reset ENTERED
35    f()
36}