From cf11f499b37e557e41f5eeac157c623d08e0068d Mon Sep 17 00:00:00 2001 From: joboet Date: Wed, 10 Jul 2024 14:12:58 +0200 Subject: [PATCH 1/2] std: implement the `once_wait` feature --- library/std/src/sync/once.rs | 41 ++++++ library/std/src/sync/once_lock.rs | 28 ++++ library/std/src/sys/sync/once/futex.rs | 124 ++++++++++++----- library/std/src/sys/sync/once/no_threads.rs | 6 + library/std/src/sys/sync/once/queue.rs | 142 ++++++++++++-------- 5 files changed, 247 insertions(+), 94 deletions(-) diff --git a/library/std/src/sync/once.rs b/library/std/src/sync/once.rs index 9d969af8c6d84..17c56c8501658 100644 --- a/library/std/src/sync/once.rs +++ b/library/std/src/sync/once.rs @@ -264,6 +264,47 @@ impl Once { self.inner.is_completed() } + /// Blocks the current thread until initialization has completed. + /// + /// # Example + /// + /// ```rust + /// #![feature(once_wait)] + /// + /// use std::sync::Once; + /// use std::thread; + /// + /// static READY: Once = Once::new(); + /// + /// let thread = thread::spawn(|| { + /// READY.wait(); + /// println!("everything is ready"); + /// }); + /// + /// READY.call_once(|| println!("performing setup")); + /// ``` + /// + /// # Panics + /// + /// If this [`Once`] has been poisoned because an initialization closure has + /// panicked, this method will also panic. Use [`wait_force`](Self::wait_force) + /// if this behaviour is not desired. + #[unstable(feature = "once_wait", issue = "127527")] + pub fn wait(&self) { + if !self.inner.is_completed() { + self.inner.wait(false); + } + } + + /// Blocks the current thread until initialization has completed, ignoring + /// poisoning. + #[unstable(feature = "once_wait", issue = "127527")] + pub fn wait_force(&self) { + if !self.inner.is_completed() { + self.inner.wait(true); + } + } + /// Returns the current state of the `Once` instance. /// /// Since this takes a mutable reference, no initialization can currently diff --git a/library/std/src/sync/once_lock.rs b/library/std/src/sync/once_lock.rs index 60e43a1cde10e..efc1f415edfc1 100644 --- a/library/std/src/sync/once_lock.rs +++ b/library/std/src/sync/once_lock.rs @@ -167,6 +167,34 @@ impl OnceLock { } } + /// Blocks the current thread until the cell is initialized. + /// + /// # Example + /// + /// Waiting for a computation on another thread to finish: + /// ```rust + /// #![feature(once_wait)] + /// + /// use std::thread; + /// use std::sync::OnceLock; + /// + /// let value = OnceLock::new(); + /// + /// thread::scope(|s| { + /// s.spawn(|| value.set(1 + 1)); + /// + /// let result = value.wait(); + /// assert_eq!(result, &2); + /// }) + /// ``` + #[inline] + #[unstable(feature = "once_wait", issue = "127527")] + pub fn wait(&self) -> &T { + self.once.wait_force(); + + unsafe { self.get_unchecked() } + } + /// Sets the contents of this cell to `value`. /// /// May block if another thread is currently attempting to initialize the cell. The cell is diff --git a/library/std/src/sys/sync/once/futex.rs b/library/std/src/sys/sync/once/futex.rs index e683777803f02..2c8a054282b01 100644 --- a/library/std/src/sys/sync/once/futex.rs +++ b/library/std/src/sys/sync/once/futex.rs @@ -6,7 +6,7 @@ use crate::sync::once::ExclusiveState; use crate::sys::futex::{futex_wait, futex_wake_all}; // On some platforms, the OS is very nice and handles the waiter queue for us. -// This means we only need one atomic value with 5 states: +// This means we only need one atomic value with 4 states: /// No initialization has run yet, and no thread is currently using the Once. const INCOMPLETE: u32 = 0; @@ -17,16 +17,20 @@ const POISONED: u32 = 1; /// Some thread is currently attempting to run initialization. It may succeed, /// so all future threads need to wait for it to finish. const RUNNING: u32 = 2; -/// Some thread is currently attempting to run initialization and there are threads -/// waiting for it to finish. -const QUEUED: u32 = 3; /// Initialization has completed and all future calls should finish immediately. -const COMPLETE: u32 = 4; +const COMPLETE: u32 = 3; -// Threads wait by setting the state to QUEUED and calling `futex_wait` on the state +// An additional bit indicates whether there are waiting threads: + +/// May only be set if the state is not COMPLETE. +const QUEUED: u32 = 4; + +// Threads wait by setting the QUEUED bit and calling `futex_wait` on the state // variable. When the running thread finishes, it will wake all waiting threads using // `futex_wake_all`. +const STATE_MASK: u32 = 0b11; + pub struct OnceState { poisoned: bool, set_state_to: Cell, @@ -45,7 +49,7 @@ impl OnceState { } struct CompletionGuard<'a> { - state: &'a AtomicU32, + state_and_queued: &'a AtomicU32, set_state_on_drop_to: u32, } @@ -54,32 +58,32 @@ impl<'a> Drop for CompletionGuard<'a> { // Use release ordering to propagate changes to all threads checking // up on the Once. `futex_wake_all` does its own synchronization, hence // we do not need `AcqRel`. - if self.state.swap(self.set_state_on_drop_to, Release) == QUEUED { - futex_wake_all(self.state); + if self.state_and_queued.swap(self.set_state_on_drop_to, Release) & QUEUED != 0 { + futex_wake_all(self.state_and_queued); } } } pub struct Once { - state: AtomicU32, + state_and_queued: AtomicU32, } impl Once { #[inline] pub const fn new() -> Once { - Once { state: AtomicU32::new(INCOMPLETE) } + Once { state_and_queued: AtomicU32::new(INCOMPLETE) } } #[inline] pub fn is_completed(&self) -> bool { // Use acquire ordering to make all initialization changes visible to the // current thread. - self.state.load(Acquire) == COMPLETE + self.state_and_queued.load(Acquire) == COMPLETE } #[inline] pub(crate) fn state(&mut self) -> ExclusiveState { - match *self.state.get_mut() { + match *self.state_and_queued.get_mut() { INCOMPLETE => ExclusiveState::Incomplete, POISONED => ExclusiveState::Poisoned, COMPLETE => ExclusiveState::Complete, @@ -87,31 +91,73 @@ impl Once { } } - // This uses FnMut to match the API of the generic implementation. As this - // implementation is quite light-weight, it is generic over the closure and - // so avoids the cost of dynamic dispatch. #[cold] #[track_caller] - pub fn call(&self, ignore_poisoning: bool, f: &mut impl FnMut(&public::OnceState)) { - let mut state = self.state.load(Acquire); + pub fn wait(&self, ignore_poisoning: bool) { + let mut state_and_queued = self.state_and_queued.load(Acquire); loop { + let state = state_and_queued & STATE_MASK; + let queued = state_and_queued & QUEUED != 0; match state { + COMPLETE => return, + POISONED if !ignore_poisoning => { + // Panic to propagate the poison. + panic!("Once instance has previously been poisoned"); + } + _ => { + // Set the QUEUED bit if it has not already been set. + if !queued { + state_and_queued += QUEUED; + if let Err(new) = self.state_and_queued.compare_exchange_weak( + state, + state_and_queued, + Relaxed, + Acquire, + ) { + state_and_queued = new; + continue; + } + } + + futex_wait(&self.state_and_queued, state_and_queued, None); + state_and_queued = self.state_and_queued.load(Acquire); + } + } + } + } + + #[cold] + #[track_caller] + pub fn call(&self, ignore_poisoning: bool, f: &mut dyn FnMut(&public::OnceState)) { + let mut state_and_queued = self.state_and_queued.load(Acquire); + loop { + let state = state_and_queued & STATE_MASK; + let queued = state_and_queued & QUEUED != 0; + match state { + COMPLETE => return, POISONED if !ignore_poisoning => { // Panic to propagate the poison. panic!("Once instance has previously been poisoned"); } INCOMPLETE | POISONED => { // Try to register the current thread as the one running. - if let Err(new) = - self.state.compare_exchange_weak(state, RUNNING, Acquire, Acquire) - { - state = new; + let next = RUNNING + if queued { QUEUED } else { 0 }; + if let Err(new) = self.state_and_queued.compare_exchange_weak( + state_and_queued, + next, + Acquire, + Acquire, + ) { + state_and_queued = new; continue; } + // `waiter_queue` will manage other waiting threads, and // wake them up on drop. - let mut waiter_queue = - CompletionGuard { state: &self.state, set_state_on_drop_to: POISONED }; + let mut waiter_queue = CompletionGuard { + state_and_queued: &self.state_and_queued, + set_state_on_drop_to: POISONED, + }; // Run the function, letting it know if we're poisoned or not. let f_state = public::OnceState { inner: OnceState { @@ -123,21 +169,27 @@ impl Once { waiter_queue.set_state_on_drop_to = f_state.inner.set_state_to.get(); return; } - RUNNING | QUEUED => { - // Set the state to QUEUED if it is not already. - if state == RUNNING - && let Err(new) = - self.state.compare_exchange_weak(RUNNING, QUEUED, Relaxed, Acquire) - { - state = new; - continue; + _ => { + // All other values must be RUNNING. + assert!(state == RUNNING); + + // Set the QUEUED bit if it is not already set. + if !queued { + state_and_queued += QUEUED; + if let Err(new) = self.state_and_queued.compare_exchange_weak( + state, + state_and_queued, + Relaxed, + Acquire, + ) { + state_and_queued = new; + continue; + } } - futex_wait(&self.state, QUEUED, None); - state = self.state.load(Acquire); + futex_wait(&self.state_and_queued, state_and_queued, None); + state_and_queued = self.state_and_queued.load(Acquire); } - COMPLETE => return, - _ => unreachable!("state is never set to invalid values"), } } } diff --git a/library/std/src/sys/sync/once/no_threads.rs b/library/std/src/sys/sync/once/no_threads.rs index 11fde1888ba7c..12c1d9f5a6c98 100644 --- a/library/std/src/sys/sync/once/no_threads.rs +++ b/library/std/src/sys/sync/once/no_threads.rs @@ -55,6 +55,12 @@ impl Once { } } + #[cold] + #[track_caller] + pub fn wait(&self, _ignore_poisoning: bool) { + panic!("not implementable on this target"); + } + #[cold] #[track_caller] pub fn call(&self, ignore_poisoning: bool, f: &mut impl FnMut(&public::OnceState)) { diff --git a/library/std/src/sys/sync/once/queue.rs b/library/std/src/sys/sync/once/queue.rs index b04d252f8b91f..7a020c94080bf 100644 --- a/library/std/src/sys/sync/once/queue.rs +++ b/library/std/src/sys/sync/once/queue.rs @@ -56,20 +56,21 @@ // allowed, so no need for `SeqCst`. use crate::cell::Cell; -use crate::sync::atomic::{AtomicBool, AtomicPtr, Ordering}; +use crate::sync::atomic::Ordering::{AcqRel, Acquire, Release}; +use crate::sync::atomic::{AtomicBool, AtomicPtr}; use crate::sync::once::ExclusiveState; use crate::thread::{self, Thread}; use crate::{fmt, ptr, sync as public}; -type Masked = (); +type StateAndQueue = *mut (); pub struct Once { - state_and_queue: AtomicPtr, + state_and_queue: AtomicPtr<()>, } pub struct OnceState { poisoned: bool, - set_state_on_drop_to: Cell<*mut Masked>, + set_state_on_drop_to: Cell, } // Four states that a Once can be in, encoded into the lower bits of @@ -81,7 +82,8 @@ const COMPLETE: usize = 0x3; // Mask to learn about the state. All other bits are the queue of waiters if // this is in the RUNNING state. -const STATE_MASK: usize = 0x3; +const STATE_MASK: usize = 0b11; +const QUEUE_MASK: usize = !STATE_MASK; // Representation of a node in the linked list of waiters, used while in the // RUNNING state. @@ -93,15 +95,23 @@ const STATE_MASK: usize = 0x3; struct Waiter { thread: Cell>, signaled: AtomicBool, - next: *const Waiter, + next: Cell<*const Waiter>, } // Head of a linked list of waiters. // Every node is a struct on the stack of a waiting thread. // Will wake up the waiters when it gets dropped, i.e. also on panic. struct WaiterQueue<'a> { - state_and_queue: &'a AtomicPtr, - set_state_on_drop_to: *mut Masked, + state_and_queue: &'a AtomicPtr<()>, + set_state_on_drop_to: StateAndQueue, +} + +fn to_queue(current: StateAndQueue) -> *const Waiter { + current.mask(QUEUE_MASK).cast() +} + +fn to_state(current: StateAndQueue) -> usize { + current.addr() & STATE_MASK } impl Once { @@ -117,7 +127,7 @@ impl Once { // operations visible to us, and, this being a fast path, weaker // ordering helps with performance. This `Acquire` synchronizes with // `Release` operations on the slow path. - self.state_and_queue.load(Ordering::Acquire).addr() == COMPLETE + self.state_and_queue.load(Acquire).addr() == COMPLETE } #[inline] @@ -130,6 +140,25 @@ impl Once { } } + #[cold] + #[track_caller] + pub fn wait(&self, ignore_poisoning: bool) { + let mut current = self.state_and_queue.load(Acquire); + loop { + let state = to_state(current); + match state { + COMPLETE => return, + POISONED if !ignore_poisoning => { + // Panic to propagate the poison. + panic!("Once instance has previously been poisoned"); + } + _ => { + current = wait(&self.state_and_queue, current); + } + } + } + } + // This is a non-generic function to reduce the monomorphization cost of // using `call_once` (this isn't exactly a trivial or small implementation). // @@ -144,9 +173,10 @@ impl Once { #[cold] #[track_caller] pub fn call(&self, ignore_poisoning: bool, init: &mut dyn FnMut(&public::OnceState)) { - let mut state_and_queue = self.state_and_queue.load(Ordering::Acquire); + let mut current = self.state_and_queue.load(Acquire); loop { - match state_and_queue.addr() { + let state = to_state(current); + match state { COMPLETE => break, POISONED if !ignore_poisoning => { // Panic to propagate the poison. @@ -154,16 +184,16 @@ impl Once { } POISONED | INCOMPLETE => { // Try to register this thread as the one RUNNING. - let exchange_result = self.state_and_queue.compare_exchange( - state_and_queue, - ptr::without_provenance_mut(RUNNING), - Ordering::Acquire, - Ordering::Acquire, - ); - if let Err(old) = exchange_result { - state_and_queue = old; + if let Err(new) = self.state_and_queue.compare_exchange_weak( + current, + current.mask(QUEUE_MASK).wrapping_byte_add(RUNNING), + Acquire, + Acquire, + ) { + current = new; continue; } + // `waiter_queue` will manage other waiting threads, and // wake them up on drop. let mut waiter_queue = WaiterQueue { @@ -174,54 +204,53 @@ impl Once { // poisoned or not. let init_state = public::OnceState { inner: OnceState { - poisoned: state_and_queue.addr() == POISONED, + poisoned: state == POISONED, set_state_on_drop_to: Cell::new(ptr::without_provenance_mut(COMPLETE)), }, }; init(&init_state); waiter_queue.set_state_on_drop_to = init_state.inner.set_state_on_drop_to.get(); - break; + return; } _ => { // All other values must be RUNNING with possibly a // pointer to the waiter queue in the more significant bits. - assert!(state_and_queue.addr() & STATE_MASK == RUNNING); - wait(&self.state_and_queue, state_and_queue); - state_and_queue = self.state_and_queue.load(Ordering::Acquire); + assert!(state == RUNNING); + current = wait(&self.state_and_queue, current); } } } } } -fn wait(state_and_queue: &AtomicPtr, mut current_state: *mut Masked) { - // Note: the following code was carefully written to avoid creating a - // mutable reference to `node` that gets aliased. +fn wait(state_and_queue: &AtomicPtr<()>, mut current: StateAndQueue) -> StateAndQueue { + let node = &Waiter { + thread: Cell::new(Some(thread::current())), + signaled: AtomicBool::new(false), + next: Cell::new(ptr::null()), + }; + loop { - // Don't queue this thread if the status is no longer running, - // otherwise we will not be woken up. - if current_state.addr() & STATE_MASK != RUNNING { - return; + let state = to_state(current); + let queue = to_queue(current); + + // If initialization has finished, return. + if matches!(state, POISONED | COMPLETE) { + return current; } - // Create the node for our current thread. - let node = Waiter { - thread: Cell::new(Some(thread::current())), - signaled: AtomicBool::new(false), - next: current_state.with_addr(current_state.addr() & !STATE_MASK) as *const Waiter, - }; - let me = core::ptr::addr_of!(node) as *const Masked as *mut Masked; + // Update the node for our current thread. + node.next.set(queue); // Try to slide in the node at the head of the linked list, making sure // that another thread didn't just replace the head of the linked list. - let exchange_result = state_and_queue.compare_exchange( - current_state, - me.with_addr(me.addr() | RUNNING), - Ordering::Release, - Ordering::Relaxed, - ); - if let Err(old) = exchange_result { - current_state = old; + if let Err(new) = state_and_queue.compare_exchange_weak( + current, + ptr::from_ref(node).wrapping_byte_add(state) as StateAndQueue, + Release, + Acquire, + ) { + current = new; continue; } @@ -230,14 +259,15 @@ fn wait(state_and_queue: &AtomicPtr, mut current_state: *mut Masked) { // would drop our `Waiter` node and leave a hole in the linked list // (and a dangling reference). Guard against spurious wakeups by // reparking ourselves until we are signaled. - while !node.signaled.load(Ordering::Acquire) { + while !node.signaled.load(Acquire) { // If the managing thread happens to signal and unpark us before we // can park ourselves, the result could be this thread never gets // unparked. Luckily `park` comes with the guarantee that if it got // an `unpark` just before on an unparked thread it does not park. thread::park(); } - break; + + return state_and_queue.load(Acquire); } } @@ -251,11 +281,10 @@ impl fmt::Debug for Once { impl Drop for WaiterQueue<'_> { fn drop(&mut self) { // Swap out our state with however we finished. - let state_and_queue = - self.state_and_queue.swap(self.set_state_on_drop_to, Ordering::AcqRel); + let current = self.state_and_queue.swap(self.set_state_on_drop_to, AcqRel); // We should only ever see an old state which was RUNNING. - assert_eq!(state_and_queue.addr() & STATE_MASK, RUNNING); + assert_eq!(current.addr() & STATE_MASK, RUNNING); // Walk the entire linked list of waiters and wake them up (in lifo // order, last to register is first to wake up). @@ -264,16 +293,13 @@ impl Drop for WaiterQueue<'_> { // free `node` if there happens to be has a spurious wakeup. // So we have to take out the `thread` field and copy the pointer to // `next` first. - let mut queue = - state_and_queue.with_addr(state_and_queue.addr() & !STATE_MASK) as *const Waiter; + let mut queue = to_queue(current); while !queue.is_null() { - let next = (*queue).next; + let next = (*queue).next.get(); let thread = (*queue).thread.take().unwrap(); - (*queue).signaled.store(true, Ordering::Release); - // ^- FIXME (maybe): This is another case of issue #55005 - // `store()` has a potentially dangling ref to `signaled`. - queue = next; + (*queue).signaled.store(true, Release); thread.unpark(); + queue = next; } } } From 1d49aad8445881ac060cfae0e351207c05f82dfd Mon Sep 17 00:00:00 2001 From: joboet Date: Fri, 12 Jul 2024 13:17:43 +0200 Subject: [PATCH 2/2] std: fix busy-waiting in `Once::wait_force`, add more tests --- library/std/src/sync/once/tests.rs | 47 ++++++++++++++++++++++++++ library/std/src/sys/sync/once/queue.rs | 12 ++++--- 2 files changed, 55 insertions(+), 4 deletions(-) diff --git a/library/std/src/sync/once/tests.rs b/library/std/src/sync/once/tests.rs index d43dabc1cf137..ce96468aeb6e1 100644 --- a/library/std/src/sync/once/tests.rs +++ b/library/std/src/sync/once/tests.rs @@ -1,5 +1,8 @@ use super::Once; +use crate::sync::atomic::AtomicBool; +use crate::sync::atomic::Ordering::Relaxed; use crate::sync::mpsc::channel; +use crate::time::Duration; use crate::{panic, thread}; #[test] @@ -113,3 +116,47 @@ fn wait_for_force_to_finish() { assert!(t1.join().is_ok()); assert!(t2.join().is_ok()); } + +#[test] +fn wait() { + for _ in 0..50 { + let val = AtomicBool::new(false); + let once = Once::new(); + + thread::scope(|s| { + for _ in 0..4 { + s.spawn(|| { + once.wait(); + assert!(val.load(Relaxed)); + }); + } + + once.call_once(|| val.store(true, Relaxed)); + }); + } +} + +#[test] +fn wait_on_poisoned() { + let once = Once::new(); + + panic::catch_unwind(|| once.call_once(|| panic!())).unwrap_err(); + panic::catch_unwind(|| once.wait()).unwrap_err(); +} + +#[test] +fn wait_force_on_poisoned() { + let once = Once::new(); + + thread::scope(|s| { + panic::catch_unwind(|| once.call_once(|| panic!())).unwrap_err(); + + s.spawn(|| { + thread::sleep(Duration::from_millis(100)); + + once.call_once_force(|_| {}); + }); + + once.wait_force(); + }) +} diff --git a/library/std/src/sys/sync/once/queue.rs b/library/std/src/sys/sync/once/queue.rs index 7a020c94080bf..86f72c82008bc 100644 --- a/library/std/src/sys/sync/once/queue.rs +++ b/library/std/src/sys/sync/once/queue.rs @@ -153,7 +153,7 @@ impl Once { panic!("Once instance has previously been poisoned"); } _ => { - current = wait(&self.state_and_queue, current); + current = wait(&self.state_and_queue, current, !ignore_poisoning); } } } @@ -216,14 +216,18 @@ impl Once { // All other values must be RUNNING with possibly a // pointer to the waiter queue in the more significant bits. assert!(state == RUNNING); - current = wait(&self.state_and_queue, current); + current = wait(&self.state_and_queue, current, true); } } } } } -fn wait(state_and_queue: &AtomicPtr<()>, mut current: StateAndQueue) -> StateAndQueue { +fn wait( + state_and_queue: &AtomicPtr<()>, + mut current: StateAndQueue, + return_on_poisoned: bool, +) -> StateAndQueue { let node = &Waiter { thread: Cell::new(Some(thread::current())), signaled: AtomicBool::new(false), @@ -235,7 +239,7 @@ fn wait(state_and_queue: &AtomicPtr<()>, mut current: StateAndQueue) -> StateAnd let queue = to_queue(current); // If initialization has finished, return. - if matches!(state, POISONED | COMPLETE) { + if state == COMPLETE || (return_on_poisoned && state == POISONED) { return current; }