1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811
//! A collection of lazy initialized values that are created by `Future`s.
//!
//! [OnceCell]'s API should be familiar to anyone who has used the
//! [`once_cell`](https://crates.io/crates/once_cell) crate or the proposed `std::cell::OnceCell`.
//! It provides an async version of a cell that can only be initialized once, permitting tasks to
//! wait on the initialization if it is already running instead of racing multiple initialization
//! tasks.
//!
//! Unlike threads, tasks can be cancelled at any point where they block. [OnceCell] deals with
//! this by allowing another initializer to run if the task currently initializing the cell is
//! dropped. This also allows for fallible initialization using [OnceCell::get_or_try_init], and
//! for the initializing `Future` to contain borrows or use references to thread-local data.
//!
//! [Lazy] takes the opposite approach: it wraps a single `Future` which is cooperatively run to
//! completion by any polling task. This requires that the initialization function be independent
//! of the calling context, but will never restart an initializing function just because the
//! surrounding task was cancelled.
//!
//! # Overhead
//!
//! Both cells use two `usize`s to store state and do not retain any allocations after
//! initialization is complete. [OnceCell] and [Lazy] only allocate if there is contention.
//!
//! # Features
//!
//! ## The `critical-section` feature
//!
//! If this feature is enabled, the [`critical-section`](https://crates.io/crates/critical-section)
//! crate is used instead of an `std` mutex. You must depend on that crate and select a locking
//! implementation; see [its documentation](https://docs.rs/critical-section/) for details.
//!
//! ## The `unpin` feature
//!
//! This feature enables the `unpin` module which contains an alternative API for [Lazy] that does
//! not rely on pinning the object during initialization, even for futures that are not [Unpin].
//! In general, prefer the types in the crate root and, if needed, box futures to make them unpin.
//!
//! ## The `std` feature
//!
//! This is currently a no-op, but might in the future be used to expose APIs that depends on
//! types only in `std`. It does *not* control the locking implementation.
#![cfg_attr(feature = "critical-section", no_std)]
extern crate alloc;
#[cfg(any(not(feature = "critical-section"), feature = "std"))]
extern crate std;
use alloc::{boxed::Box, vec, vec::Vec};
use core::{
cell::UnsafeCell,
convert::Infallible,
future::Future,
panic::{RefUnwindSafe, UnwindSafe},
pin::Pin,
ptr,
sync::atomic::{AtomicPtr, AtomicUsize, Ordering},
task,
};
#[cfg(feature = "critical-section")]
struct Mutex<T> {
data: UnsafeCell<T>,
locked: core::sync::atomic::AtomicBool,
}
#[cfg(feature = "critical-section")]
impl<T> Mutex<T> {
const fn new(data: T) -> Self {
Mutex { data: UnsafeCell::new(data), locked: core::sync::atomic::AtomicBool::new(false) }
}
}
#[cfg(not(feature = "critical-section"))]
use std::sync::Mutex;
#[cfg(feature = "critical-section")]
fn with_lock<T, R>(mutex: &Mutex<T>, f: impl FnOnce(&mut T) -> R) -> R {
struct Guard<'a, T>(&'a Mutex<T>);
impl<'a, T> Drop for Guard<'a, T> {
fn drop(&mut self) {
self.0.locked.store(false, Ordering::Relaxed);
}
}
critical_section::with(|_| {
if mutex.locked.swap(true, Ordering::Relaxed) {
// Note: this can in theory happen if the delegated Clone impl on a Waker provided in
// an initialization context turns around and tries to initialize the same cell. This
// is an absurd thing to do, but it's safe so we can't assume nobody will ever do it.
panic!("Attempted reentrant locking");
}
let guard = Guard(mutex);
// Safety: we just checked that we were the one to set `locked` to true, and the data in
// this Mutex will only be accessed while the lock is true. We use Relaxed memory ordering
// instead of Acquire/Release because critical_section::with itself must provide an
// Acquire/Release barrier around its closure, and also guarantees that there will not be
// more than one such closure executing at a time.
let rv = unsafe { f(&mut *mutex.data.get()) };
drop(guard);
rv
})
}
#[cfg(not(feature = "critical-section"))]
fn with_lock<T, R>(mutex: &Mutex<T>, f: impl FnOnce(&mut T) -> R) -> R {
f(&mut *mutex.lock().unwrap())
}
/// Types that do not rely on pinning during initialization.
///
/// This module is only built if the `unpin` crate feature is enabled.
///
/// This module contains [OnceFuture](unpin::OnceFuture) and its wrappers [Lazy](unpin::Lazy) and
/// [ConstLazy](unpin::ConstLazy), which provide lazy initialization without requiring the
/// resulting structure be pinned.
///
/// This is the API exposed by the 0.3 version of this crate for `Lazy`.
#[cfg(feature = "unpin")]
pub mod unpin;
/// A cell which can be written to only once.
///
/// This allows initialization using an async closure that borrows from its environment.
///
/// ```
/// # async fn run() {
/// use std::rc::Rc;
/// use std::sync::Arc;
/// use async_once_cell::OnceCell;
///
/// let non_send_value = Rc::new(4);
/// let shared = Arc::new(OnceCell::new());
///
/// let value : &i32 = shared.get_or_init(async {
/// *non_send_value
/// }).await;
/// assert_eq!(value, &4);
///
/// // A second init is not called
/// let second = shared.get_or_init(async {
/// unreachable!()
/// }).await;
/// assert_eq!(second, &4);
///
/// # }
/// ```
#[derive(Debug)]
pub struct OnceCell<T> {
value: UnsafeCell<Option<T>>,
inner: Inner,
}
// Safety: our UnsafeCell should be treated like an RwLock<T>
unsafe impl<T: Sync + Send> Sync for OnceCell<T> {}
unsafe impl<T: Send> Send for OnceCell<T> {}
impl<T> Unpin for OnceCell<T> {}
impl<T: RefUnwindSafe + UnwindSafe> RefUnwindSafe for OnceCell<T> {}
impl<T: UnwindSafe> UnwindSafe for OnceCell<T> {}
/// Monomorphic portion of the state
#[derive(Debug)]
struct Inner {
state: AtomicUsize,
queue: AtomicPtr<Queue>,
}
/// Transient state during initialization
///
/// Unlike the sync OnceCell, this cannot be a linked list through stack frames, because Futures
/// can be freed at any point by any thread. Instead, this structure is allocated on the heap
/// during the first initialization call and freed after the value is set (or when the OnceCell is
/// dropped, if the value never gets set).
struct Queue {
wakers: Mutex<Option<Vec<task::Waker>>>,
}
/// This is somewhat like Arc<Queue>, but holds the refcount in Inner instead of Queue so it can be
/// freed once the cell's initialization is complete.
struct QueueRef<'a> {
inner: &'a Inner,
queue: *const Queue,
}
// Safety: the queue is a reference (only the lack of a valid lifetime requires it to be a pointer)
unsafe impl<'a> Sync for QueueRef<'a> {}
unsafe impl<'a> Send for QueueRef<'a> {}
#[derive(Debug)]
struct QuickInitGuard<'a>(&'a Inner);
/// A Future that waits for acquisition of a QueueHead
struct QueueWaiter<'a> {
guard: Option<QueueRef<'a>>,
}
/// A guard for the actual initialization of the OnceCell
struct QueueHead<'a> {
guard: QueueRef<'a>,
}
const NEW: usize = 0x0;
const QINIT_BIT: usize = 1 + (usize::MAX >> 2);
const READY_BIT: usize = 1 + (usize::MAX >> 1);
impl Inner {
const fn new() -> Self {
Inner { state: AtomicUsize::new(NEW), queue: AtomicPtr::new(ptr::null_mut()) }
}
const fn new_ready() -> Self {
Inner { state: AtomicUsize::new(READY_BIT), queue: AtomicPtr::new(ptr::null_mut()) }
}
/// Initialize the queue (if needed) and return a waiter that can be polled to get a QueueHead
/// that gives permission to initialize the OnceCell.
///
/// The Queue referenced in the returned QueueRef will not be freed until the cell is populated
/// and all references have been dropped. If any references remain, further calls to
/// initialize will return the existing queue.
#[cold]
fn initialize(&self, try_quick: bool) -> Result<QueueWaiter, QuickInitGuard> {
if try_quick {
if self
.state
.compare_exchange(NEW, QINIT_BIT, Ordering::Acquire, Ordering::Relaxed)
.is_ok()
{
// On success, we know that there were no other QueueRef objects active, and we
// just set QINIT_BIT which makes us the only party allowed to create a QueueHead.
// This remains true even if the queue is created later.
return Err(QuickInitGuard(self));
}
}
// Increment the queue's reference count. This ensures that queue won't be freed until we exit.
let prev_state = self.state.fetch_add(1, Ordering::Acquire);
// Note: unlike Arc, refcount overflow is impossible. The only way to increment the
// refcount is by calling poll on the Future returned by get_or_try_init, which is !Unpin.
// The poll call requires a Pinned pointer to this Future, and the contract of Pin requires
// Drop to be called on any !Unpin value that was pinned before the memory is reused.
// Because the Drop impl of QueueRef decrements the refcount, an overflow would require
// more than (usize::MAX / 4) QueueRef objects in memory, which is impossible as these
// objects take up more than 4 bytes.
let mut guard = QueueRef { inner: self, queue: self.queue.load(Ordering::Acquire) };
if guard.queue.is_null() && prev_state & READY_BIT == 0 {
let wakers = Mutex::new(None);
// Race with other callers of initialize to create the queue
let new_queue = Box::into_raw(Box::new(Queue { wakers }));
match self.queue.compare_exchange(
ptr::null_mut(),
new_queue,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_null) => {
// Normal case: it was actually set. The Release part of AcqRel orders this
// with all Acquires on the queue.
guard.queue = new_queue;
}
Err(actual) => {
// we lost the race, but we have the (non-null) value now.
guard.queue = actual;
// Safety: we just allocated it, and nobody else has seen it
unsafe {
drop(Box::from_raw(new_queue));
}
}
}
}
Ok(QueueWaiter { guard: Some(guard) })
}
fn set_ready(&self) {
// This Release pairs with the Acquire any time we check READY_BIT, and ensures that the
// writes to the cell's value are visible to the cell's readers.
let prev_state = self.state.fetch_or(READY_BIT, Ordering::Release);
debug_assert_eq!(prev_state & READY_BIT, 0, "Invalid state: someone else set READY_BIT");
}
}
impl<'a> Drop for QueueRef<'a> {
fn drop(&mut self) {
// Release the reference to queue
let prev_state = self.inner.state.fetch_sub(1, Ordering::Release);
// Note: as of now, self.queue may be invalid
let curr_state = prev_state - 1;
if curr_state == READY_BIT || curr_state == READY_BIT | QINIT_BIT {
// We just removed the only waiter on an initialized cell. This means the
// queue is no longer needed. Acquire the queue again so we can free it.
let queue = self.inner.queue.swap(ptr::null_mut(), Ordering::Acquire);
if !queue.is_null() {
// Safety: the last guard is being freed, and queue is only used by guard-holders.
// Due to the swap, we are the only one who is freeing this particular queue.
unsafe {
drop(Box::from_raw(queue));
}
}
}
}
}
impl<'a> Drop for QuickInitGuard<'a> {
fn drop(&mut self) {
let prev_state = self.0.state.load(Ordering::Relaxed);
if prev_state == QINIT_BIT | READY_BIT || prev_state == QINIT_BIT {
let target = prev_state & !QINIT_BIT;
// Try to finish the fast path of initialization if possible.
if self
.0
.state
.compare_exchange(prev_state, target, Ordering::Relaxed, Ordering::Relaxed)
.is_ok()
{
// If init succeeded, the Release in set_ready already ordered the value. If init
// failed, we made no writes that need to be ordered and there are no waiters to
// wake, so we can leave the state at NEW.
if target == READY_BIT {
// It's possible (though unlikely) that someone created the queue but abandoned
// their QueueRef before we finished our poll, resulting in us not observing
// them. No wakes are needed in this case because there are no waiting tasks,
// but we should still clean up the allocation.
let queue = self.0.queue.swap(ptr::null_mut(), Ordering::Relaxed);
if !queue.is_null() {
// Synchronize with both the fetch_sub that lowered the refcount and the
// queue initialization.
core::sync::atomic::fence(Ordering::Acquire);
// Safety: we observed no active QueueRefs, and queue is only used by
// guard-holders. Due to the swap, we are the only one who is freeing this
// particular queue.
unsafe {
drop(Box::from_raw(queue));
}
}
}
return;
}
}
// Slow path: get a guard, create the QueueHead we should have been holding, then drop it
// so that the tasks are woken as intended. This is needed regardless of if we succeeded
// or not - either waiters need to run init themselves, or they need to read the value we
// set.
//
// The guard is guaranteed to have been created with no QueueHead available because
// QINIT_BIT is still set.
let waiter = self.0.initialize(false).expect("Got a QuickInitGuard in slow init");
let guard = waiter.guard.expect("No guard available even without polling");
if guard.queue.is_null() {
// The queue was already freed by someone else before we got our QueueRef (this must
// have happened between the load of prev_state and initialize, because otherwise we
// would have taken the fast path). This implies that all other tasks have noticed
// READY_BIT and do not need waking, so there is nothing left for us to do except
// release our reference.
drop(guard);
} else {
// Safety: the guard holds a place on the waiter list and we just checked that the
// queue is non-null. It will remain valid until guard is dropped.
let queue = unsafe { &*guard.queue };
with_lock(&queue.wakers, |lock| {
// Ensure that nobody else can grab the QueueHead between when we release QINIT_BIT
// and when our QueueHead is dropped.
lock.get_or_insert_with(Vec::new);
// Allow someone else to take the head position once we drop it. Ordering is
// handled by the Mutex.
self.0.state.fetch_and(!QINIT_BIT, Ordering::Relaxed);
});
// Safety: we just took the head position, and we were the QuickInitGuard
drop(QueueHead { guard })
}
}
}
impl Drop for Inner {
fn drop(&mut self) {
let queue = *self.queue.get_mut();
if !queue.is_null() {
// Safety: nobody else could have a reference
unsafe {
drop(Box::from_raw(queue));
}
}
}
}
impl<'a> Future for QueueWaiter<'a> {
type Output = Option<QueueHead<'a>>;
fn poll(
mut self: Pin<&mut Self>,
cx: &mut task::Context<'_>,
) -> task::Poll<Option<QueueHead<'a>>> {
let guard = self.guard.as_ref().expect("Polled future after finished");
// Fast path for waiters that get notified after the value is set
let state = guard.inner.state.load(Ordering::Acquire);
if state & READY_BIT != 0 {
return task::Poll::Ready(None);
}
// Safety: the guard holds a place on the waiter list and we just checked that the state is
// not ready, so the queue is non-null and will remain valid until guard is dropped.
let queue = unsafe { &*guard.queue };
let rv = with_lock(&queue.wakers, |lock| {
// Another task might have called set_ready() and dropped its QueueHead between our
// optimistic lock-free check and our lock acquisition. Don't return a QueueHead unless we
// know for sure that we are allowed to initialize.
let state = guard.inner.state.load(Ordering::Acquire);
if state & READY_BIT != 0 {
return task::Poll::Ready(None);
}
match lock.as_mut() {
None if state & QINIT_BIT == 0 => {
// take the head position and start a waker queue
*lock = Some(Vec::new());
task::Poll::Ready(Some(()))
}
None => {
// Someone else has a QuickInitGuard; they will wake us when they finish.
let waker = cx.waker().clone();
*lock = Some(vec![waker]);
task::Poll::Pending
}
Some(wakers) => {
// Wait for the QueueHead to be dropped
let my_waker = cx.waker();
for waker in wakers.iter() {
if waker.will_wake(my_waker) {
return task::Poll::Pending;
}
}
wakers.push(my_waker.clone());
task::Poll::Pending
}
}
});
// Safety: If rv is Ready/Some, we know:
// - we are holding a QueueRef (in guard) that prevents state from being 0
// - creating a new QuickInitGuard requires the state to be 0
// - we just checked QINIT_BIT and saw there isn't a QuickInitGuard active
// - the queue was None, meaning there are no current QueueHeads
// - we just set the queue to Some, claiming the head
//
// If rv is Ready/None, this is due to READY_BIT being set.
// If rv is Pending, we have a waker in the queue.
rv.map(|o| o.map(|()| QueueHead { guard: self.guard.take().unwrap() }))
}
}
impl<'a> Drop for QueueHead<'a> {
fn drop(&mut self) {
// Safety: if queue is not null, then it is valid as long as the guard is alive
if let Some(queue) = unsafe { self.guard.queue.as_ref() } {
// Take the waker queue so the next QueueWaiter can make a new one
let wakers = with_lock(&queue.wakers, |w| w.take())
.expect("QueueHead dropped without a waker list");
for waker in wakers {
waker.wake();
}
}
}
}
impl<T> OnceCell<T> {
/// Creates a new empty cell.
pub const fn new() -> Self {
Self { value: UnsafeCell::new(None), inner: Inner::new() }
}
/// Creates a new cell with the given contents.
///
/// If value is `None`, this is equivalent to `new()`.
pub const fn new_with(value: Option<T>) -> Self {
let inner = match value {
Some(_) => Inner::new_ready(),
None => Inner::new(),
};
Self { value: UnsafeCell::new(value), inner }
}
/// Gets the contents of the cell, initializing it with `init` if the cell was empty.
///
/// Many tasks may call `get_or_init` concurrently with different initializing futures, but
/// it is guaranteed that only one future will be executed as long as the resulting future is
/// polled to completion.
///
/// If `init` panics, the panic is propagated to the caller, and the cell remains uninitialized.
///
/// If the Future returned by this function is dropped prior to completion, the cell remains
/// uninitialized, and another `init` function will be started (if any are available).
///
/// It is an error to reentrantly initialize the cell from `init`. The current implementation
/// deadlocks, but will recover if the offending task is dropped or if the future is actually
/// able to proceed despite the reentrant call never returning.
pub async fn get_or_init(&self, init: impl Future<Output = T>) -> &T {
match self.get_or_try_init(async move { Ok::<T, Infallible>(init.await) }).await {
Ok(t) => t,
Err(e) => match e {},
}
}
/// Gets the contents of the cell, initializing it with `init` if the cell was empty. If the
/// cell was empty and `init` failed, an error is returned.
///
/// If `init` panics or returns an error, the panic or error is propagated to the caller, and
/// the cell remains uninitialized. In this case, another `init` function from a concurrent
/// caller will be selected to execute, if one is available.
///
/// If the Future returned by this function is dropped prior to completion, the cell remains
/// uninitialized, and another `init` function will be started.
///
/// It is an error to reentrantly initialize the cell from `init`. The current implementation
/// deadlocks, but will recover if the offending task is dropped or if the future is actually
/// able to proceed despite the reentrant call never returning.
pub async fn get_or_try_init<E>(
&self,
init: impl Future<Output = Result<T, E>>,
) -> Result<&T, E> {
let state = self.inner.state.load(Ordering::Acquire);
if state & READY_BIT == 0 {
self.init_slow(state == NEW, init).await?;
}
// Safety: initialized on all paths
Ok(unsafe { (&*self.value.get()).as_ref().unwrap() })
}
#[cold]
async fn init_slow<E>(
&self,
try_quick: bool,
init: impl Future<Output = Result<T, E>>,
) -> Result<(), E> {
match self.inner.initialize(try_quick) {
Err(guard) => {
// Try to proceed assuming no contention.
let value = init.await?;
// Safety: the guard acts like QueueHead even if there is contention.
unsafe {
*self.value.get() = Some(value);
}
self.inner.set_ready();
drop(guard);
}
Ok(guard) => {
if let Some(init_lock) = guard.await {
// We hold the QueueHead, so we know that nobody else has successfully run an init
// poll and that nobody else can start until it is dropped. On error, panic, or
// drop of this Future, the head will be passed to another waiter.
let value = init.await?;
// Safety: We still hold the head, so nobody else can write to value
unsafe {
*self.value.get() = Some(value);
}
// mark the cell ready before giving up the head
init_lock.guard.inner.set_ready();
// drop of QueueHead notifies other Futures
// drop of QueueRef (might) free the Queue
} else {
// someone initialized it while waiting on the queue
}
}
}
Ok(())
}
/// Gets the reference to the underlying value.
///
/// Returns `None` if the cell is empty or being initialized. This method never blocks.
pub fn get(&self) -> Option<&T> {
let state = self.inner.state.load(Ordering::Acquire);
if state & READY_BIT == 0 {
None
} else {
unsafe { (&*self.value.get()).as_ref() }
}
}
/// Gets a mutable reference to the underlying value.
pub fn get_mut(&mut self) -> Option<&mut T> {
self.value.get_mut().as_mut()
}
/// Takes the value out of this `OnceCell`, moving it back to an uninitialized state.
pub fn take(&mut self) -> Option<T> {
self.inner = Inner::new();
self.value.get_mut().take()
}
/// Consumes the OnceCell, returning the wrapped value. Returns None if the cell was empty.
pub fn into_inner(self) -> Option<T> {
self.value.into_inner()
}
}
impl<T> Default for OnceCell<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> From<T> for OnceCell<T> {
fn from(value: T) -> Self {
Self::new_with(Some(value))
}
}
#[derive(Debug)]
enum LazyState<T, F> {
Running(F),
Ready(T),
}
/// A value which is computed on demand by running a future.
///
/// Unlike [OnceCell], if a task is cancelled, the initializing future's execution will be
/// continued by other (concurrent or future) callers of [Lazy::get].
///
/// ```
/// # async fn run() {
/// use std::sync::Arc;
/// use async_once_cell::Lazy;
///
/// struct Data {
/// id: u32,
/// }
///
/// let shared = Arc::pin(Lazy::new(async move {
/// Data { id: 4 }
/// }));
///
/// assert_eq!(shared.as_ref().get().await.id, 4);
/// # }
/// ```
#[derive(Debug)]
pub struct Lazy<T, F> {
value: UnsafeCell<LazyState<T, F>>,
inner: Inner,
}
// Safety: our UnsafeCell should be treated like an RwLock<(T, F)>
unsafe impl<T: Sync + Send, F: Sync + Send> Sync for Lazy<T, F> {}
unsafe impl<T: Send, F: Send> Send for Lazy<T, F> {}
impl<T: Unpin, F: Unpin> Unpin for Lazy<T, F> {}
impl<T: RefUnwindSafe + UnwindSafe, F: RefUnwindSafe + UnwindSafe> RefUnwindSafe for Lazy<T, F> {}
impl<T: UnwindSafe, F: UnwindSafe> UnwindSafe for Lazy<T, F> {}
impl<T, F> Lazy<T, F>
where
F: Future<Output = T>,
{
/// Creates a new lazy value with the given initializing future.
pub const fn new(future: F) -> Self {
Self::from_future(future)
}
/// Forces the evaluation of this lazy value and returns a reference to the result.
///
/// The [Pin::static_ref] function may be useful if this is a static value.
pub async fn get(self: Pin<&Self>) -> Pin<&T> {
let state = self.inner.state.load(Ordering::Acquire);
if state & READY_BIT == 0 {
self.init_slow(state == NEW).await;
}
// Safety: initialized on all paths, and pinned like self
unsafe {
match &*self.value.get() {
LazyState::Ready(v) => Pin::new_unchecked(v),
_ => unreachable!(),
}
}
}
#[cold]
async fn init_slow(self: Pin<&Self>, try_quick: bool) {
match self.inner.initialize(try_quick) {
Err(guard) => {
let init = unsafe {
match &mut *self.value.get() {
LazyState::Running(f) => Pin::new_unchecked(f),
_ => unreachable!(),
}
};
let value = init.await;
// Safety: the guard acts like QueueHead even if there is contention.
// This overwrites the pinned future, dropping it in place
unsafe {
*self.value.get() = LazyState::Ready(value);
}
self.inner.set_ready();
drop(guard);
}
Ok(guard) => {
if let Some(init_lock) = guard.await {
let init = unsafe {
match &mut *self.value.get() {
LazyState::Running(f) => Pin::new_unchecked(f),
_ => unreachable!(),
}
};
// We hold the QueueHead, so we know that nobody else has successfully run an init
// poll and that nobody else can start until it is dropped. On error, panic, or
// drop of this Future, the head will be passed to another waiter.
let value = init.await;
// Safety: We still hold the head, so nobody else can write to value
// This overwrites the pinned future, dropping it in place
unsafe {
*self.value.get() = LazyState::Ready(value);
}
// mark the cell ready before giving up the head
init_lock.guard.inner.set_ready();
// drop of QueueHead notifies other Futures
// drop of QueueRef (might) free the Queue
} else {
// someone initialized it while waiting on the queue
}
}
}
}
}
impl<T, F> Lazy<T, F>
where
F: Future<Output = T> + Unpin,
{
/// Forces the evaluation of this lazy value and returns a reference to the result.
///
/// Unlike [Self::get], this does not require pinning the object.
pub async fn get_unpin(&self) -> &T {
// The get() function itself does not use the fact that T is pinned, and Pin::deref already
// exposes a &T from Pin<&T> (although not with the right lifetime).
unsafe { Pin::into_inner_unchecked(Pin::new_unchecked(self).get().await) }
}
}
impl<T, F> Lazy<T, F> {
/// Creates a new lazy value with the given initializing future.
///
/// This is equivalent to [Self::new] but with no type bound.
pub const fn from_future(future: F) -> Self {
Self { value: UnsafeCell::new(LazyState::Running(future)), inner: Inner::new() }
}
/// Creates an already-initialized lazy value.
pub const fn with_value(value: T) -> Self {
Self { value: UnsafeCell::new(LazyState::Ready(value)), inner: Inner::new_ready() }
}
/// Gets the value without blocking or starting the initialization.
pub fn try_get(&self) -> Option<&T> {
let state = self.inner.state.load(Ordering::Acquire);
if state & READY_BIT == 0 {
None
} else {
match unsafe { &*self.value.get() } {
LazyState::Ready(v) => Some(v),
_ => unreachable!(),
}
}
}
/// Gets the value without blocking or starting the initialization.
///
/// This requires mutable access to self, so rust's aliasing rules prevent any concurrent
/// access and allow violating the usual rules for accessing this cell.
pub fn try_get_mut(self: Pin<&mut Self>) -> Option<Pin<&mut T>> {
unsafe {
match self.get_unchecked_mut().value.get_mut() {
LazyState::Ready(v) => Some(Pin::new_unchecked(v)),
_ => None,
}
}
}
/// Gets the value without blocking or starting the initialization.
///
/// This requires mutable access to self, so rust's aliasing rules prevent any concurrent
/// access and allow violating the usual rules for accessing this cell.
pub fn try_get_mut_unpin(&mut self) -> Option<&mut T> {
match self.value.get_mut() {
LazyState::Ready(v) => Some(v),
_ => None,
}
}
/// Gets the value if it was set.
pub fn into_inner(self) -> Option<T> {
match self.value.into_inner() {
LazyState::Ready(v) => Some(v),
_ => None,
}
}
}