Merge #178

178: replace parking_lot with parking_lot_core r=matklad a=matklad

This is primaraliy to enable #177 for pl as well

Co-authored-by: Aleksey Kladov <aleksey.kladov@gmail.com>

CHANGELOG.md

@@ -1,9 +1,14 @@
 # Changelog
 
+## Unreleased
+
+- 
+
 ## 1.11
 
-- Add `OnceCell::with_value` to create initialized `OnceCell` at compile time.
+- Add `OnceCell::with_value` to create initialized `OnceCell` in `const` context.
 - Improve `Clone` implementation for `OnceCell`.
+- Rewrite `parking_lot` version on top of `parking_lot_core`, for even smaller cells!
 
 ## 1.10
 

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "once_cell"
-version = "1.10.0"
+version = "1.11.0"
 authors = ["Aleksey Kladov <aleksey.kladov@gmail.com>"]
 license = "MIT OR Apache-2.0"
 edition = "2018"
@@ -22,7 +22,7 @@ members = ["xtask"]
 # Uses parking_lot to implement once_cell::sync::OnceCell.
 # This makes not speed difference, but makes each OnceCell<T>
 # for up to 16 bytes smaller, depending on the size of the T.
-parking_lot = { version = "0.12", optional = true, default_features = false }
+parking_lot_core = { version = "0.9.3", optional = true, default_features = false }
 
 # To be used in order to enable the race feature on targets
 # that do not have atomics
@@ -48,6 +48,8 @@ race = []
 # At the moment, this feature is unused.
 unstable = []
 
+parking_lot = ["parking_lot_core"]
+
 [[example]]
 name = "bench"
 required-features = ["std"]

src/imp_pl.rs

@@ -2,19 +2,18 @@ use std::{
     cell::UnsafeCell,
     hint,
     panic::{RefUnwindSafe, UnwindSafe},
-    sync::atomic::{AtomicBool, Ordering},
+    sync::atomic::{AtomicU8, Ordering},
 };
 
-use parking_lot::Mutex;
-
-use crate::take_unchecked;
-
 pub(crate) struct OnceCell<T> {
-    mutex: Mutex<()>,
-    is_initialized: AtomicBool,
+    state: AtomicU8,
     value: UnsafeCell<Option<T>>,
 }
 
+const INCOMPLETE: u8 = 0x0;
+const RUNNING: u8 = 0x1;
+const COMPLETE: u8 = 0x2;
+
 // Why do we need `T: Send`?
 // Thread A creates a `OnceCell` and shares it with
 // scoped thread B, which fills the cell, which is
