Skip to content

Commit

Permalink
replace parking_lot with parking_lot_core
Browse files Browse the repository at this point in the history
  • Loading branch information
matklad committed May 18, 2022
1 parent b57c774 commit 662fcf9
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 32 deletions.
4 changes: 3 additions & 1 deletion Cargo.toml
Expand Up @@ -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
Expand All @@ -48,6 +48,8 @@ race = []
# At the moment, this feature is unused.
unstable = []

parking_lot = ["parking_lot_core"]

[[example]]
name = "bench"
required-features = ["std"]
Expand Down
82 changes: 51 additions & 31 deletions src/imp_pl.rs
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -121,18 +112,47 @@ 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);
}
}

// 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) {
let key = state as *const _ as usize;
loop {
let exchange =
state.compare_exchange(INCOMPLETE, RUNNING, Ordering::Acquire, Ordering::Acquire);
match exchange {
Ok(_) => {
let mut guard = Guard { state, new_state: INCOMPLETE };
if init() {
guard.new_state = COMPLETE;
}
unsafe {
parking_lot_core::unpark_all(key, parking_lot_core::DEFAULT_UNPARK_TOKEN);
}
return;
}
Err(COMPLETE) => return,
Err(RUNNING) => unsafe {
parking_lot_core::park(
key,
|| state.load(Ordering::Relaxed) == RUNNING,
|| (),
|_, _| (),
parking_lot_core::DEFAULT_PARK_TOKEN,
None,
);
},
Err(_) => debug_assert!(false),
}
}
}
Expand All @@ -141,5 +161,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>());
}

0 comments on commit 662fcf9

Please sign in to comment.