Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

replace parking_lot with parking_lot_core #178

Merged
merged 5 commits into from May 19, 2022
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
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
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This can probably be a plain Relaxed load.

The use of this API can only ever be used in an informational manner (it is 100% racy no matter how the returned value is used…) so it doesn't matter if it spuriously returns false for a little while when the once cell is already initialized.

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This one I am rather sure should be Acquire, as it guards get_unchecked accesses to the value.

once_cell/src/lib.rs

Lines 864 to 867 in b57c774

pub fn get(&self) -> Option<&T> {
if self.0.is_initialized() {
// Safe b/c value is initialized.
Some(unsafe { self.get_unchecked() })

}

/// 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);
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
state.compare_exchange(INCOMPLETE, RUNNING, Ordering::Acquire, Ordering::Acquire);
state.compare_exchange(INCOMPLETE, RUNNING, Ordering::Acquire, Ordering::Relaxed);

You could also consider compare_exchange_weak, though I don't know if there's any benefit in doing so since you will need to continue on Err(INCOMPLETE) anyway.

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hm, this I think also needs to be acquire, otherwise Err(COMPLETE) => return, branch won't synchronize-with the corresponding Release store, and subsequent get_unchecked would read garbage. Both cases are great to check if loom works though!

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>());
}