1use super::{Semaphore, SemaphorePermit, TryAcquireError};
2use crate::loom::cell::UnsafeCell;
3use std::error::Error;
4use std::fmt;
5use std::future::Future;
6use std::mem::MaybeUninit;
7use std::ops::Drop;
8use std::ptr;
9use std::sync::atomic::{AtomicBool, Ordering};
10
11// This file contains an implementation of an OnceCell. The principle
12// behind the safety 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/// ```
70pub struct OnceCell<T> {
71 value_set: AtomicBool,
72 value: UnsafeCell<MaybeUninit<T>>,
73 semaphore: Semaphore,
74}
75
76impl<T> Default for OnceCell<T> {
77 fn default() -> OnceCell<T> {
78 OnceCell::new()
79 }
80}
81
82impl<T: fmt::Debug> fmt::Debug for OnceCell<T> {
83 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
84 fmt&mut DebugStruct<'_, '_>.debug_struct("OnceCell")
85 .field(name:"value", &self.get())
86 .finish()
87 }
88}
89
90impl<T: Clone> Clone for OnceCell<T> {
91 fn clone(&self) -> OnceCell<T> {
92 OnceCell::new_with(self.get().cloned())
93 }
94}
95
96impl<T: PartialEq> PartialEq for OnceCell<T> {
97 fn eq(&self, other: &OnceCell<T>) -> bool {
98 self.get() == other.get()
99 }
100}
101
102impl<T: Eq> Eq for OnceCell<T> {}
103
104impl<T> Drop for OnceCell<T> {
105 fn drop(&mut self) {
106 if self.initialized_mut() {
107 unsafe {
108 self.value
109 .with_mut(|ptr: *mut MaybeUninit| ptr::drop_in_place((*ptr).as_mut_ptr()));
110 };
111 }
112 }
113}
114
115impl<T> From<T> for OnceCell<T> {
116 fn from(value: T) -> Self {
117 OnceCell {
118 value_set: AtomicBool::new(true),
119 value: UnsafeCell::new(data:MaybeUninit::new(val:value)),
120 semaphore: Semaphore::new_closed(),
121 }
122 }
123}
124
125impl<T> OnceCell<T> {
126 /// Creates a new empty `OnceCell` instance.
127 pub fn new() -> Self {
128 OnceCell {
129 value_set: AtomicBool::new(false),
130 value: UnsafeCell::new(MaybeUninit::uninit()),
131 semaphore: Semaphore::new(1),
132 }
133 }
134
135 /// Creates a new empty `OnceCell` instance.
136 ///
137 /// Equivalent to `OnceCell::new`, except that it can be used in static
138 /// variables.
139 ///
140 /// When using the `tracing` [unstable feature], a `OnceCell` created with
141 /// `const_new` will not be instrumented. As such, it will not be visible
142 /// in [`tokio-console`]. Instead, [`OnceCell::new`] should be used to
143 /// create an instrumented object if that is needed.
144 ///
145 /// # Example
146 ///
147 /// ```
148 /// use tokio::sync::OnceCell;
149 ///
150 /// static ONCE: OnceCell<u32> = OnceCell::const_new();
151 ///
152 /// async fn get_global_integer() -> &'static u32 {
153 /// ONCE.get_or_init(|| async {
154 /// 1 + 1
155 /// }).await
156 /// }
157 ///
158 /// #[tokio::main]
159 /// async fn main() {
160 /// let result = get_global_integer().await;
161 /// assert_eq!(*result, 2);
162 /// }
163 /// ```
164 ///
165 /// [`tokio-console`]: https://github.com/tokio-rs/console
166 /// [unstable feature]: crate#unstable-features
167 #[cfg(not(all(loom, test)))]
168 pub const fn const_new() -> Self {
169 OnceCell {
170 value_set: AtomicBool::new(false),
171 value: UnsafeCell::new(MaybeUninit::uninit()),
172 semaphore: Semaphore::const_new(1),
173 }
174 }
175
176 /// Creates a new `OnceCell` that contains the provided value, if any.
177 ///
178 /// If the `Option` is `None`, this is equivalent to `OnceCell::new`.
179 ///
180 /// [`OnceCell::new`]: crate::sync::OnceCell::new
181 // Once https://github.com/rust-lang/rust/issues/73255 lands
182 // and tokio MSRV is bumped to the rustc version with it stabilised,
183 // we can make this function available in const context,
184 // by creating `Semaphore::const_new_closed`.
185 pub fn new_with(value: Option<T>) -> Self {
186 if let Some(v) = value {
187 OnceCell::from(v)
188 } else {
189 OnceCell::new()
190 }
191 }
192
193 /// Creates a new `OnceCell` that contains the provided value.
194 ///
195 /// # Example
196 ///
197 /// When using the `tracing` [unstable feature], a `OnceCell` created with
198 /// `const_new_with` will not be instrumented. As such, it will not be
199 /// visible in [`tokio-console`]. Instead, [`OnceCell::new_with`] should be
200 /// used to create an instrumented object if that is needed.
201 ///
202 /// ```
203 /// use tokio::sync::OnceCell;
204 ///
205 /// static ONCE: OnceCell<u32> = OnceCell::const_new_with(1);
206 ///
207 /// async fn get_global_integer() -> &'static u32 {
208 /// ONCE.get_or_init(|| async {
209 /// 1 + 1
210 /// }).await
211 /// }
212 ///
213 /// #[tokio::main]
214 /// async fn main() {
215 /// let result = get_global_integer().await;
216 /// assert_eq!(*result, 1);
217 /// }
218 /// ```
219 ///
220 /// [`tokio-console`]: https://github.com/tokio-rs/console
221 /// [unstable feature]: crate#unstable-features
222 #[cfg(not(all(loom, test)))]
223 pub const fn const_new_with(value: T) -> Self {
224 OnceCell {
225 value_set: AtomicBool::new(true),
226 value: UnsafeCell::new(MaybeUninit::new(value)),
227 semaphore: Semaphore::const_new_closed(),
228 }
229 }
230
231 /// Returns `true` if the `OnceCell` currently contains a value, and `false`
232 /// otherwise.
233 pub fn initialized(&self) -> bool {
234 // Using acquire ordering so any threads that read a true from this
235 // atomic is able to read the value.
236 self.value_set.load(Ordering::Acquire)
237 }
238
239 /// Returns `true` if the `OnceCell` currently contains a value, and `false`
240 /// otherwise.
241 fn initialized_mut(&mut self) -> bool {
242 *self.value_set.get_mut()
243 }
244
245 // SAFETY: The OnceCell must not be empty.
246 unsafe fn get_unchecked(&self) -> &T {
247 &*self.value.with(|ptr| (*ptr).as_ptr())
248 }
249
250 // SAFETY: The OnceCell must not be empty.
251 unsafe fn get_unchecked_mut(&mut self) -> &mut T {
252 &mut *self.value.with_mut(|ptr| (*ptr).as_mut_ptr())
253 }
254
255 fn set_value(&self, value: T, permit: SemaphorePermit<'_>) -> &T {
256 // SAFETY: We are holding the only permit on the semaphore.
257 unsafe {
258 self.value.with_mut(|ptr| (*ptr).as_mut_ptr().write(value));
259 }
260
261 // Using release ordering so any threads that read a true from this
262 // atomic is able to read the value we just stored.
263 self.value_set.store(true, Ordering::Release);
264 self.semaphore.close();
265 permit.forget();
266
267 // SAFETY: We just initialized the cell.
268 unsafe { self.get_unchecked() }
269 }
270
271 /// Returns a reference to the value currently stored in the `OnceCell`, or
272 /// `None` if the `OnceCell` is empty.
273 pub fn get(&self) -> Option<&T> {
274 if self.initialized() {
275 Some(unsafe { self.get_unchecked() })
276 } else {
277 None
278 }
279 }
280
281 /// Returns a mutable reference to the value currently stored in the
282 /// `OnceCell`, or `None` if the `OnceCell` is empty.
283 ///
284 /// Since this call borrows the `OnceCell` mutably, it is safe to mutate the
285 /// value inside the `OnceCell` — the mutable borrow statically guarantees
286 /// no other references exist.
287 pub fn get_mut(&mut self) -> Option<&mut T> {
288 if self.initialized_mut() {
289 Some(unsafe { self.get_unchecked_mut() })
290 } else {
291 None
292 }
293 }
294
295 /// Sets the value of the `OnceCell` to the given value if the `OnceCell` is
296 /// empty.
297 ///
298 /// If the `OnceCell` already has a value, this call will fail with an
299 /// [`SetError::AlreadyInitializedError`].
300 ///
301 /// If the `OnceCell` is empty, but some other task is currently trying to
302 /// set the value, this call will fail with [`SetError::InitializingError`].
303 ///
304 /// [`SetError::AlreadyInitializedError`]: crate::sync::SetError::AlreadyInitializedError
305 /// [`SetError::InitializingError`]: crate::sync::SetError::InitializingError
306 pub fn set(&self, value: T) -> Result<(), SetError<T>> {
307 if self.initialized() {
308 return Err(SetError::AlreadyInitializedError(value));
309 }
310
311 // Another task might be initializing the cell, in which case
312 // `try_acquire` will return an error. If we succeed to acquire the
313 // permit, then we can set the value.
314 match self.semaphore.try_acquire() {
315 Ok(permit) => {
316 debug_assert!(!self.initialized());
317 self.set_value(value, permit);
318 Ok(())
319 }
320 Err(TryAcquireError::NoPermits) => {
321 // Some other task is holding the permit. That task is
322 // currently trying to initialize the value.
323 Err(SetError::InitializingError(value))
324 }
325 Err(TryAcquireError::Closed) => {
326 // The semaphore was closed. Some other task has initialized
327 // the value.
328 Err(SetError::AlreadyInitializedError(value))
329 }
330 }
331 }
332
333 /// Gets the value currently in the `OnceCell`, or initialize it with the
334 /// given asynchronous operation.
335 ///
336 /// If some other task is currently working on initializing the `OnceCell`,
337 /// this call will wait for that other task to finish, then return the value
338 /// that the other task produced.
339 ///
340 /// If the provided operation is cancelled or panics, the initialization
341 /// attempt is cancelled. If there are other tasks waiting for the value to
342 /// be initialized, one of them will start another attempt at initializing
343 /// the value.
344 ///
345 /// This will deadlock if `f` tries to initialize the cell recursively.
346 pub async fn get_or_init<F, Fut>(&self, f: F) -> &T
347 where
348 F: FnOnce() -> Fut,
349 Fut: Future<Output = T>,
350 {
351 crate::trace::async_trace_leaf().await;
352
353 if self.initialized() {
354 // SAFETY: The OnceCell has been fully initialized.
355 unsafe { self.get_unchecked() }
356 } else {
357 // Here we try to acquire the semaphore permit. Holding the permit
358 // will allow us to set the value of the OnceCell, and prevents
359 // other tasks from initializing the OnceCell while we are holding
360 // it.
361 match self.semaphore.acquire().await {
362 Ok(permit) => {
363 debug_assert!(!self.initialized());
364
365 // If `f()` panics or `select!` is called, this
366 // `get_or_init` call is aborted and the semaphore permit is
367 // dropped.
368 let value = f().await;
369
370 self.set_value(value, permit)
371 }
372 Err(_) => {
373 debug_assert!(self.initialized());
374
375 // SAFETY: The semaphore has been closed. This only happens
376 // when the OnceCell is fully initialized.
377 unsafe { self.get_unchecked() }
378 }
379 }
380 }
381 }
382
383 /// Gets the value currently in the `OnceCell`, or initialize it with the
384 /// given asynchronous operation.
385 ///
386 /// If some other task is currently working on initializing the `OnceCell`,
387 /// this call will wait for that other task to finish, then return the value
388 /// that the other task produced.
389 ///
390 /// If the provided operation returns an error, is cancelled or panics, the
391 /// initialization attempt is cancelled. If there are other tasks waiting
392 /// for the value to be initialized, one of them will start another attempt
393 /// at initializing the value.
394 ///
395 /// This will deadlock if `f` tries to initialize the cell recursively.
396 pub async fn get_or_try_init<E, F, Fut>(&self, f: F) -> Result<&T, E>
397 where
398 F: FnOnce() -> Fut,
399 Fut: Future<Output = Result<T, E>>,
400 {
401 crate::trace::async_trace_leaf().await;
402
403 if self.initialized() {
404 // SAFETY: The OnceCell has been fully initialized.
405 unsafe { Ok(self.get_unchecked()) }
406 } else {
407 // Here we try to acquire the semaphore permit. Holding the permit
408 // will allow us to set the value of the OnceCell, and prevents
409 // other tasks from initializing the OnceCell while we are holding
410 // it.
411 match self.semaphore.acquire().await {
412 Ok(permit) => {
413 debug_assert!(!self.initialized());
414
415 // If `f()` panics or `select!` is called, this
416 // `get_or_try_init` call is aborted and the semaphore
417 // permit is dropped.
418 let value = f().await;
419
420 match value {
421 Ok(value) => Ok(self.set_value(value, permit)),
422 Err(e) => Err(e),
423 }
424 }
425 Err(_) => {
426 debug_assert!(self.initialized());
427
428 // SAFETY: The semaphore has been closed. This only happens
429 // when the OnceCell is fully initialized.
430 unsafe { Ok(self.get_unchecked()) }
431 }
432 }
433 }
434 }
435
436 /// Takes the value from the cell, destroying the cell in the process.
437 /// Returns `None` if the cell is empty.
438 pub fn into_inner(mut self) -> Option<T> {
439 if self.initialized_mut() {
440 // Set to uninitialized for the destructor of `OnceCell` to work properly
441 *self.value_set.get_mut() = false;
442 Some(unsafe { self.value.with(|ptr| ptr::read(ptr).assume_init()) })
443 } else {
444 None
445 }
446 }
447
448 /// Takes ownership of the current value, leaving the cell empty. Returns
449 /// `None` if the cell is empty.
450 pub fn take(&mut self) -> Option<T> {
451 std::mem::take(self).into_inner()
452 }
453}
454
455// Since `get` gives us access to immutable references of the OnceCell, OnceCell
456// can only be Sync if T is Sync, otherwise OnceCell would allow sharing
457// references of !Sync values across threads. We need T to be Send in order for
458// OnceCell to by Sync because we can use `set` on `&OnceCell<T>` to send values
459// (of type T) across threads.
460unsafe impl<T: Sync + Send> Sync for OnceCell<T> {}
461
462// Access to OnceCell's value is guarded by the semaphore permit
463// and atomic operations on `value_set`, so as long as T itself is Send
464// it's safe to send it to another thread
465unsafe impl<T: Send> Send for OnceCell<T> {}
466
467/// Errors that can be returned from [`OnceCell::set`].
468///
469/// [`OnceCell::set`]: crate::sync::OnceCell::set
470#[derive(Debug, PartialEq, Eq)]
471pub enum SetError<T> {
472 /// The cell was already initialized when [`OnceCell::set`] was called.
473 ///
474 /// [`OnceCell::set`]: crate::sync::OnceCell::set
475 AlreadyInitializedError(T),
476
477 /// The cell is currently being initialized.
478 InitializingError(T),
479}
480
481impl<T> fmt::Display for SetError<T> {
482 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
483 match self {
484 SetError::AlreadyInitializedError(_) => write!(f, "AlreadyInitializedError"),
485 SetError::InitializingError(_) => write!(f, "InitializingError"),
486 }
487 }
488}
489
490impl<T: fmt::Debug> Error for SetError<T> {}
491
492impl<T> SetError<T> {
493 /// Whether `SetError` is `SetError::AlreadyInitializedError`.
494 pub fn is_already_init_err(&self) -> bool {
495 match self {
496 SetError::AlreadyInitializedError(_) => true,
497 SetError::InitializingError(_) => false,
498 }
499 }
500
501 /// Whether `SetError` is `SetError::InitializingError`
502 pub fn is_initializing_err(&self) -> bool {
503 match self {
504 SetError::AlreadyInitializedError(_) => false,
505 SetError::InitializingError(_) => true,
506 }
507 }
508}
509