@@ -28,25 +27,17 @@ impl<T: UnwindSafe> UnwindSafe for OnceCell<T> {}
 
 impl<T> OnceCell<T> {
     pub(crate) const fn new() -> OnceCell<T> {
-        OnceCell {
-            mutex: parking_lot::const_mutex(()),
-            is_initialized: AtomicBool::new(false),
-            value: UnsafeCell::new(None),
-        }
+        OnceCell { state: AtomicU8::new(INCOMPLETE), value: UnsafeCell::new(None) }
     }
 
     pub(crate) const fn with_value(value: T) -> OnceCell<T> {
-        OnceCell {
-            mutex: parking_lot::const_mutex(()),
-            is_initialized: AtomicBool::new(true),
-            value: UnsafeCell::new(Some(value)),
-        }
+        OnceCell { state: AtomicU8::new(COMPLETE), value: UnsafeCell::new(Some(value)) }
     }
 
     /// Safety: synchronizes with store to value via Release/Acquire.
     #[inline]
     pub(crate) fn is_initialized(&self) -> bool {
-        self.is_initialized.load(Ordering::Acquire)
+        self.state.load(Ordering::Acquire) == COMPLETE
     }
 
     /// Safety: synchronizes with store to value via `is_initialized` or mutex
@@ -59,7 +50,7 @@ impl<T> OnceCell<T> {
         let mut f = Some(f);
         let mut res: Result<(), E> = Ok(());
         let slot: *mut Option<T> = self.value.get();
-        initialize_inner(&self.mutex, &self.is_initialized, &mut || {
+        initialize_inner(&self.state, &mut || {
             // We are calling user-supplied function and need to be careful.
             // - if it returns Err, we unlock mutex and return without touching anything
             // - if it panics, we unlock mutex and propagate panic without touching anything
@@ -68,7 +59,7 @@ impl<T> OnceCell<T> {
             //   but that is more complicated
             // - finally, if it returns Ok, we store the value and store the flag with
             //   `Release`, which synchronizes with `Acquire`s.
-            let f = unsafe { take_unchecked(&mut f) };
+            let f = unsafe { crate::take_unchecked(&mut f) };
             match f() {
                 Ok(value) => unsafe {
                     // Safe b/c we have a unique access and no panic may happen
@@ -121,18 +112,48 @@ impl<T> OnceCell<T> {
     }
 }
 
+struct Guard<'a> {
+    state: &'a AtomicU8,
+    new_state: u8,
+}
+
+impl<'a> Drop for Guard<'a> {
+    fn drop(&mut self) {
+        self.state.store(self.new_state, Ordering::Release);
+        unsafe {
+            let key = self.state as *const AtomicU8 as usize;
+            parking_lot_core::unpark_all(key, parking_lot_core::DEFAULT_UNPARK_TOKEN);
+        }
+    }
+}
+
 // Note: this is intentionally monomorphic
 #[inline(never)]
-fn initialize_inner(
-    mutex: &Mutex<()>,
-    is_initialized: &AtomicBool,
-    init: &mut dyn FnMut() -> bool,
-) {
-    let _guard = mutex.lock();
-
-    if !is_initialized.load(Ordering::Acquire) {
-        if init() {
-            is_initialized.store(true, Ordering::Release);
+fn initialize_inner(state: &AtomicU8, init: &mut dyn FnMut() -> bool) {
+    loop {
+        let exchange =
+            state.compare_exchange_weak(INCOMPLETE, RUNNING, Ordering::Acquire, Ordering::Acquire);
+        match exchange {
+            Ok(_) => {
+                let mut guard = Guard { state, new_state: INCOMPLETE };
+                if init() {
+                    guard.new_state = COMPLETE;
+                }
+                return;
+            }
+            Err(COMPLETE) => return,
+            Err(RUNNING) => unsafe {
+                let key = state as *const AtomicU8 as usize;
+                parking_lot_core::park(
+                    key,
+                    || state.load(Ordering::Relaxed) == RUNNING,
+                    || (),
+                    |_, _| (),
+                    parking_lot_core::DEFAULT_PARK_TOKEN,
+                    None,
+                );
+            },
+            Err(_) => debug_assert!(false),
         }
     }
 }
@@ -141,5 +162,5 @@ fn initialize_inner(
 fn test_size() {
     use std::mem::size_of;
 
-    assert_eq!(size_of::<OnceCell<bool>>(), 2 * size_of::<bool>() + size_of::<u8>());
+    assert_eq!(size_of::<OnceCell<bool>>(), 1 * size_of::<bool>() + size_of::<u8>());
 }

src/lib.rs

@@ -311,13 +311,18 @@
 //! At the moment, `unsync` has an additional benefit that reentrant initialization causes a panic, which
 //! might be easier to debug than a deadlock.
 //!
+//! **Does this crate support async?**
+//!
+//! No, but you can use [`async_once_cell`](https://crates.io/crates/async_once_cell) instead.
+//!
 //! # Related crates
 //!
 //! * [double-checked-cell](https://github.com/niklasf/double-checked-cell)
 //! * [lazy-init](https://crates.io/crates/lazy-init)
 //! * [lazycell](https://crates.io/crates/lazycell)
 //! * [mitochondria](https://crates.io/crates/mitochondria)
 //! * [lazy_static](https://crates.io/crates/lazy_static)
+//! * [async_once_cell](https://crates.io/crates/async_once_cell)
 //!
 //! Most of this crate's functionality is available in `std` in nightly Rust.
 //! See the [tracking issue](https://github.com/rust-lang/rust/issues/74465).

tests/it.rs

@@ -319,6 +319,28 @@ mod sync {
         assert_eq!(cell.get(), Some(&"hello".to_string()));
     }
 
+    #[test]
+    #[cfg_attr(miri, ignore)] // miri doesn't support Barrier
+    fn get_or_init_stress() {
+        use std::sync::Barrier;
+        let n_threads = 1_000;
+        let n_cells = 1_000;
+        let cells: Vec<_> = std::iter::repeat_with(|| (Barrier::new(n_threads), OnceCell::new()))
+            .take(n_cells)
+            .collect();
+        scope(|s| {
+            for _ in 0..n_threads {
+                s.spawn(|_| {
+                    for (i, (b, s)) in cells.iter().enumerate() {
+                        b.wait();
+                        let j = s.get_or_init(|| i);
+                        assert_eq!(*j, i);
+                    }
+                });
+            }
+        }).unwrap();
+    }
+
     #[test]
     fn from_impl() {
         assert_eq!(OnceCell::from("value").get(), Some(&"value"));