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