std: implement the once_wait feature

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

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

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 {
+        if self.state_and_queued.swap(self.set_state_on_drop_to, Release) & QUEUED != 0 {
-            futex_wake_all(self.state);
+            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)) {
+    pub fn wait(&self, ignore_poisoning: bool) {
-        let mut state = self.state.load(Acquire);
+        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) =
+                    let next = RUNNING + if queued { QUEUED } else { 0 };
-                        self.state.compare_exchange_weak(state, RUNNING, Acquire, Acquire)
+                    if let Err(new) = self.state_and_queued.compare_exchange_weak(
-                    {
+                        state_and_queued,
-                        state = new;
+                        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 =
+                    let mut waiter_queue = CompletionGuard {
-                        CompletionGuard { state: &self.state, set_state_on_drop_to: POISONED };
+                        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.
+                    // All other values must be RUNNING.
-                    if state == RUNNING
+                    assert!(state == RUNNING);
-                        && let Err(new) =
+
-                            self.state.compare_exchange_weak(RUNNING, QUEUED, Relaxed, Acquire)
+                    // Set the QUEUED bit if it is not already set.
-                    {
+                    if !queued {
-                        state = new;
+                        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);
                     }
-                COMPLETE => return,
+
-                _ => unreachable!("state is never set to invalid values"),
+                    futex_wait(&self.state_and_queued, state_and_queued, None);
+                    state_and_queued = self.state_and_queued.load(Acquire);
+                }
             }
         }
     }

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)) {

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>,
+    state_and_queue: &'a AtomicPtr<()>,
-    set_state_on_drop_to: *mut Masked,
+    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(
+                    if let Err(new) = self.state_and_queue.compare_exchange_weak(
-                        state_and_queue,
+                        current,
-                        ptr::without_provenance_mut(RUNNING),
+                        current.mask(QUEUE_MASK).wrapping_byte_add(RUNNING),
-                        Ordering::Acquire,
+                        Acquire,
-                        Ordering::Acquire,
+                        Acquire,
-                    );
+                    ) {
-                    if let Err(old) = exchange_result {
+                        current = new;
-                        state_and_queue = old;
                         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);
+                    assert!(state == RUNNING);
-                    wait(&self.state_and_queue, state_and_queue);
+                    current = wait(&self.state_and_queue, current);
-                    state_and_queue = self.state_and_queue.load(Ordering::Acquire);
                 }
             }
         }
     }
 }
 
-fn wait(state_and_queue: &AtomicPtr<Masked>, mut current_state: *mut Masked) {
+fn wait(state_and_queue: &AtomicPtr<()>, mut current: StateAndQueue) -> StateAndQueue {
-    // Note: the following code was carefully written to avoid creating a
+    let node = &Waiter {
-    // mutable reference to `node` that gets aliased.
-    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;
-        }
-
-        // 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,
+        next: Cell::new(ptr::null()),
     };
-        let me = core::ptr::addr_of!(node) as *const Masked as *mut Masked;
+
+    loop {
+        let state = to_state(current);
+        let queue = to_queue(current);
+
+        // If initialization has finished, return.
+        if matches!(state, POISONED | COMPLETE) {
+            return current;
+        }
+
+        // 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(
+        if let Err(new) = state_and_queue.compare_exchange_weak(
-            current_state,
+            current,
-            me.with_addr(me.addr() | RUNNING),
+            ptr::from_ref(node).wrapping_byte_add(state) as StateAndQueue,
-            Ordering::Release,
+            Release,
-            Ordering::Relaxed,
+            Acquire,
-        );
+        ) {
-        if let Err(old) = exchange_result {
+            current = new;
-            current_state = old;
             continue;
         }
 
@@ -230,14 +259,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 +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 =
+        let current = self.state_and_queue.swap(self.set_state_on_drop_to, AcqRel);
-            self.state_and_queue.swap(self.set_state_on_drop_to, Ordering::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 =
+            let mut queue = to_queue(current);
-                state_and_queue.with_addr(state_and_queue.addr() & !STATE_MASK) as *const Waiter;
             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);
+                (*queue).signaled.store(true, Release);
-                // ^- FIXME (maybe): This is another case of issue #55005
-                // `store()` has a potentially dangling ref to `signaled`.
-                queue = next;
                 thread.unpark();
+                queue = next;
             }
         }
     }