std/sync/mpmc/
context.rs

Help
1//! Thread-local channel context.
2
3use super::select::Selected;
4use super::waker::current_thread_id;
5use crate::cell::Cell;
6use crate::ptr;
7use crate::sync::Arc;
8use crate::sync::atomic::{AtomicPtr, AtomicUsize, Ordering};
9use crate::thread::{self, Thread};
10use crate::time::Instant;
11
12/// Thread-local context.
13#[derive(Debug, Clone)]
14pub struct Context {
15    inner: Arc<Inner>,
16}
17
18/// Inner representation of `Context`.
19#[derive(Debug)]
20struct Inner {
21    /// Selected operation.
22    select: AtomicUsize,
23
24    /// A slot into which another thread may store a pointer to its `Packet`.
25    packet: AtomicPtr<()>,
26
27    /// Thread handle.
28    thread: Thread,
29
30    /// Thread id.
31    thread_id: usize,
32}
33
34impl Context {
35    /// Creates a new context for the duration of the closure.
36    #[inline]
37    pub fn with<F, R>(f: F) -> R
38    where
39        F: FnOnce(&Context) -> R,
40    {
41        thread_local! {
42            /// Cached thread-local context.
43            static CONTEXT: Cell<Option<Context>> = Cell::new(Some(Context::new()));
44        }
45
46        let mut f = Some(f);
47        let mut f = |cx: &Context| -> R {
48            let f = f.take().unwrap();
49            f(cx)
50        };
51
52        CONTEXT
53            .try_with(|cell| match cell.take() {
54                None => f(&Context::new()),
55                Some(cx) => {
56                    cx.reset();
57                    let res = f(&cx);
58                    cell.set(Some(cx));
59                    res
60                }
61            })
62            .unwrap_or_else(|_| f(&Context::new()))
63    }
64
65    /// Creates a new `Context`.
66    #[cold]
67    fn new() -> Context {
68        Context {
69            inner: Arc::new(Inner {
70                select: AtomicUsize::new(Selected::Waiting.into()),
71                packet: AtomicPtr::new(ptr::null_mut()),
72                thread: thread::current_or_unnamed(),
73                thread_id: current_thread_id(),
74            }),
75        }
76    }
77
78    /// Resets `select` and `packet`.
79    #[inline]
80    fn reset(&self) {
81        self.inner.select.store(Selected::Waiting.into(), Ordering::Release);
82        self.inner.packet.store(ptr::null_mut(), Ordering::Release);
83    }
84
85    /// Attempts to select an operation.
86    ///
87    /// On failure, the previously selected operation is returned.
88    #[inline]
89    pub fn try_select(&self, select: Selected) -> Result<(), Selected> {
90        self.inner
91            .select
92            .compare_exchange(
93                Selected::Waiting.into(),
94                select.into(),
95                Ordering::AcqRel,
96                Ordering::Acquire,
97            )
98            .map(|_| ())
99            .map_err(|e| e.into())
100    }
101
102    /// Stores a packet.
103    ///
104    /// This method must be called after `try_select` succeeds and there is a packet to provide.
105    #[inline]
106    pub fn store_packet(&self, packet: *mut ()) {
107        if !packet.is_null() {
108            self.inner.packet.store(packet, Ordering::Release);
109        }
110    }
111
112    /// Waits until an operation is selected and returns it.
113    ///
114    /// If the deadline is reached, `Selected::Aborted` will be selected.
115    ///
116    /// # Safety
117    /// This may only be called from the thread this `Context` belongs to.
118    #[inline]
119    pub unsafe fn wait_until(&self, deadline: Option<Instant>) -> Selected {
120        loop {
121            // Check whether an operation has been selected.
122            let sel = Selected::from(self.inner.select.load(Ordering::Acquire));
123            if sel != Selected::Waiting {
124                return sel;
125            }
126
127            // If there's a deadline, park the current thread until the deadline is reached.
128            if let Some(end) = deadline {
129                let now = Instant::now();
130
131                if now < end {
132                    // SAFETY: guaranteed by caller.
133                    unsafe { self.inner.thread.park_timeout(end - now) };
134                } else {
135                    // The deadline has been reached. Try aborting select.
136                    return match self.try_select(Selected::Aborted) {
137                        Ok(()) => Selected::Aborted,
138                        Err(s) => s,
139                    };
140                }
141            } else {
142                // SAFETY: guaranteed by caller.
143                unsafe { self.inner.thread.park() };
144            }
145        }
146    }
147
148    /// Unparks the thread this context belongs to.
149    #[inline]
150    pub fn unpark(&self) {
151        self.inner.thread.unpark();
152    }
153
154    /// Returns the id of the thread this context belongs to.
155    #[inline]
156    pub fn thread_id(&self) -> usize {
157        self.inner.thread_id
158    }
159}
std/sync/mpmc/context.rs

std/sync/mpmc/
context.rs
