context.rs - source
Description: Source of the Rust file `library/std/src/sync/mpmc/context.rs`.

#
std/sync/mpmc/
context.rs

    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::{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: Atomic<usize>,
    23
    24    /// A slot into which another thread may store a pointer to its `Packet`.
    25    packet: Atomic<*mut ()>,
    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}