tracing_core/spin/
once.rs

1use core::cell::UnsafeCell;
2use core::fmt;
3use core::hint::spin_loop;
4use core::sync::atomic::Ordering;
5
6#[cfg(feature = "portable-atomic")]
7use portable_atomic::AtomicUsize;
8
9#[cfg(not(feature = "portable-atomic"))]
10use core::sync::atomic::AtomicUsize;
11
12/// A synchronization primitive which can be used to run a one-time global
13/// initialization. Unlike its std equivalent, this is generalized so that the
14/// closure returns a value and it is stored. Once therefore acts something like
15/// a future, too.
16pub struct Once<T> {
17    state: AtomicUsize,
18    data: UnsafeCell<Option<T>>, // TODO remove option and use mem::uninitialized
19}
20
21impl<T: fmt::Debug> fmt::Debug for Once<T> {
22    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
23        match self.r#try() {
24            Some(s) => write!(f, "Once {{ data: ")
25                .and_then(|()| s.fmt(f))
26                .and_then(|()| write!(f, "}}")),
27            None => write!(f, "Once {{ <uninitialized> }}"),
28        }
29    }
30}
31
32// Same unsafe impls as `std::sync::RwLock`, because this also allows for
33// concurrent reads.
34unsafe impl<T: Send + Sync> Sync for Once<T> {}
35unsafe impl<T: Send> Send for Once<T> {}
36
37// Four states that a Once can be in, encoded into the lower bits of `state` in
38// the Once structure.
39const INCOMPLETE: usize = 0x0;
40const RUNNING: usize = 0x1;
41const COMPLETE: usize = 0x2;
42const PANICKED: usize = 0x3;
43
44use core::hint::unreachable_unchecked as unreachable;
45
46impl<T> Once<T> {
47    /// Initialization constant of `Once`.
48    pub const INIT: Self = Once {
49        state: AtomicUsize::new(INCOMPLETE),
50        data: UnsafeCell::new(None),
51    };
52
53    /// Creates a new `Once` value.
54    pub const fn new() -> Once<T> {
55        Self::INIT
56    }
57
58    fn force_get<'a>(&'a self) -> &'a T {
59        match unsafe { &*self.data.get() }.as_ref() {
60            None => unsafe { unreachable() },
61            Some(p) => p,
62        }
63    }
64
65    /// Performs an initialization routine once and only once. The given closure
66    /// will be executed if this is the first time `call_once` has been called,
67    /// and otherwise the routine will *not* be invoked.
68    ///
69    /// This method will block the calling thread if another initialization
70    /// routine is currently running.
71    ///
72    /// When this function returns, it is guaranteed that some initialization
73    /// has run and completed (it may not be the closure specified). The
74    /// returned pointer will point to the result from the closure that was
75    /// run.
76    pub fn call_once<'a, F>(&'a self, builder: F) -> &'a T
77    where
78        F: FnOnce() -> T,
79    {
80        let mut status = self.state.load(Ordering::SeqCst);
81
82        if status == INCOMPLETE {
83            status = match self.state.compare_exchange(
84                INCOMPLETE,
85                RUNNING,
86                Ordering::SeqCst,
87                Ordering::SeqCst,
88            ) {
89                Ok(status) => {
90                    debug_assert_eq!(
91                        status, INCOMPLETE,
92                        "if compare_exchange succeeded, previous status must be incomplete",
93                    );
94                    // We init
95                    // We use a guard (Finish) to catch panics caused by builder
96                    let mut finish = Finish {
97                        state: &self.state,
98                        panicked: true,
99                    };
100                    unsafe { *self.data.get() = Some(builder()) };
101                    finish.panicked = false;
102
103                    self.state.store(COMPLETE, Ordering::SeqCst);
104
105                    // This next line is strictly an optimization
106                    return self.force_get();
107                }
108                Err(status) => status,
109            }
110        }
111
112        loop {
113            match status {
114                INCOMPLETE => unreachable!(),
115                RUNNING => {
116                    // We spin
117                    spin_loop();
118                    status = self.state.load(Ordering::SeqCst)
119                }
120                PANICKED => panic!("Once has panicked"),
121                COMPLETE => return self.force_get(),
122                _ => unsafe { unreachable() },
123            }
124        }
125    }
126
127    /// Returns a pointer iff the `Once` was previously initialized
128    pub fn r#try<'a>(&'a self) -> Option<&'a T> {
129        match self.state.load(Ordering::SeqCst) {
130            COMPLETE => Some(self.force_get()),
131            _ => None,
132        }
133    }
134
135    /// Like try, but will spin if the `Once` is in the process of being
136    /// initialized
137    pub fn wait<'a>(&'a self) -> Option<&'a T> {
138        loop {
139            match self.state.load(Ordering::SeqCst) {
140                INCOMPLETE => return None,
141
142                RUNNING => {
143                    spin_loop() // We spin
144                }
145                COMPLETE => return Some(self.force_get()),
146                PANICKED => panic!("Once has panicked"),
147                _ => unsafe { unreachable() },
148            }
149        }
150    }
151}
152
153struct Finish<'a> {
154    state: &'a AtomicUsize,
155    panicked: bool,
156}
157
158impl<'a> Drop for Finish<'a> {
159    fn drop(&mut self) {
160        if self.panicked {
161            self.state.store(PANICKED, Ordering::SeqCst);
162        }
163    }
164}
tracing_core/spin/once.rs

tracing_core/spin/
once.rs
