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/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/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<T> OnceLock<T> {
}
}
+ /// 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<u32>,
@@ -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..86f72c82008bc 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<Masked>,
+ state_and_queue: AtomicPtr<()>,
}
pub struct OnceState {
poisoned: bool,
- set_state_on_drop_to: Cell<*mut Masked>,
+ set_state_on_drop_to: Cell<StateAndQueue>,
}
// 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<Option<Thread>>,
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<Masked>,
- 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, !ignore_poisoning);
+ }
+ }
+ }
+ }
+
// 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,57 @@ 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, true);
}
}
}
}
}
-fn wait(state_and_queue: &AtomicPtr<Masked>, 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,
+ return_on_poisoned: bool,
+) -> 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 state == COMPLETE || (return_on_poisoned && state == POISONED) {
+ 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 +263,15 @@ fn wait(state_and_queue: &AtomicPtr<Masked>, 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 +285,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 +297,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;
}
}
}