• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 //! Thread-safe, non-blocking, "first one wins" flavor of `OnceCell`.
2 //!
3 //! If two threads race to initialize a type from the `race` module, they
4 //! don't block, execute initialization function together, but only one of
5 //! them stores the result.
6 //!
7 //! This module does not require `std` feature.
8 //!
9 //! # Atomic orderings
10 //!
11 //! All types in this module use `Acquire` and `Release`
12 //! [atomic orderings](Ordering) for all their operations. While this is not
13 //! strictly necessary for types other than `OnceBox`, it is useful for users as
14 //! it allows them to be certain that after `get` or `get_or_init` returns on
15 //! one thread, any side-effects caused by the setter thread prior to them
16 //! calling `set` or `get_or_init` will be made visible to that thread; without
17 //! it, it's possible for it to appear as if they haven't happened yet from the
18 //! getter thread's perspective. This is an acceptable tradeoff to make since
19 //! `Acquire` and `Release` have very little performance overhead on most
20 //! architectures versus `Relaxed`.
21 
22 #[cfg(not(feature = "portable-atomic"))]
23 use core::sync::atomic;
24 #[cfg(feature = "portable-atomic")]
25 use portable_atomic as atomic;
26 
27 use atomic::{AtomicPtr, AtomicUsize, Ordering};
28 use core::cell::UnsafeCell;
29 use core::marker::PhantomData;
30 use core::num::NonZeroUsize;
31 use core::ptr;
32 
33 /// A thread-safe cell which can be written to only once.
34 #[derive(Default, Debug)]
35 pub struct OnceNonZeroUsize {
36     inner: AtomicUsize,
37 }
38 
39 impl OnceNonZeroUsize {
40     /// Creates a new empty cell.
41     #[inline]
new() -> OnceNonZeroUsize42     pub const fn new() -> OnceNonZeroUsize {
43         OnceNonZeroUsize { inner: AtomicUsize::new(0) }
44     }
45 
46     /// Gets the underlying value.
47     #[inline]
get(&self) -> Option<NonZeroUsize>48     pub fn get(&self) -> Option<NonZeroUsize> {
49         let val = self.inner.load(Ordering::Acquire);
50         NonZeroUsize::new(val)
51     }
52 
53     /// Get the reference to the underlying value, without checking if the cell
54     /// is initialized.
55     ///
56     /// # Safety
57     ///
58     /// Caller must ensure that the cell is in initialized state, and that
59     /// the contents are acquired by (synchronized to) this thread.
get_unchecked(&self) -> NonZeroUsize60     pub unsafe fn get_unchecked(&self) -> NonZeroUsize {
61         #[inline(always)]
62         fn as_const_ptr(r: &AtomicUsize) -> *const usize {
63             use core::mem::align_of;
64 
65             let p: *const AtomicUsize = r;
66             // SAFETY: "This type has the same size and bit validity as
67             // the underlying integer type, usize. However, the alignment of
68             // this type is always equal to its size, even on targets where
69             // usize has a lesser alignment."
70             const _ALIGNMENT_COMPATIBLE: () =
71                 assert!(align_of::<AtomicUsize>() % align_of::<usize>() == 0);
72             p.cast::<usize>()
73         }
74 
75         // TODO(MSRV-1.70): Use `AtomicUsize::as_ptr().cast_const()`
76         // See https://github.com/rust-lang/rust/issues/138246.
77         let p = as_const_ptr(&self.inner);
78 
79         // SAFETY: The caller is responsible for ensuring that the value
80         // was initialized and that the contents have been acquired by
81         // this thread. Assuming that, we can assume there will be no
82         // conflicting writes to the value since the value will never
83         // change once initialized. This relies on the statement in
84         // https://doc.rust-lang.org/1.83.0/core/sync/atomic/ that "(A
85         // `compare_exchange` or `compare_exchange_weak` that does not
86         // succeed is not considered a write."
87         let val = unsafe { p.read() };
88 
89         // SAFETY: The caller is responsible for ensuring the value is
90         // initialized and thus not zero.
91         unsafe { NonZeroUsize::new_unchecked(val) }
92     }
93 
94     /// Sets the contents of this cell to `value`.
95     ///
96     /// Returns `Ok(())` if the cell was empty and `Err(())` if it was
97     /// full.
98     #[inline]
set(&self, value: NonZeroUsize) -> Result<(), ()>99     pub fn set(&self, value: NonZeroUsize) -> Result<(), ()> {
100         let exchange =
101             self.inner.compare_exchange(0, value.get(), Ordering::AcqRel, Ordering::Acquire);
102         match exchange {
103             Ok(_) => Ok(()),
104             Err(_) => Err(()),
105         }
106     }
107 
108     /// Gets the contents of the cell, initializing it with `f` if the cell was
109     /// empty.
110     ///
111     /// If several threads concurrently run `get_or_init`, more than one `f` can
112     /// be called. However, all threads will return the same value, produced by
113     /// some `f`.
get_or_init<F>(&self, f: F) -> NonZeroUsize where F: FnOnce() -> NonZeroUsize,114     pub fn get_or_init<F>(&self, f: F) -> NonZeroUsize
115     where
116         F: FnOnce() -> NonZeroUsize,
117     {
118         enum Void {}
119         match self.get_or_try_init(|| Ok::<NonZeroUsize, Void>(f())) {
120             Ok(val) => val,
121             Err(void) => match void {},
122         }
123     }
124 
125     /// Gets the contents of the cell, initializing it with `f` if
126     /// the cell was empty. If the cell was empty and `f` failed, an
127     /// error is returned.
128     ///
129     /// If several threads concurrently run `get_or_init`, more than one `f` can
130     /// be called. However, all threads will return the same value, produced by
131     /// some `f`.
get_or_try_init<F, E>(&self, f: F) -> Result<NonZeroUsize, E> where F: FnOnce() -> Result<NonZeroUsize, E>,132     pub fn get_or_try_init<F, E>(&self, f: F) -> Result<NonZeroUsize, E>
133     where
134         F: FnOnce() -> Result<NonZeroUsize, E>,
135     {
136         let val = self.inner.load(Ordering::Acquire);
137         match NonZeroUsize::new(val) {
138             Some(it) => Ok(it),
139             None => self.init(f),
140         }
141     }
142 
143     #[cold]
144     #[inline(never)]
init<E>(&self, f: impl FnOnce() -> Result<NonZeroUsize, E>) -> Result<NonZeroUsize, E>145     fn init<E>(&self, f: impl FnOnce() -> Result<NonZeroUsize, E>) -> Result<NonZeroUsize, E> {
146         let mut val = f()?.get();
147         let exchange = self.inner.compare_exchange(0, val, Ordering::AcqRel, Ordering::Acquire);
148         if let Err(old) = exchange {
149             val = old;
150         }
151         Ok(unsafe { NonZeroUsize::new_unchecked(val) })
152     }
153 }
154 
155 /// A thread-safe cell which can be written to only once.
156 #[derive(Default, Debug)]
157 pub struct OnceBool {
158     inner: OnceNonZeroUsize,
159 }
160 
161 impl OnceBool {
162     /// Creates a new empty cell.
163     #[inline]
new() -> OnceBool164     pub const fn new() -> OnceBool {
165         OnceBool { inner: OnceNonZeroUsize::new() }
166     }
167 
168     /// Gets the underlying value.
169     #[inline]
get(&self) -> Option<bool>170     pub fn get(&self) -> Option<bool> {
171         self.inner.get().map(OnceBool::from_usize)
172     }
173 
174     /// Sets the contents of this cell to `value`.
175     ///
176     /// Returns `Ok(())` if the cell was empty and `Err(())` if it was
177     /// full.
178     #[inline]
set(&self, value: bool) -> Result<(), ()>179     pub fn set(&self, value: bool) -> Result<(), ()> {
180         self.inner.set(OnceBool::to_usize(value))
181     }
182 
183     /// Gets the contents of the cell, initializing it with `f` if the cell was
184     /// empty.
185     ///
186     /// If several threads concurrently run `get_or_init`, more than one `f` can
187     /// be called. However, all threads will return the same value, produced by
188     /// some `f`.
get_or_init<F>(&self, f: F) -> bool where F: FnOnce() -> bool,189     pub fn get_or_init<F>(&self, f: F) -> bool
190     where
191         F: FnOnce() -> bool,
192     {
193         OnceBool::from_usize(self.inner.get_or_init(|| OnceBool::to_usize(f())))
194     }
195 
196     /// Gets the contents of the cell, initializing it with `f` if
197     /// the cell was empty. If the cell was empty and `f` failed, an
198     /// error is returned.
199     ///
200     /// If several threads concurrently run `get_or_init`, more than one `f` can
201     /// be called. However, all threads will return the same value, produced by
202     /// some `f`.
get_or_try_init<F, E>(&self, f: F) -> Result<bool, E> where F: FnOnce() -> Result<bool, E>,203     pub fn get_or_try_init<F, E>(&self, f: F) -> Result<bool, E>
204     where
205         F: FnOnce() -> Result<bool, E>,
206     {
207         self.inner.get_or_try_init(|| f().map(OnceBool::to_usize)).map(OnceBool::from_usize)
208     }
209 
210     #[inline]
from_usize(value: NonZeroUsize) -> bool211     fn from_usize(value: NonZeroUsize) -> bool {
212         value.get() == 1
213     }
214 
215     #[inline]
to_usize(value: bool) -> NonZeroUsize216     fn to_usize(value: bool) -> NonZeroUsize {
217         unsafe { NonZeroUsize::new_unchecked(if value { 1 } else { 2 }) }
218     }
219 }
220 
221 /// A thread-safe cell which can be written to only once.
222 pub struct OnceRef<'a, T> {
223     inner: AtomicPtr<T>,
224     ghost: PhantomData<UnsafeCell<&'a T>>,
225 }
226 
227 // TODO: Replace UnsafeCell with SyncUnsafeCell once stabilized
228 unsafe impl<'a, T: Sync> Sync for OnceRef<'a, T> {}
229 
230 impl<'a, T> core::fmt::Debug for OnceRef<'a, T> {
fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result231     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
232         write!(f, "OnceRef({:?})", self.inner)
233     }
234 }
235 
236 impl<'a, T> Default for OnceRef<'a, T> {
default() -> Self237     fn default() -> Self {
238         Self::new()
239     }
240 }
241 
242 impl<'a, T> OnceRef<'a, T> {
243     /// Creates a new empty cell.
new() -> OnceRef<'a, T>244     pub const fn new() -> OnceRef<'a, T> {
245         OnceRef { inner: AtomicPtr::new(ptr::null_mut()), ghost: PhantomData }
246     }
247 
248     /// Gets a reference to the underlying value.
get(&self) -> Option<&'a T>249     pub fn get(&self) -> Option<&'a T> {
250         let ptr = self.inner.load(Ordering::Acquire);
251         unsafe { ptr.as_ref() }
252     }
253 
254     /// Sets the contents of this cell to `value`.
255     ///
256     /// Returns `Ok(())` if the cell was empty and `Err(value)` if it was
257     /// full.
set(&self, value: &'a T) -> Result<(), ()>258     pub fn set(&self, value: &'a T) -> Result<(), ()> {
259         let ptr = value as *const T as *mut T;
260         let exchange =
261             self.inner.compare_exchange(ptr::null_mut(), ptr, Ordering::AcqRel, Ordering::Acquire);
262         match exchange {
263             Ok(_) => Ok(()),
264             Err(_) => Err(()),
265         }
266     }
267 
268     /// Gets the contents of the cell, initializing it with `f` if the cell was
269     /// empty.
270     ///
271     /// If several threads concurrently run `get_or_init`, more than one `f` can
272     /// be called. However, all threads will return the same value, produced by
273     /// some `f`.
get_or_init<F>(&self, f: F) -> &'a T where F: FnOnce() -> &'a T,274     pub fn get_or_init<F>(&self, f: F) -> &'a T
275     where
276         F: FnOnce() -> &'a T,
277     {
278         enum Void {}
279         match self.get_or_try_init(|| Ok::<&'a T, Void>(f())) {
280             Ok(val) => val,
281             Err(void) => match void {},
282         }
283     }
284 
285     /// Gets the contents of the cell, initializing it with `f` if
286     /// the cell was empty. If the cell was empty and `f` failed, an
287     /// error is returned.
288     ///
289     /// If several threads concurrently run `get_or_init`, more than one `f` can
290     /// be called. However, all threads will return the same value, produced by
291     /// some `f`.
get_or_try_init<F, E>(&self, f: F) -> Result<&'a T, E> where F: FnOnce() -> Result<&'a T, E>,292     pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&'a T, E>
293     where
294         F: FnOnce() -> Result<&'a T, E>,
295     {
296         let mut ptr = self.inner.load(Ordering::Acquire);
297 
298         if ptr.is_null() {
299             // TODO replace with `cast_mut` when MSRV reaches 1.65.0 (also in `set`)
300             ptr = f()? as *const T as *mut T;
301             let exchange = self.inner.compare_exchange(
302                 ptr::null_mut(),
303                 ptr,
304                 Ordering::AcqRel,
305                 Ordering::Acquire,
306             );
307             if let Err(old) = exchange {
308                 ptr = old;
309             }
310         }
311 
312         Ok(unsafe { &*ptr })
313     }
314 
315     /// ```compile_fail
316     /// use once_cell::race::OnceRef;
317     ///
318     /// let mut l = OnceRef::new();
319     ///
320     /// {
321     ///     let y = 2;
322     ///     let mut r = OnceRef::new();
323     ///     r.set(&y).unwrap();
324     ///     core::mem::swap(&mut l, &mut r);
325     /// }
326     ///
327     /// // l now contains a dangling reference to y
328     /// eprintln!("uaf: {}", l.get().unwrap());
329     /// ```
_dummy()330     fn _dummy() {}
331 }
332 
333 #[cfg(feature = "alloc")]
334 pub use self::once_box::OnceBox;
335 
336 #[cfg(feature = "alloc")]
337 mod once_box {
338     use super::atomic::{AtomicPtr, Ordering};
339     use core::{marker::PhantomData, ptr};
340 
341     use alloc::boxed::Box;
342 
343     /// A thread-safe cell which can be written to only once.
344     pub struct OnceBox<T> {
345         inner: AtomicPtr<T>,
346         ghost: PhantomData<Option<Box<T>>>,
347     }
348 
349     impl<T> core::fmt::Debug for OnceBox<T> {
fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result350         fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
351             write!(f, "OnceBox({:?})", self.inner.load(Ordering::Relaxed))
352         }
353     }
354 
355     impl<T> Default for OnceBox<T> {
default() -> Self356         fn default() -> Self {
357             Self::new()
358         }
359     }
360 
361     impl<T> Drop for OnceBox<T> {
drop(&mut self)362         fn drop(&mut self) {
363             let ptr = *self.inner.get_mut();
364             if !ptr.is_null() {
365                 drop(unsafe { Box::from_raw(ptr) })
366             }
367         }
368     }
369 
370     impl<T> OnceBox<T> {
371         /// Creates a new empty cell.
new() -> OnceBox<T>372         pub const fn new() -> OnceBox<T> {
373             OnceBox { inner: AtomicPtr::new(ptr::null_mut()), ghost: PhantomData }
374         }
375 
376         /// Creates a new cell with the given value.
with_value(value: Box<T>) -> Self377         pub fn with_value(value: Box<T>) -> Self {
378             OnceBox { inner: AtomicPtr::new(Box::into_raw(value)), ghost: PhantomData }
379         }
380 
381         /// Gets a reference to the underlying value.
get(&self) -> Option<&T>382         pub fn get(&self) -> Option<&T> {
383             let ptr = self.inner.load(Ordering::Acquire);
384             if ptr.is_null() {
385                 return None;
386             }
387             Some(unsafe { &*ptr })
388         }
389 
390         /// Sets the contents of this cell to `value`.
391         ///
392         /// Returns `Ok(())` if the cell was empty and `Err(value)` if it was
393         /// full.
set(&self, value: Box<T>) -> Result<(), Box<T>>394         pub fn set(&self, value: Box<T>) -> Result<(), Box<T>> {
395             let ptr = Box::into_raw(value);
396             let exchange = self.inner.compare_exchange(
397                 ptr::null_mut(),
398                 ptr,
399                 Ordering::AcqRel,
400                 Ordering::Acquire,
401             );
402             if exchange.is_err() {
403                 let value = unsafe { Box::from_raw(ptr) };
404                 return Err(value);
405             }
406             Ok(())
407         }
408 
409         /// Gets the contents of the cell, initializing it with `f` if the cell was
410         /// empty.
411         ///
412         /// If several threads concurrently run `get_or_init`, more than one `f` can
413         /// be called. However, all threads will return the same value, produced by
414         /// some `f`.
get_or_init<F>(&self, f: F) -> &T where F: FnOnce() -> Box<T>,415         pub fn get_or_init<F>(&self, f: F) -> &T
416         where
417             F: FnOnce() -> Box<T>,
418         {
419             enum Void {}
420             match self.get_or_try_init(|| Ok::<Box<T>, Void>(f())) {
421                 Ok(val) => val,
422                 Err(void) => match void {},
423             }
424         }
425 
426         /// Gets the contents of the cell, initializing it with `f` if
427         /// the cell was empty. If the cell was empty and `f` failed, an
428         /// error is returned.
429         ///
430         /// If several threads concurrently run `get_or_init`, more than one `f` can
431         /// be called. However, all threads will return the same value, produced by
432         /// some `f`.
get_or_try_init<F, E>(&self, f: F) -> Result<&T, E> where F: FnOnce() -> Result<Box<T>, E>,433         pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E>
434         where
435             F: FnOnce() -> Result<Box<T>, E>,
436         {
437             let mut ptr = self.inner.load(Ordering::Acquire);
438 
439             if ptr.is_null() {
440                 let val = f()?;
441                 ptr = Box::into_raw(val);
442                 let exchange = self.inner.compare_exchange(
443                     ptr::null_mut(),
444                     ptr,
445                     Ordering::AcqRel,
446                     Ordering::Acquire,
447                 );
448                 if let Err(old) = exchange {
449                     drop(unsafe { Box::from_raw(ptr) });
450                     ptr = old;
451                 }
452             };
453             Ok(unsafe { &*ptr })
454         }
455     }
456 
457     unsafe impl<T: Sync + Send> Sync for OnceBox<T> {}
458 
459     impl<T: Clone> Clone for OnceBox<T> {
clone(&self) -> Self460         fn clone(&self) -> Self {
461             match self.get() {
462                 Some(value) => OnceBox::with_value(Box::new(value.clone())),
463                 None => OnceBox::new(),
464             }
465         }
466     }
467 
468     /// ```compile_fail
469     /// struct S(*mut ());
470     /// unsafe impl Sync for S {}
471     ///
472     /// fn share<T: Sync>(_: &T) {}
473     /// share(&once_cell::race::OnceBox::<S>::new());
474     /// ```
_dummy()475     fn _dummy() {}
476 }
477