1 //! Synchronization primitives for one-time evaluation.
2 
3 use core::{
4     cell::UnsafeCell,
5     mem::MaybeUninit,
6     sync::atomic::{AtomicUsize, Ordering},
7     fmt,
8 };
9 
10 /// A primitive that provides lazy one-time initialization.
11 ///
12 /// Unlike its `std::sync` equivalent, this is generalized such that the closure returns a
13 /// value to be stored by the [`Once`] (`std::sync::Once` can be trivially emulated with
14 /// `Once<()>`).
15 ///
16 /// Because [`Once::new`] is `const`, this primitive may be used to safely initialize statics.
17 ///
18 /// # Examples
19 ///
20 /// ```
21 /// use spin;
22 ///
23 /// static START: spin::Once<()> = spin::Once::new();
24 ///
25 /// START.call_once(|| {
26 ///     // run initialization here
27 /// });
28 /// ```
29 pub struct Once<T> {
30     state: AtomicUsize,
31     data: UnsafeCell<MaybeUninit<T>>,
32 }
33 
34 impl<T: fmt::Debug> fmt::Debug for Once<T> {
fmt(&self, f: &mut fmt::Formatter) -> fmt::Result35     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
36         match self.get() {
37             Some(s) => write!(f, "Once {{ data: ")
38 				.and_then(|()| s.fmt(f))
39 				.and_then(|()| write!(f, "}}")),
40             None => write!(f, "Once {{ <uninitialized> }}")
41         }
42     }
43 }
44 
45 // Same unsafe impls as `std::sync::RwLock`, because this also allows for
46 // concurrent reads.
47 unsafe impl<T: Send + Sync> Sync for Once<T> {}
48 unsafe impl<T: Send> Send for Once<T> {}
49 
50 // Four states that a Once can be in, encoded into the lower bits of `state` in
51 // the Once structure.
52 const INCOMPLETE: usize = 0x0;
53 const RUNNING: usize = 0x1;
54 const COMPLETE: usize = 0x2;
55 const PANICKED: usize = 0x3;
56 
57 use core::hint::unreachable_unchecked as unreachable;
58 
59 impl<T> Once<T> {
60     /// Initialization constant of [`Once`].
61     #[allow(clippy::declare_interior_mutable_const)]
62     pub const INIT: Self = Self {
63         state: AtomicUsize::new(INCOMPLETE),
64         data: UnsafeCell::new(MaybeUninit::uninit()),
65     };
66 
67     /// Creates a new [`Once`].
new() -> Once<T>68     pub const fn new() -> Once<T> {
69         Self::INIT
70     }
71 
72     /// Creates a new initialized [`Once`].
initialized(data: T) -> Once<T>73     pub const fn initialized(data: T) -> Once<T> {
74         Self {
75             state: AtomicUsize::new(COMPLETE),
76             data: UnsafeCell::new(MaybeUninit::new(data)),
77         }
78     }
79 
80     /// Get a reference to the initialized instance. Must only be called once COMPLETE.
force_get(&self) -> &T81     unsafe fn force_get(&self) -> &T {
82         // SAFETY:
83         // * `UnsafeCell`/inner deref: data never changes again
84         // * `MaybeUninit`/outer deref: data was initialized
85         &*(*self.data.get()).as_ptr()
86     }
87 
88     /// Get a reference to the initialized instance. Must only be called once COMPLETE.
force_get_mut(&mut self) -> &mut T89     unsafe fn force_get_mut(&mut self) -> &mut T {
90         // SAFETY:
91         // * `UnsafeCell`/inner deref: data never changes again
92         // * `MaybeUninit`/outer deref: data was initialized
93         &mut *(*self.data.get()).as_mut_ptr()
94     }
95 
96     /// Get a reference to the initialized instance. Must only be called once COMPLETE.
force_into_inner(self) -> T97     unsafe fn force_into_inner(self) -> T {
98         // SAFETY:
99         // * `UnsafeCell`/inner deref: data never changes again
100         // * `MaybeUninit`/outer deref: data was initialized
101         (*self.data.get()).as_ptr().read()
102     }
103 
104     /// Performs an initialization routine once and only once. The given closure
105     /// will be executed if this is the first time `call_once` has been called,
106     /// and otherwise the routine will *not* be invoked.
107     ///
108     /// This method will block the calling thread if another initialization
109     /// routine is currently running.
110     ///
111     /// When this function returns, it is guaranteed that some initialization
112     /// has run and completed (it may not be the closure specified). The
113     /// returned pointer will point to the result from the closure that was
114     /// run.
115     ///
116     /// # Panics
117     ///
118     /// This function will panic if the [`Once`] previously panicked while attempting
119     /// to initialize. This is similar to the poisoning behaviour of `std::sync`'s
120     /// primitives.
121     ///
122     /// # Examples
123     ///
124     /// ```
125     /// use spin;
126     ///
127     /// static INIT: spin::Once<usize> = spin::Once::new();
128     ///
129     /// fn get_cached_val() -> usize {
130     ///     *INIT.call_once(expensive_computation)
131     /// }
132     ///
133     /// fn expensive_computation() -> usize {
134     ///     // ...
135     /// # 2
136     /// }
137     /// ```
call_once<F: FnOnce() -> T>(&self, f: F) -> &T138     pub fn call_once<F: FnOnce() -> T>(&self, f: F) -> &T {
139         let mut status = self.state.load(Ordering::SeqCst);
140 
141         if status == INCOMPLETE {
142             status = self.state.compare_and_swap(
143                 INCOMPLETE,
144                 RUNNING,
145                 Ordering::SeqCst,
146             );
147 
148             if status == INCOMPLETE { // We init
149                 // We use a guard (Finish) to catch panics caused by builder
150                 let mut finish = Finish { state: &self.state, panicked: true };
151                 unsafe {
152                     // SAFETY:
153                     // `UnsafeCell`/deref: currently the only accessor, mutably
154                     // and immutably by cas exclusion.
155                     // `write`: pointer comes from `MaybeUninit`.
156                     (*self.data.get()).as_mut_ptr().write(f())
157                 };
158                 finish.panicked = false;
159 
160                 status = COMPLETE;
161                 self.state.store(status, Ordering::SeqCst);
162 
163                 // This next line is strictly an optimization
164                 return unsafe { self.force_get() };
165             }
166         }
167 
168         self
169             .poll()
170             .unwrap_or_else(|| unreachable!("Encountered INCOMPLETE when polling Once"))
171     }
172 
173     /// Returns a reference to the inner value if the [`Once`] has been initialized.
get(&self) -> Option<&T>174     pub fn get(&self) -> Option<&T> {
175         match self.state.load(Ordering::SeqCst) {
176             COMPLETE => Some(unsafe { self.force_get() }),
177             _ => None,
178         }
179     }
180 
181     /// Returns a mutable reference to the inner value if the [`Once`] has been initialized.
182     ///
183     /// Because this method requires a mutable reference to the [`Once`], no synchronization
184     /// overhead is required to access the inner value. In effect, it is zero-cost.
get_mut(&mut self) -> Option<&mut T>185     pub fn get_mut(&mut self) -> Option<&mut T> {
186         match *self.state.get_mut() {
187             COMPLETE => Some(unsafe { self.force_get_mut() }),
188             _ => None,
189         }
190     }
191 
192     /// Returns a the inner value if the [`Once`] has been initialized.
193     ///
194     /// Because this method requires ownershup of the [`Once`], no synchronization overhead
195     /// is required to access the inner value. In effect, it is zero-cost.
try_into_inner(mut self) -> Option<T>196     pub fn try_into_inner(mut self) -> Option<T> {
197         match *self.state.get_mut() {
198             COMPLETE => Some(unsafe { self.force_into_inner() }),
199             _ => None,
200         }
201     }
202 
203     /// Returns a reference to the inner value if the [`Once`] has been initialized.
is_completed(&self) -> bool204     pub fn is_completed(&self) -> bool {
205         self.state.load(Ordering::SeqCst) == COMPLETE
206     }
207 
208     /// Spins until the [`Once`] contains a value.
209     ///
210     /// Note that in releases prior to `0.7`, this function had the behaviour of [`Once::poll`].
211     ///
212     /// # Panics
213     ///
214     /// This function will panic if the [`Once`] previously panicked while attempting
215     /// to initialize. This is similar to the poisoning behaviour of `std::sync`'s
216     /// primitives.
wait(&self) -> &T217     pub fn wait(&self) -> &T {
218         loop {
219             match self.poll() {
220                 Some(x) => break x,
221                 None => crate::relax(),
222             }
223         }
224     }
225 
226     /// Like [`Once::get`], but will spin if the [`Once`] is in the process of being
227     /// initialized. If initialization has not even begun, `None` will be returned.
228     ///
229     /// Note that in releases prior to `0.7`, this function was named `wait`.
230     ///
231     /// # Panics
232     ///
233     /// This function will panic if the [`Once`] previously panicked while attempting
234     /// to initialize. This is similar to the poisoning behaviour of `std::sync`'s
235     /// primitives.
poll(&self) -> Option<&T>236     pub fn poll(&self) -> Option<&T> {
237         loop {
238             match self.state.load(Ordering::SeqCst) {
239                 INCOMPLETE => return None,
240                 RUNNING => crate::relax(), // We spin
241                 COMPLETE => return Some(unsafe { self.force_get() }),
242                 PANICKED => panic!("Once previously poisoned by a panicked"),
243                 _ => unsafe { unreachable() },
244             }
245         }
246     }
247 }
248 
249 impl<T> From<T> for Once<T> {
from(data: T) -> Self250     fn from(data: T) -> Self {
251         Self::initialized(data)
252     }
253 }
254 
255 impl<T> Drop for Once<T> {
drop(&mut self)256     fn drop(&mut self) {
257         if self.state.load(Ordering::SeqCst) == COMPLETE {
258             unsafe {
259                 //TODO: Use MaybeUninit::assume_init_drop once stabilised
260                 core::ptr::drop_in_place((*self.data.get()).as_mut_ptr());
261             }
262         }
263     }
264 }
265 
266 struct Finish<'a> {
267     state: &'a AtomicUsize,
268     panicked: bool,
269 }
270 
271 impl<'a> Drop for Finish<'a> {
drop(&mut self)272     fn drop(&mut self) {
273         if self.panicked {
274             self.state.store(PANICKED, Ordering::SeqCst);
275         }
276     }
277 }
278 
279 #[cfg(test)]
280 mod tests {
281     use std::prelude::v1::*;
282 
283     use std::sync::mpsc::channel;
284     use std::thread;
285     use super::Once;
286 
287     #[test]
smoke_once()288     fn smoke_once() {
289         static O: Once<()> = Once::new();
290         let mut a = 0;
291         O.call_once(|| a += 1);
292         assert_eq!(a, 1);
293         O.call_once(|| a += 1);
294         assert_eq!(a, 1);
295     }
296 
297     #[test]
smoke_once_value()298     fn smoke_once_value() {
299         static O: Once<usize> = Once::new();
300         let a = O.call_once(|| 1);
301         assert_eq!(*a, 1);
302         let b = O.call_once(|| 2);
303         assert_eq!(*b, 1);
304     }
305 
306     #[test]
stampede_once()307     fn stampede_once() {
308         static O: Once<()> = Once::new();
309         static mut RUN: bool = false;
310 
311         let (tx, rx) = channel();
312         for _ in 0..10 {
313             let tx = tx.clone();
314             thread::spawn(move|| {
315                 for _ in 0..4 { thread::yield_now() }
316                 unsafe {
317                     O.call_once(|| {
318                         assert!(!RUN);
319                         RUN = true;
320                     });
321                     assert!(RUN);
322                 }
323                 tx.send(()).unwrap();
324             });
325         }
326 
327         unsafe {
328             O.call_once(|| {
329                 assert!(!RUN);
330                 RUN = true;
331             });
332             assert!(RUN);
333         }
334 
335         for _ in 0..10 {
336             rx.recv().unwrap();
337         }
338     }
339 
340     #[test]
get()341     fn get() {
342         static INIT: Once<usize> = Once::new();
343 
344         assert!(INIT.get().is_none());
345         INIT.call_once(|| 2);
346         assert_eq!(INIT.get().map(|r| *r), Some(2));
347     }
348 
349     #[test]
get_no_wait()350     fn get_no_wait() {
351         static INIT: Once<usize> = Once::new();
352 
353         assert!(INIT.get().is_none());
354         thread::spawn(move|| {
355             INIT.call_once(|| loop { });
356         });
357         assert!(INIT.get().is_none());
358     }
359 
360 
361     #[test]
poll()362     fn poll() {
363         static INIT: Once<usize> = Once::new();
364 
365         assert!(INIT.poll().is_none());
366         INIT.call_once(|| 3);
367         assert_eq!(INIT.poll().map(|r| *r), Some(3));
368     }
369 
370 
371     #[test]
wait()372     fn wait() {
373         static INIT: Once<usize> = Once::new();
374 
375         std::thread::spawn(|| {
376             assert_eq!(*INIT.wait(), 3);
377             assert!(INIT.is_completed());
378         });
379 
380         for _ in 0..4 { thread::yield_now() }
381 
382         assert!(INIT.poll().is_none());
383         INIT.call_once(|| 3);
384     }
385 
386     #[test]
387     #[ignore = "Android uses panic_abort"]
panic()388     fn panic() {
389         use ::std::panic;
390 
391         static INIT: Once<()> = Once::new();
392 
393         // poison the once
394         let t = panic::catch_unwind(|| {
395             INIT.call_once(|| panic!());
396         });
397         assert!(t.is_err());
398 
399         // poisoning propagates
400         let t = panic::catch_unwind(|| {
401             INIT.call_once(|| {});
402         });
403         assert!(t.is_err());
404     }
405 
406     #[test]
init_constant()407     fn init_constant() {
408         static O: Once<()> = Once::INIT;
409         let mut a = 0;
410         O.call_once(|| a += 1);
411         assert_eq!(a, 1);
412         O.call_once(|| a += 1);
413         assert_eq!(a, 1);
414     }
415 
416     static mut CALLED: bool = false;
417 
418     struct DropTest {}
419 
420     impl Drop for DropTest {
drop(&mut self)421         fn drop(&mut self) {
422             unsafe {
423                 CALLED = true;
424             }
425         }
426     }
427 
428     #[test]
drop()429     fn drop() {
430         unsafe {
431             CALLED = false;
432         }
433 
434         {
435             let once = Once::new();
436             once.call_once(|| DropTest {});
437         }
438 
439         assert!(unsafe {
440             CALLED
441         });
442     }
443 
444     #[test]
skip_uninit_drop()445     fn skip_uninit_drop() {
446         unsafe {
447             CALLED = false;
448         }
449 
450         {
451             let once = Once::<DropTest>::new();
452         }
453 
454         assert!(unsafe {
455             !CALLED
456         });
457     }
458 }
459