From 43be119f1cf43663000eed472693a5e3cd72e9ff Mon Sep 17 00:00:00 2001 From: Markus Reiter Date: Wed, 7 Sep 2022 04:58:33 +0200 Subject: [PATCH] Change `critical-section` implementation to not use spin-lock. --- Cargo.toml | 3 ++ src/imp_cs.rs | 93 ++++++++++----------------------------------------- src/lib.rs | 1 + tests/it.rs | 42 +++++++++++++++++++---- 4 files changed, 56 insertions(+), 83 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index bb40a01..1d4f1f4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -88,3 +88,6 @@ required-features = ["std"] [package.metadata.docs.rs] all-features = true + +[patch.crates-io] +critical-section = { git = "https://github.com/reitermarkus/critical-section", branch = "lock-poisoning" } diff --git a/src/imp_cs.rs b/src/imp_cs.rs index 53848f8..56aadb0 100644 --- a/src/imp_cs.rs +++ b/src/imp_cs.rs @@ -1,23 +1,13 @@ -#[cfg(feature = "atomic-polyfill")] -use atomic_polyfill as atomic; -#[cfg(not(feature = "atomic-polyfill"))] -use core::sync::atomic; - -use atomic::{AtomicU8, Ordering}; - use core::panic::{RefUnwindSafe, UnwindSafe}; +use critical_section::{CriticalSection, Mutex}; + use crate::unsync::OnceCell as UnsyncOnceCell; pub(crate) struct OnceCell { - state: AtomicU8, - value: UnsyncOnceCell, + value: Mutex>, } -const INCOMPLETE: u8 = 0; -const RUNNING: u8 = 1; -const COMPLETE: u8 = 2; - // Why do we need `T: Send`? // Thread A creates a `OnceCell` and shares it with // scoped thread B, which fills the cell, which is @@ -31,16 +21,18 @@ impl UnwindSafe for OnceCell {} impl OnceCell { pub(crate) const fn new() -> OnceCell { - OnceCell { state: AtomicU8::new(INCOMPLETE), value: UnsyncOnceCell::new() } + OnceCell { value: Mutex::new(UnsyncOnceCell::new()) } } pub(crate) const fn with_value(value: T) -> OnceCell { - OnceCell { state: AtomicU8::new(COMPLETE), value: UnsyncOnceCell::with_value(value) } + OnceCell { value: Mutex::new(UnsyncOnceCell::with_value(value)) } } #[inline] pub(crate) fn is_initialized(&self) -> bool { - self.state.load(Ordering::Acquire) == COMPLETE + // The only write access is synchronized in `initialize` with a + // critical section, so read-only access is valid everywhere else. + unsafe { self.value.borrow(CriticalSection::new()).get().is_some() } } #[cold] @@ -48,30 +40,10 @@ impl OnceCell { where F: FnOnce() -> Result, { - let mut f = Some(f); - let mut res: Result<(), E> = Ok(()); - let slot: &UnsyncOnceCell = &self.value; - initialize_inner(&self.state, &mut || { - let f = unsafe { crate::take_unchecked(&mut f) }; - match f() { - Ok(value) => { - unsafe { crate::unwrap_unchecked(slot.set(value).ok()) }; - true - } - Err(err) => { - res = Err(err); - false - } - } - }); - res - } - - #[cold] - pub(crate) fn wait(&self) { - while !self.is_initialized() { - core::hint::spin_loop(); - } + critical_section::with(|cs| { + let cell = self.value.borrow(cs); + cell.get_or_try_init(f).map(|_| ()) + }) } /// Get the reference to the underlying value, without checking if the cell @@ -83,49 +55,18 @@ impl OnceCell { /// the contents are acquired by (synchronized to) this thread. pub(crate) unsafe fn get_unchecked(&self) -> &T { debug_assert!(self.is_initialized()); - self.value.get_unchecked() + // The only write access is synchronized in `initialize` with a + // critical section, so read-only access is valid everywhere else. + self.value.borrow(CriticalSection::new()).get_unchecked() } #[inline] pub(crate) fn get_mut(&mut self) -> Option<&mut T> { - self.value.get_mut() + self.value.get_mut().get_mut() } #[inline] pub(crate) fn into_inner(self) -> Option { - self.value.into_inner() - } -} - -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); - } -} - -#[inline(never)] -fn initialize_inner(state: &AtomicU8, init: &mut dyn FnMut() -> bool) { - loop { - match state.compare_exchange_weak(INCOMPLETE, RUNNING, Ordering::Acquire, Ordering::Acquire) - { - Ok(_) => { - let mut guard = Guard { state, new_state: INCOMPLETE }; - if init() { - guard.new_state = COMPLETE; - } - return; - } - Err(COMPLETE) => return, - Err(RUNNING) | Err(INCOMPLETE) => core::hint::spin_loop(), - Err(_) => { - debug_assert!(false); - unsafe { core::hint::unreachable_unchecked() } - } - } + self.value.into_inner().into_inner() } } diff --git a/src/lib.rs b/src/lib.rs index 11dcac6..64f09e5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -959,6 +959,7 @@ pub mod sync { /// assert_eq!(*value, 92); /// t.join().unwrap(); /// ``` + #[cfg(not(feature = "critical-section"))] pub fn wait(&self) -> &T { if !self.0.is_initialized() { self.0.wait() diff --git a/tests/it.rs b/tests/it.rs index 410b93b..5107e08 100644 --- a/tests/it.rs +++ b/tests/it.rs @@ -207,7 +207,7 @@ mod unsync { } #[test] - #[cfg(feature = "std")] + #[cfg(all(feature = "std", not(feature = "critical-section")))] fn lazy_poisoning() { let x: Lazy = Lazy::new(|| panic!("kaboom")); for _ in 0..2 { @@ -249,10 +249,13 @@ mod unsync { } } -#[cfg(feature = "std")] +#[cfg(feature = "sync")] mod sync { use std::sync::atomic::{AtomicUsize, Ordering::SeqCst}; + #[cfg(feature = "critical-section")] + use core::cell::Cell; + use crossbeam_utils::thread::scope; use once_cell::sync::{Lazy, OnceCell}; @@ -354,6 +357,7 @@ mod sync { assert_eq!(cell.get(), Some(&"hello".to_string())); } + #[cfg(not(feature = "critical-section"))] #[test] fn wait() { let cell: OnceCell = OnceCell::new(); @@ -365,6 +369,7 @@ mod sync { .unwrap(); } + #[cfg(not(feature = "critical-section"))] #[test] fn get_or_init_stress() { use std::sync::Barrier; @@ -430,6 +435,7 @@ mod sync { #[test] #[cfg_attr(miri, ignore)] // miri doesn't support processes + #[cfg(not(feature = "critical-section"))] fn reentrant_init() { let examples_dir = { let mut exe = std::env::current_exe().unwrap(); @@ -457,6 +463,20 @@ mod sync { } } + #[cfg(feature = "critical-section")] + #[test] + #[should_panic(expected = "reentrant init")] + fn reentrant_init() { + let x: OnceCell> = OnceCell::new(); + let dangling_ref: Cell> = Cell::new(None); + x.get_or_init(|| { + let r = x.get_or_init(|| Box::new(92)); + dangling_ref.set(Some(r)); + Box::new(62) + }); + eprintln!("use after free: {:?}", dangling_ref.get().unwrap()); + } + #[test] fn lazy_new() { let called = AtomicUsize::new(0); @@ -610,6 +630,7 @@ mod sync { assert_eq!(fib[5], 8) } + #[cfg(not(feature = "critical-section"))] #[test] fn once_cell_does_not_leak_partially_constructed_boxes() { let n_tries = if cfg!(miri) { 10 } else { 100 }; @@ -636,6 +657,7 @@ mod sync { } } + #[cfg(not(feature = "critical-section"))] #[test] fn get_does_not_block() { use std::sync::Barrier; @@ -671,12 +693,11 @@ mod sync { #[cfg(feature = "race")] mod race { + #[cfg(not(feature = "critical-section"))] + use std::sync::Barrier; use std::{ num::NonZeroUsize, - sync::{ - atomic::{AtomicUsize, Ordering::SeqCst}, - Barrier, - }, + sync::atomic::{AtomicUsize, Ordering::SeqCst}, }; use crossbeam_utils::thread::scope; @@ -728,6 +749,7 @@ mod race { assert_eq!(cell.get(), Some(val1)); } + #[cfg(not(feature = "critical-section"))] #[test] fn once_non_zero_usize_first_wins() { let val1 = NonZeroUsize::new(92).unwrap(); @@ -807,12 +829,16 @@ mod race { #[cfg(all(feature = "race", feature = "alloc"))] mod race_once_box { + #[cfg(not(feature = "critical-section"))] + use std::sync::Barrier; use std::sync::{ atomic::{AtomicUsize, Ordering::SeqCst}, - Arc, Barrier, + Arc, }; + #[cfg(not(feature = "critical-section"))] use crossbeam_utils::thread::scope; + use once_cell::race::OnceBox; #[derive(Default)] @@ -842,6 +868,7 @@ mod race_once_box { } } + #[cfg(not(feature = "critical-section"))] #[test] fn once_box_smoke_test() { let heap = Heap::default(); @@ -896,6 +923,7 @@ mod race_once_box { assert_eq!(heap.total(), 0); } + #[cfg(not(feature = "critical-section"))] #[test] fn once_box_first_wins() { let cell = OnceBox::new();