1 use super::{Semaphore, SemaphorePermit, TryAcquireError}; 2 use crate::loom::cell::UnsafeCell; 3 use std::error::Error; 4 use std::fmt; 5 use std::future::Future; 6 use std::mem::MaybeUninit; 7 use std::ops::Drop; 8 use std::ptr; 9 use std::sync::atomic::{AtomicBool, Ordering}; 10 11 // This file contains an implementation of an OnceCell. The principle 12 // behind the safety the of the cell is that any thread with an `&OnceCell` may 13 // access the `value` field according the following rules: 14 // 15 // 1. When `value_set` is false, the `value` field may be modified by the 16 // thread holding the permit on the semaphore. 17 // 2. When `value_set` is true, the `value` field may be accessed immutably by 18 // any thread. 19 // 20 // It is an invariant that if the semaphore is closed, then `value_set` is true. 21 // The reverse does not necessarily hold — but if not, the semaphore may not 22 // have any available permits. 23 // 24 // A thread with a `&mut OnceCell` may modify the value in any way it wants as 25 // long as the invariants are upheld. 26 27 /// A thread-safe cell that can be written to only once. 28 /// 29 /// A `OnceCell` is typically used for global variables that need to be 30 /// initialized once on first use, but need no further changes. The `OnceCell` 31 /// in Tokio allows the initialization procedure to be asynchronous. 32 /// 33 /// # Examples 34 /// 35 /// ``` 36 /// use tokio::sync::OnceCell; 37 /// 38 /// async fn some_computation() -> u32 { 39 /// 1 + 1 40 /// } 41 /// 42 /// static ONCE: OnceCell<u32> = OnceCell::const_new(); 43 /// 44 /// #[tokio::main] 45 /// async fn main() { 46 /// let result = ONCE.get_or_init(some_computation).await; 47 /// assert_eq!(*result, 2); 48 /// } 49 /// ``` 50 /// 51 /// It is often useful to write a wrapper method for accessing the value. 52 /// 53 /// ``` 54 /// use tokio::sync::OnceCell; 55 /// 56 /// static ONCE: OnceCell<u32> = OnceCell::const_new(); 57 /// 58 /// async fn get_global_integer() -> &'static u32 { 59 /// ONCE.get_or_init(|| async { 60 /// 1 + 1 61 /// }).await 62 /// } 63 /// 64 /// #[tokio::main] 65 /// async fn main() { 66 /// let result = get_global_integer().await; 67 /// assert_eq!(*result, 2); 68 /// } 69 /// ``` 70 pub struct OnceCell<T> { 71 value_set: AtomicBool, 72 value: UnsafeCell<MaybeUninit<T>>, 73 semaphore: Semaphore, 74 } 75 76 impl<T> Default for OnceCell<T> { default() -> OnceCell<T>77 fn default() -> OnceCell<T> { 78 OnceCell::new() 79 } 80 } 81 82 impl<T: fmt::Debug> fmt::Debug for OnceCell<T> { fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result83 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { 84 fmt.debug_struct("OnceCell") 85 .field("value", &self.get()) 86 .finish() 87 } 88 } 89 90 impl<T: Clone> Clone for OnceCell<T> { clone(&self) -> OnceCell<T>91 fn clone(&self) -> OnceCell<T> { 92 OnceCell::new_with(self.get().cloned()) 93 } 94 } 95 96 impl<T: PartialEq> PartialEq for OnceCell<T> { eq(&self, other: &OnceCell<T>) -> bool97 fn eq(&self, other: &OnceCell<T>) -> bool { 98 self.get() == other.get() 99 } 100 } 101 102 impl<T: Eq> Eq for OnceCell<T> {} 103 104 impl<T> Drop for OnceCell<T> { drop(&mut self)105 fn drop(&mut self) { 106 if self.initialized_mut() { 107 unsafe { 108 self.value 109 .with_mut(|ptr| ptr::drop_in_place((&mut *ptr).as_mut_ptr())); 110 }; 111 } 112 } 113 } 114 115 impl<T> From<T> for OnceCell<T> { from(value: T) -> Self116 fn from(value: T) -> Self { 117 let semaphore = Semaphore::new(0); 118 semaphore.close(); 119 OnceCell { 120 value_set: AtomicBool::new(true), 121 value: UnsafeCell::new(MaybeUninit::new(value)), 122 semaphore, 123 } 124 } 125 } 126 127 impl<T> OnceCell<T> { 128 /// Creates a new empty `OnceCell` instance. new() -> Self129 pub fn new() -> Self { 130 OnceCell { 131 value_set: AtomicBool::new(false), 132 value: UnsafeCell::new(MaybeUninit::uninit()), 133 semaphore: Semaphore::new(1), 134 } 135 } 136 137 /// Creates a new `OnceCell` that contains the provided value, if any. 138 /// 139 /// If the `Option` is `None`, this is equivalent to `OnceCell::new`. 140 /// 141 /// [`OnceCell::new`]: crate::sync::OnceCell::new new_with(value: Option<T>) -> Self142 pub fn new_with(value: Option<T>) -> Self { 143 if let Some(v) = value { 144 OnceCell::from(v) 145 } else { 146 OnceCell::new() 147 } 148 } 149 150 /// Creates a new empty `OnceCell` instance. 151 /// 152 /// Equivalent to `OnceCell::new`, except that it can be used in static 153 /// variables. 154 /// 155 /// # Example 156 /// 157 /// ``` 158 /// use tokio::sync::OnceCell; 159 /// 160 /// static ONCE: OnceCell<u32> = OnceCell::const_new(); 161 /// 162 /// async fn get_global_integer() -> &'static u32 { 163 /// ONCE.get_or_init(|| async { 164 /// 1 + 1 165 /// }).await 166 /// } 167 /// 168 /// #[tokio::main] 169 /// async fn main() { 170 /// let result = get_global_integer().await; 171 /// assert_eq!(*result, 2); 172 /// } 173 /// ``` 174 #[cfg(all(feature = "parking_lot", not(all(loom, test))))] 175 #[cfg_attr(docsrs, doc(cfg(feature = "parking_lot")))] const_new() -> Self176 pub const fn const_new() -> Self { 177 OnceCell { 178 value_set: AtomicBool::new(false), 179 value: UnsafeCell::new(MaybeUninit::uninit()), 180 semaphore: Semaphore::const_new(1), 181 } 182 } 183 184 /// Returns `true` if the `OnceCell` currently contains a value, and `false` 185 /// otherwise. initialized(&self) -> bool186 pub fn initialized(&self) -> bool { 187 // Using acquire ordering so any threads that read a true from this 188 // atomic is able to read the value. 189 self.value_set.load(Ordering::Acquire) 190 } 191 192 /// Returns `true` if the `OnceCell` currently contains a value, and `false` 193 /// otherwise. initialized_mut(&mut self) -> bool194 fn initialized_mut(&mut self) -> bool { 195 *self.value_set.get_mut() 196 } 197 198 // SAFETY: The OnceCell must not be empty. get_unchecked(&self) -> &T199 unsafe fn get_unchecked(&self) -> &T { 200 &*self.value.with(|ptr| (*ptr).as_ptr()) 201 } 202 203 // SAFETY: The OnceCell must not be empty. get_unchecked_mut(&mut self) -> &mut T204 unsafe fn get_unchecked_mut(&mut self) -> &mut T { 205 &mut *self.value.with_mut(|ptr| (*ptr).as_mut_ptr()) 206 } 207 set_value(&self, value: T, permit: SemaphorePermit<'_>) -> &T208 fn set_value(&self, value: T, permit: SemaphorePermit<'_>) -> &T { 209 // SAFETY: We are holding the only permit on the semaphore. 210 unsafe { 211 self.value.with_mut(|ptr| (*ptr).as_mut_ptr().write(value)); 212 } 213 214 // Using release ordering so any threads that read a true from this 215 // atomic is able to read the value we just stored. 216 self.value_set.store(true, Ordering::Release); 217 self.semaphore.close(); 218 permit.forget(); 219 220 // SAFETY: We just initialized the cell. 221 unsafe { self.get_unchecked() } 222 } 223 224 /// Returns a reference to the value currently stored in the `OnceCell`, or 225 /// `None` if the `OnceCell` is empty. get(&self) -> Option<&T>226 pub fn get(&self) -> Option<&T> { 227 if self.initialized() { 228 Some(unsafe { self.get_unchecked() }) 229 } else { 230 None 231 } 232 } 233 234 /// Returns a mutable reference to the value currently stored in the 235 /// `OnceCell`, or `None` if the `OnceCell` is empty. 236 /// 237 /// Since this call borrows the `OnceCell` mutably, it is safe to mutate the 238 /// value inside the `OnceCell` — the mutable borrow statically guarantees 239 /// no other references exist. get_mut(&mut self) -> Option<&mut T>240 pub fn get_mut(&mut self) -> Option<&mut T> { 241 if self.initialized_mut() { 242 Some(unsafe { self.get_unchecked_mut() }) 243 } else { 244 None 245 } 246 } 247 248 /// Set the value of the `OnceCell` to the given value if the `OnceCell` is 249 /// empty. 250 /// 251 /// If the `OnceCell` already has a value, this call will fail with an 252 /// [`SetError::AlreadyInitializedError`]. 253 /// 254 /// If the `OnceCell` is empty, but some other task is currently trying to 255 /// set the value, this call will fail with [`SetError::InitializingError`]. 256 /// 257 /// [`SetError::AlreadyInitializedError`]: crate::sync::SetError::AlreadyInitializedError 258 /// [`SetError::InitializingError`]: crate::sync::SetError::InitializingError set(&self, value: T) -> Result<(), SetError<T>>259 pub fn set(&self, value: T) -> Result<(), SetError<T>> { 260 if self.initialized() { 261 return Err(SetError::AlreadyInitializedError(value)); 262 } 263 264 // Another task might be initializing the cell, in which case 265 // `try_acquire` will return an error. If we succeed to acquire the 266 // permit, then we can set the value. 267 match self.semaphore.try_acquire() { 268 Ok(permit) => { 269 debug_assert!(!self.initialized()); 270 self.set_value(value, permit); 271 Ok(()) 272 } 273 Err(TryAcquireError::NoPermits) => { 274 // Some other task is holding the permit. That task is 275 // currently trying to initialize the value. 276 Err(SetError::InitializingError(value)) 277 } 278 Err(TryAcquireError::Closed) => { 279 // The semaphore was closed. Some other task has initialized 280 // the value. 281 Err(SetError::AlreadyInitializedError(value)) 282 } 283 } 284 } 285 286 /// Get the value currently in the `OnceCell`, or initialize it with the 287 /// given asynchronous operation. 288 /// 289 /// If some other task is currently working on initializing the `OnceCell`, 290 /// this call will wait for that other task to finish, then return the value 291 /// that the other task produced. 292 /// 293 /// If the provided operation is cancelled or panics, the initialization 294 /// attempt is cancelled. If there are other tasks waiting for the value to 295 /// be initialized, one of them will start another attempt at initializing 296 /// the value. 297 /// 298 /// This will deadlock if `f` tries to initialize the cell recursively. get_or_init<F, Fut>(&self, f: F) -> &T where F: FnOnce() -> Fut, Fut: Future<Output = T>,299 pub async fn get_or_init<F, Fut>(&self, f: F) -> &T 300 where 301 F: FnOnce() -> Fut, 302 Fut: Future<Output = T>, 303 { 304 if self.initialized() { 305 // SAFETY: The OnceCell has been fully initialized. 306 unsafe { self.get_unchecked() } 307 } else { 308 // Here we try to acquire the semaphore permit. Holding the permit 309 // will allow us to set the value of the OnceCell, and prevents 310 // other tasks from initializing the OnceCell while we are holding 311 // it. 312 match self.semaphore.acquire().await { 313 Ok(permit) => { 314 debug_assert!(!self.initialized()); 315 316 // If `f()` panics or `select!` is called, this 317 // `get_or_init` call is aborted and the semaphore permit is 318 // dropped. 319 let value = f().await; 320 321 self.set_value(value, permit) 322 } 323 Err(_) => { 324 debug_assert!(self.initialized()); 325 326 // SAFETY: The semaphore has been closed. This only happens 327 // when the OnceCell is fully initialized. 328 unsafe { self.get_unchecked() } 329 } 330 } 331 } 332 } 333 334 /// Get the value currently in the `OnceCell`, or initialize it with the 335 /// given asynchronous operation. 336 /// 337 /// If some other task is currently working on initializing the `OnceCell`, 338 /// this call will wait for that other task to finish, then return the value 339 /// that the other task produced. 340 /// 341 /// If the provided operation returns an error, is cancelled or panics, the 342 /// initialization attempt is cancelled. If there are other tasks waiting 343 /// for the value to be initialized, one of them will start another attempt 344 /// at initializing the value. 345 /// 346 /// This will deadlock if `f` tries to initialize the cell recursively. get_or_try_init<E, F, Fut>(&self, f: F) -> Result<&T, E> where F: FnOnce() -> Fut, Fut: Future<Output = Result<T, E>>,347 pub async fn get_or_try_init<E, F, Fut>(&self, f: F) -> Result<&T, E> 348 where 349 F: FnOnce() -> Fut, 350 Fut: Future<Output = Result<T, E>>, 351 { 352 if self.initialized() { 353 // SAFETY: The OnceCell has been fully initialized. 354 unsafe { Ok(self.get_unchecked()) } 355 } else { 356 // Here we try to acquire the semaphore permit. Holding the permit 357 // will allow us to set the value of the OnceCell, and prevents 358 // other tasks from initializing the OnceCell while we are holding 359 // it. 360 match self.semaphore.acquire().await { 361 Ok(permit) => { 362 debug_assert!(!self.initialized()); 363 364 // If `f()` panics or `select!` is called, this 365 // `get_or_try_init` call is aborted and the semaphore 366 // permit is dropped. 367 let value = f().await; 368 369 match value { 370 Ok(value) => Ok(self.set_value(value, permit)), 371 Err(e) => Err(e), 372 } 373 } 374 Err(_) => { 375 debug_assert!(self.initialized()); 376 377 // SAFETY: The semaphore has been closed. This only happens 378 // when the OnceCell is fully initialized. 379 unsafe { Ok(self.get_unchecked()) } 380 } 381 } 382 } 383 } 384 385 /// Take the value from the cell, destroying the cell in the process. 386 /// Returns `None` if the cell is empty. into_inner(mut self) -> Option<T>387 pub fn into_inner(mut self) -> Option<T> { 388 if self.initialized_mut() { 389 // Set to uninitialized for the destructor of `OnceCell` to work properly 390 *self.value_set.get_mut() = false; 391 Some(unsafe { self.value.with(|ptr| ptr::read(ptr).assume_init()) }) 392 } else { 393 None 394 } 395 } 396 397 /// Takes ownership of the current value, leaving the cell empty. Returns 398 /// `None` if the cell is empty. take(&mut self) -> Option<T>399 pub fn take(&mut self) -> Option<T> { 400 std::mem::take(self).into_inner() 401 } 402 } 403 404 // Since `get` gives us access to immutable references of the OnceCell, OnceCell 405 // can only be Sync if T is Sync, otherwise OnceCell would allow sharing 406 // references of !Sync values across threads. We need T to be Send in order for 407 // OnceCell to by Sync because we can use `set` on `&OnceCell<T>` to send values 408 // (of type T) across threads. 409 unsafe impl<T: Sync + Send> Sync for OnceCell<T> {} 410 411 // Access to OnceCell's value is guarded by the semaphore permit 412 // and atomic operations on `value_set`, so as long as T itself is Send 413 // it's safe to send it to another thread 414 unsafe impl<T: Send> Send for OnceCell<T> {} 415 416 /// Errors that can be returned from [`OnceCell::set`]. 417 /// 418 /// [`OnceCell::set`]: crate::sync::OnceCell::set 419 #[derive(Debug, PartialEq)] 420 pub enum SetError<T> { 421 /// The cell was already initialized when [`OnceCell::set`] was called. 422 /// 423 /// [`OnceCell::set`]: crate::sync::OnceCell::set 424 AlreadyInitializedError(T), 425 426 /// The cell is currently being initialized. 427 InitializingError(T), 428 } 429 430 impl<T> fmt::Display for SetError<T> { fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result431 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 432 match self { 433 SetError::AlreadyInitializedError(_) => write!(f, "AlreadyInitializedError"), 434 SetError::InitializingError(_) => write!(f, "InitializingError"), 435 } 436 } 437 } 438 439 impl<T: fmt::Debug> Error for SetError<T> {} 440 441 impl<T> SetError<T> { 442 /// Whether `SetError` is `SetError::AlreadyInitializedError`. is_already_init_err(&self) -> bool443 pub fn is_already_init_err(&self) -> bool { 444 match self { 445 SetError::AlreadyInitializedError(_) => true, 446 SetError::InitializingError(_) => false, 447 } 448 } 449 450 /// Whether `SetError` is `SetError::InitializingError` is_initializing_err(&self) -> bool451 pub fn is_initializing_err(&self) -> bool { 452 match self { 453 SetError::AlreadyInitializedError(_) => false, 454 SetError::InitializingError(_) => true, 455 } 456 } 457 } 458