summaryrefslogtreecommitdiffstats
path: root/tokio/src/runtime/context.rs
blob: d3bd698229143f0824a279db8b6853d012640d26 (plain)
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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
//! Thread local runtime context
use crate::runtime::Spawner;
use std::cell::RefCell;

thread_local! {
    static CONTEXT: RefCell<Option<ThreadContext>> = RefCell::new(None)
}

/// ThreadContext makes Runtime context accessible to each Runtime thread.
#[derive(Debug, Clone)]
pub(crate) struct ThreadContext {
    /// Handles to the executor.
    spawner: Spawner,

    /// Handles to the I/O drivers
    io_handle: crate::runtime::io::Handle,

    /// Handles to the time drivers
    time_handle: crate::runtime::time::Handle,

    /// Source of `Instant::now()`
    clock: Option<crate::runtime::time::Clock>,
}

impl Default for ThreadContext {
    fn default() -> Self {
        ThreadContext {
            spawner: Spawner::Shell,
            #[cfg(all(feature = "io-driver", not(loom)))]
            io_handle: None,
            #[cfg(any(not(feature = "io-driver"), loom))]
            io_handle: (),
            #[cfg(all(feature = "time", not(loom)))]
            time_handle: None,
            #[cfg(any(not(feature = "time"), loom))]
            time_handle: (),
            clock: None,
        }
    }
}

impl ThreadContext {
    /// Construct a new [`ThreadContext`]
    ///
    /// [`ThreadContext`]: struct.ThreadContext.html
    pub(crate) fn new(
        spawner: Spawner,
        io_handle: crate::runtime::io::Handle,
        time_handle: crate::runtime::time::Handle,
        clock: Option<crate::runtime::time::Clock>,
    ) -> Self {
        ThreadContext {
            spawner,
            #[cfg(all(feature = "io-driver", not(loom)))]
            io_handle,
            #[cfg(any(not(feature = "io-driver"), loom))]
            io_handle,
            #[cfg(all(feature = "time", not(loom)))]
            time_handle,
            #[cfg(any(not(feature = "time"), loom))]
            time_handle,
            clock,
        }
    }

    /// Clone the current [`ThreadContext`] if one is set, otherwise construct a new [`ThreadContext`].
    ///
    /// [`ThreadContext`]: struct.ThreadContext.html
    #[allow(dead_code)]
    pub(crate) fn clone_current() -> Self {
        CONTEXT.with(|ctx| ctx.borrow().clone().unwrap_or_else(Default::default))
    }

    /// Set this [`ThreadContext`] as the current active [`ThreadContext`].
    ///
    /// [`ThreadContext`]: struct.ThreadContext.html
    pub(crate) fn enter(self) -> ThreadContextDropGuard {
        CONTEXT.with(|ctx| {
            let previous = ctx.borrow_mut().replace(self);
            ThreadContextDropGuard { previous }
        })
    }

    #[cfg(all(feature = "io-driver", not(loom)))]
    pub(crate) fn io_handle() -> crate::runtime::io::Handle {
        CONTEXT.with(|ctx| match *ctx.borrow() {
            Some(ref ctx) => ctx.io_handle.clone(),
            None => None,
        })
    }

    #[cfg(all(feature = "time", not(loom)))]
    pub(crate) fn time_handle() -> crate::runtime::time::Handle {
        CONTEXT.with(|ctx| match *ctx.borrow() {
            Some(ref ctx) => ctx.time_handle.clone(),
            None => None,
        })
    }

    #[cfg(feature = "rt-core")]
    pub(crate) fn spawn_handle() -> Option<Spawner> {
        CONTEXT.with(|ctx| match *ctx.borrow() {
            Some(ref ctx) => Some(ctx.spawner.clone()),
            None => None,
        })
    }

    #[cfg(all(feature = "test-util", feature = "time"))]
    pub(crate) fn clock() -> Option<crate::runtime::time::Clock> {
        CONTEXT.with(
            |ctx| match ctx.borrow().as_ref().map(|ctx| ctx.clock.clone()) {
                Some(Some(clock)) => Some(clock),
                _ => None,
            },
        )
    }
}

/// [`ThreadContextDropGuard`] will replace the `previous` thread context on drop.
///
/// [`ThreadContextDropGuard`]: struct.ThreadContextDropGuard.html
#[derive(Debug)]
pub(crate) struct ThreadContextDropGuard {
    previous: Option<ThreadContext>,
}

impl Drop for ThreadContextDropGuard {
    fn drop(&mut self) {
        CONTEXT.with(|ctx| match self.previous.clone() {
            Some(prev) => ctx.borrow_mut().replace(prev),
            None => ctx.borrow_mut().take(),
        });
    }
}