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 (feature = "critical-section" )] |
23 | use atomic_polyfill as atomic; |
24 | #[cfg (not(feature = "critical-section" ))] |
25 | use core::sync::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 ] |
42 | pub const fn new() -> OnceNonZeroUsize { |
43 | OnceNonZeroUsize { inner: AtomicUsize::new(0) } |
44 | } |
45 | |
46 | /// Gets the underlying value. |
47 | #[inline ] |
48 | pub fn get(&self) -> Option<NonZeroUsize> { |
49 | let val = self.inner.load(Ordering::Acquire); |
50 | NonZeroUsize::new(val) |
51 | } |
52 | |
53 | /// Sets the contents of this cell to `value`. |
54 | /// |
55 | /// Returns `Ok(())` if the cell was empty and `Err(())` if it was |
56 | /// full. |
57 | #[inline ] |
58 | pub fn set(&self, value: NonZeroUsize) -> Result<(), ()> { |
59 | let exchange = |
60 | self.inner.compare_exchange(0, value.get(), Ordering::AcqRel, Ordering::Acquire); |
61 | match exchange { |
62 | Ok(_) => Ok(()), |
63 | Err(_) => Err(()), |
64 | } |
65 | } |
66 | |
67 | /// Gets the contents of the cell, initializing it with `f` if the cell was |
68 | /// empty. |
69 | /// |
70 | /// If several threads concurrently run `get_or_init`, more than one `f` can |
71 | /// be called. However, all threads will return the same value, produced by |
72 | /// some `f`. |
73 | pub fn get_or_init<F>(&self, f: F) -> NonZeroUsize |
74 | where |
75 | F: FnOnce() -> NonZeroUsize, |
76 | { |
77 | enum Void {} |
78 | match self.get_or_try_init(|| Ok::<NonZeroUsize, Void>(f())) { |
79 | Ok(val) => val, |
80 | Err(void) => match void {}, |
81 | } |
82 | } |
83 | |
84 | /// Gets the contents of the cell, initializing it with `f` if |
85 | /// the cell was empty. If the cell was empty and `f` failed, an |
86 | /// error is returned. |
87 | /// |
88 | /// If several threads concurrently run `get_or_init`, more than one `f` can |
89 | /// be called. However, all threads will return the same value, produced by |
90 | /// some `f`. |
91 | pub fn get_or_try_init<F, E>(&self, f: F) -> Result<NonZeroUsize, E> |
92 | where |
93 | F: FnOnce() -> Result<NonZeroUsize, E>, |
94 | { |
95 | let val = self.inner.load(Ordering::Acquire); |
96 | let res = match NonZeroUsize::new(val) { |
97 | Some(it) => it, |
98 | None => { |
99 | let mut val = f()?.get(); |
100 | let exchange = |
101 | self.inner.compare_exchange(0, val, Ordering::AcqRel, Ordering::Acquire); |
102 | if let Err(old) = exchange { |
103 | val = old; |
104 | } |
105 | unsafe { NonZeroUsize::new_unchecked(val) } |
106 | } |
107 | }; |
108 | Ok(res) |
109 | } |
110 | } |
111 | |
112 | /// A thread-safe cell which can be written to only once. |
113 | #[derive (Default, Debug)] |
114 | pub struct OnceBool { |
115 | inner: OnceNonZeroUsize, |
116 | } |
117 | |
118 | impl OnceBool { |
119 | /// Creates a new empty cell. |
120 | #[inline ] |
121 | pub const fn new() -> OnceBool { |
122 | OnceBool { inner: OnceNonZeroUsize::new() } |
123 | } |
124 | |
125 | /// Gets the underlying value. |
126 | #[inline ] |
127 | pub fn get(&self) -> Option<bool> { |
128 | self.inner.get().map(OnceBool::from_usize) |
129 | } |
130 | |
131 | /// Sets the contents of this cell to `value`. |
132 | /// |
133 | /// Returns `Ok(())` if the cell was empty and `Err(())` if it was |
134 | /// full. |
135 | #[inline ] |
136 | pub fn set(&self, value: bool) -> Result<(), ()> { |
137 | self.inner.set(OnceBool::to_usize(value)) |
138 | } |
139 | |
140 | /// Gets the contents of the cell, initializing it with `f` if the cell was |
141 | /// empty. |
142 | /// |
143 | /// If several threads concurrently run `get_or_init`, more than one `f` can |
144 | /// be called. However, all threads will return the same value, produced by |
145 | /// some `f`. |
146 | pub fn get_or_init<F>(&self, f: F) -> bool |
147 | where |
148 | F: FnOnce() -> bool, |
149 | { |
150 | OnceBool::from_usize(self.inner.get_or_init(|| OnceBool::to_usize(f()))) |
151 | } |
152 | |
153 | /// Gets the contents of the cell, initializing it with `f` if |
154 | /// the cell was empty. If the cell was empty and `f` failed, an |
155 | /// error is returned. |
156 | /// |
157 | /// If several threads concurrently run `get_or_init`, more than one `f` can |
158 | /// be called. However, all threads will return the same value, produced by |
159 | /// some `f`. |
160 | pub fn get_or_try_init<F, E>(&self, f: F) -> Result<bool, E> |
161 | where |
162 | F: FnOnce() -> Result<bool, E>, |
163 | { |
164 | self.inner.get_or_try_init(|| f().map(OnceBool::to_usize)).map(OnceBool::from_usize) |
165 | } |
166 | |
167 | #[inline ] |
168 | fn from_usize(value: NonZeroUsize) -> bool { |
169 | value.get() == 1 |
170 | } |
171 | |
172 | #[inline ] |
173 | fn to_usize(value: bool) -> NonZeroUsize { |
174 | unsafe { NonZeroUsize::new_unchecked(if value { 1 } else { 2 }) } |
175 | } |
176 | } |
177 | |
178 | /// A thread-safe cell which can be written to only once. |
179 | pub struct OnceRef<'a, T> { |
180 | inner: AtomicPtr<T>, |
181 | ghost: PhantomData<UnsafeCell<&'a T>>, |
182 | } |
183 | |
184 | // TODO: Replace UnsafeCell with SyncUnsafeCell once stabilized |
185 | unsafe impl<'a, T: Sync> Sync for OnceRef<'a, T> {} |
186 | |
187 | impl<'a, T> core::fmt::Debug for OnceRef<'a, T> { |
188 | fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { |
189 | write!(f, "OnceRef( {:?})" , self.inner) |
190 | } |
191 | } |
192 | |
193 | impl<'a, T> Default for OnceRef<'a, T> { |
194 | fn default() -> Self { |
195 | Self::new() |
196 | } |
197 | } |
198 | |
199 | impl<'a, T> OnceRef<'a, T> { |
200 | /// Creates a new empty cell. |
201 | pub const fn new() -> OnceRef<'a, T> { |
202 | OnceRef { inner: AtomicPtr::new(ptr::null_mut()), ghost: PhantomData } |
203 | } |
204 | |
205 | /// Gets a reference to the underlying value. |
206 | pub fn get(&self) -> Option<&'a T> { |
207 | let ptr = self.inner.load(Ordering::Acquire); |
208 | unsafe { ptr.as_ref() } |
209 | } |
210 | |
211 | /// Sets the contents of this cell to `value`. |
212 | /// |
213 | /// Returns `Ok(())` if the cell was empty and `Err(value)` if it was |
214 | /// full. |
215 | pub fn set(&self, value: &'a T) -> Result<(), ()> { |
216 | let ptr = value as *const T as *mut T; |
217 | let exchange = |
218 | self.inner.compare_exchange(ptr::null_mut(), ptr, Ordering::AcqRel, Ordering::Acquire); |
219 | match exchange { |
220 | Ok(_) => Ok(()), |
221 | Err(_) => Err(()), |
222 | } |
223 | } |
224 | |
225 | /// Gets the contents of the cell, initializing it with `f` if the cell was |
226 | /// empty. |
227 | /// |
228 | /// If several threads concurrently run `get_or_init`, more than one `f` can |
229 | /// be called. However, all threads will return the same value, produced by |
230 | /// some `f`. |
231 | pub fn get_or_init<F>(&self, f: F) -> &'a T |
232 | where |
233 | F: FnOnce() -> &'a T, |
234 | { |
235 | enum Void {} |
236 | match self.get_or_try_init(|| Ok::<&'a T, Void>(f())) { |
237 | Ok(val) => val, |
238 | Err(void) => match void {}, |
239 | } |
240 | } |
241 | |
242 | /// Gets the contents of the cell, initializing it with `f` if |
243 | /// the cell was empty. If the cell was empty and `f` failed, an |
244 | /// error is returned. |
245 | /// |
246 | /// If several threads concurrently run `get_or_init`, more than one `f` can |
247 | /// be called. However, all threads will return the same value, produced by |
248 | /// some `f`. |
249 | pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&'a T, E> |
250 | where |
251 | F: FnOnce() -> Result<&'a T, E>, |
252 | { |
253 | let mut ptr = self.inner.load(Ordering::Acquire); |
254 | |
255 | if ptr.is_null() { |
256 | // TODO replace with `cast_mut` when MSRV reaches 1.65.0 (also in `set`) |
257 | ptr = f()? as *const T as *mut T; |
258 | let exchange = self.inner.compare_exchange( |
259 | ptr::null_mut(), |
260 | ptr, |
261 | Ordering::AcqRel, |
262 | Ordering::Acquire, |
263 | ); |
264 | if let Err(old) = exchange { |
265 | ptr = old; |
266 | } |
267 | } |
268 | |
269 | Ok(unsafe { &*ptr }) |
270 | } |
271 | |
272 | /// ```compile_fail |
273 | /// use once_cell::race::OnceRef; |
274 | /// |
275 | /// let mut l = OnceRef::new(); |
276 | /// |
277 | /// { |
278 | /// let y = 2; |
279 | /// let mut r = OnceRef::new(); |
280 | /// r.set(&y).unwrap(); |
281 | /// core::mem::swap(&mut l, &mut r); |
282 | /// } |
283 | /// |
284 | /// // l now contains a dangling reference to y |
285 | /// eprintln!("uaf: {}" , l.get().unwrap()); |
286 | /// ``` |
287 | fn _dummy() {} |
288 | } |
289 | |
290 | #[cfg (feature = "alloc" )] |
291 | pub use self::once_box::OnceBox; |
292 | |
293 | #[cfg (feature = "alloc" )] |
294 | mod once_box { |
295 | use super::atomic::{AtomicPtr, Ordering}; |
296 | use core::{marker::PhantomData, ptr}; |
297 | |
298 | use alloc::boxed::Box; |
299 | |
300 | /// A thread-safe cell which can be written to only once. |
301 | pub struct OnceBox<T> { |
302 | inner: AtomicPtr<T>, |
303 | ghost: PhantomData<Option<Box<T>>>, |
304 | } |
305 | |
306 | impl<T> core::fmt::Debug for OnceBox<T> { |
307 | fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { |
308 | write!(f, "OnceBox( {:?})" , self.inner.load(Ordering::Relaxed)) |
309 | } |
310 | } |
311 | |
312 | impl<T> Default for OnceBox<T> { |
313 | fn default() -> Self { |
314 | Self::new() |
315 | } |
316 | } |
317 | |
318 | impl<T> Drop for OnceBox<T> { |
319 | fn drop(&mut self) { |
320 | let ptr = *self.inner.get_mut(); |
321 | if !ptr.is_null() { |
322 | drop(unsafe { Box::from_raw(ptr) }) |
323 | } |
324 | } |
325 | } |
326 | |
327 | impl<T> OnceBox<T> { |
328 | /// Creates a new empty cell. |
329 | pub const fn new() -> OnceBox<T> { |
330 | OnceBox { inner: AtomicPtr::new(ptr::null_mut()), ghost: PhantomData } |
331 | } |
332 | |
333 | /// Gets a reference to the underlying value. |
334 | pub fn get(&self) -> Option<&T> { |
335 | let ptr = self.inner.load(Ordering::Acquire); |
336 | if ptr.is_null() { |
337 | return None; |
338 | } |
339 | Some(unsafe { &*ptr }) |
340 | } |
341 | |
342 | /// Sets the contents of this cell to `value`. |
343 | /// |
344 | /// Returns `Ok(())` if the cell was empty and `Err(value)` if it was |
345 | /// full. |
346 | pub fn set(&self, value: Box<T>) -> Result<(), Box<T>> { |
347 | let ptr = Box::into_raw(value); |
348 | let exchange = self.inner.compare_exchange( |
349 | ptr::null_mut(), |
350 | ptr, |
351 | Ordering::AcqRel, |
352 | Ordering::Acquire, |
353 | ); |
354 | if exchange.is_err() { |
355 | let value = unsafe { Box::from_raw(ptr) }; |
356 | return Err(value); |
357 | } |
358 | Ok(()) |
359 | } |
360 | |
361 | /// Gets the contents of the cell, initializing it with `f` if the cell was |
362 | /// empty. |
363 | /// |
364 | /// If several threads concurrently run `get_or_init`, more than one `f` can |
365 | /// be called. However, all threads will return the same value, produced by |
366 | /// some `f`. |
367 | pub fn get_or_init<F>(&self, f: F) -> &T |
368 | where |
369 | F: FnOnce() -> Box<T>, |
370 | { |
371 | enum Void {} |
372 | match self.get_or_try_init(|| Ok::<Box<T>, Void>(f())) { |
373 | Ok(val) => val, |
374 | Err(void) => match void {}, |
375 | } |
376 | } |
377 | |
378 | /// Gets the contents of the cell, initializing it with `f` if |
379 | /// the cell was empty. If the cell was empty and `f` failed, an |
380 | /// error is returned. |
381 | /// |
382 | /// If several threads concurrently run `get_or_init`, more than one `f` can |
383 | /// be called. However, all threads will return the same value, produced by |
384 | /// some `f`. |
385 | pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E> |
386 | where |
387 | F: FnOnce() -> Result<Box<T>, E>, |
388 | { |
389 | let mut ptr = self.inner.load(Ordering::Acquire); |
390 | |
391 | if ptr.is_null() { |
392 | let val = f()?; |
393 | ptr = Box::into_raw(val); |
394 | let exchange = self.inner.compare_exchange( |
395 | ptr::null_mut(), |
396 | ptr, |
397 | Ordering::AcqRel, |
398 | Ordering::Acquire, |
399 | ); |
400 | if let Err(old) = exchange { |
401 | drop(unsafe { Box::from_raw(ptr) }); |
402 | ptr = old; |
403 | } |
404 | }; |
405 | Ok(unsafe { &*ptr }) |
406 | } |
407 | } |
408 | |
409 | unsafe impl<T: Sync + Send> Sync for OnceBox<T> {} |
410 | |
411 | /// ```compile_fail |
412 | /// struct S(*mut ()); |
413 | /// unsafe impl Sync for S {} |
414 | /// |
415 | /// fn share<T: Sync>(_: &T) {} |
416 | /// share(&once_cell::race::OnceBox::<S>::new()); |
417 | /// ``` |
418 | fn _dummy() {} |
419 | } |
420 | |