tokio/runtime/context/
runtime_mt.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
use super::{EnterRuntime, CONTEXT};

/// Returns true if in a runtime context.
pub(crate) fn current_enter_context() -> EnterRuntime {
    CONTEXT.with(|c| c.runtime.get())
}

/// Forces the current "entered" state to be cleared while the closure
/// is executed.
pub(crate) fn exit_runtime<F: FnOnce() -> R, R>(f: F) -> R {
    // Reset in case the closure panics
    struct Reset(EnterRuntime);

    impl Drop for Reset {
        fn drop(&mut self) {
            CONTEXT.with(|c| {
                assert!(
                    !c.runtime.get().is_entered(),
                    "closure claimed permanent executor"
                );
                c.runtime.set(self.0);
            });
        }
    }

    let was = CONTEXT.with(|c| {
        let e = c.runtime.get();
        assert!(e.is_entered(), "asked to exit when not entered");
        c.runtime.set(EnterRuntime::NotEntered);
        e
    });

    let _reset = Reset(was);
    // dropping _reset after f() will reset ENTERED
    f()
}