Revert "remove inlined lazy::Waiter in favor of sync::Once"

This reverts commit d101794.

Author: Mark-Simulacrum

src/libstd/lazy.rs

@@ -3,10 +3,12 @@
 use crate::{
     cell::{Cell, UnsafeCell},
     fmt,
+    marker::PhantomData,
     mem::{self, MaybeUninit},
     ops::{Deref, Drop},
     panic::{RefUnwindSafe, UnwindSafe},
-    sync::Once,
+    sync::atomic::{AtomicBool, AtomicUsize, Ordering},
+    thread::{self, Thread},
 };
 
 #[doc(inline)]
@@ -40,7 +42,10 @@ pub use core::lazy::*;
 /// ```
 #[unstable(feature = "once_cell", issue = "68198")]
 pub struct SyncOnceCell<T> {
-    once: Once,
+    // This `state` word is actually an encoded version of just a pointer to a
+    // `Waiter`, so we add the `PhantomData` appropriately.
+    state_and_queue: AtomicUsize,
+    _marker: PhantomData<*mut Waiter>,
     // Whether or not the value is initialized is tracked by `state_and_queue`.
     value: UnsafeCell<MaybeUninit<T>>,
 }
@@ -117,7 +122,8 @@ impl<T> SyncOnceCell<T> {
     #[unstable(feature = "once_cell", issue = "68198")]
     pub const fn new() -> SyncOnceCell<T> {
         SyncOnceCell {
-            once: Once::new(),
+            state_and_queue: AtomicUsize::new(INCOMPLETE),
+            _marker: PhantomData,
             value: UnsafeCell::new(MaybeUninit::uninit()),
         }
     }
@@ -129,7 +135,7 @@ impl<T> SyncOnceCell<T> {
     #[unstable(feature = "once_cell", issue = "68198")]
     pub fn get(&self) -> Option<&T> {
         if self.is_initialized() {
-            // Safe b/c checked is_initialized
+            // Safe b/c checked is_initialize
             Some(unsafe { self.get_unchecked() })
         } else {
             None
@@ -142,7 +148,7 @@ impl<T> SyncOnceCell<T> {
     #[unstable(feature = "once_cell", issue = "68198")]
     pub fn get_mut(&mut self) -> Option<&mut T> {
         if self.is_initialized() {
-            // Safe b/c checked is_initialized and we have a unique access
+            // Safe b/c checked is_initialize and we have a unique access
             Some(unsafe { self.get_unchecked_mut() })
         } else {
             None
@@ -344,32 +350,37 @@ impl<T> SyncOnceCell<T> {
         }
     }
 
+    /// Safety: synchronizes with store to value via Release/(Acquire|SeqCst).
     #[inline]
     fn is_initialized(&self) -> bool {
-        self.once.is_completed()
+        // An `Acquire` load is enough because that makes all the initialization
+        // operations visible to us, and, this being a fast path, weaker
+        // ordering helps with performance. This `Acquire` synchronizes with
+        // `SeqCst` operations on the slow path.
+        self.state_and_queue.load(Ordering::Acquire) == COMPLETE
     }
 
+    /// Safety: synchronizes with store to value via SeqCst read from state,
+    /// writes value only once because we never get to INCOMPLETE state after a
+    /// successful write.
     #[cold]
     fn initialize<F, E>(&self, f: F) -> Result<(), E>
     where
         F: FnOnce() -> Result<T, E>,
     {
+        let mut f = Some(f);
         let mut res: Result<(), E> = Ok(());
         let slot = &self.value;
-
-        // Ignore poisoning from other threads
-        // If another thread panics, then we'll be able to run our closure
-        self.once.call_once_force(|p| {
+        initialize_inner(&self.state_and_queue, &mut || {
+            let f = f.take().unwrap();
             match f() {
                 Ok(value) => {
                     unsafe { (&mut *slot.get()).write(value) };
+                    true
                 }
                 Err(e) => {
                     res = Err(e);
-
-                    // Treat the underlying `Once` as poisoned since we
-                    // failed to initialize our value. Calls
-                    p.poison();
+                    false
                 }
             }
         });
@@ -396,6 +407,106 @@ impl<T> Drop for SyncOnceCell<T> {
     }
 }
 
+const INCOMPLETE: usize = 0x0;
+const RUNNING: usize = 0x1;
+const COMPLETE: usize = 0x2;
+
+const STATE_MASK: usize = 0x3;
+
+// The alignment here is so that we can stash the state in the lower
+// bits of the `next` pointer
+#[repr(align(4))]
+struct Waiter {
+    thread: Cell<Option<Thread>>,
+    signaled: AtomicBool,
+    next: *const Waiter,
+}
+
+struct WaiterQueue<'a> {
+    state_and_queue: &'a AtomicUsize,
+    set_state_on_drop_to: usize,
+}
+
+impl Drop for WaiterQueue<'_> {
+    fn drop(&mut self) {
+        let state_and_queue =
+            self.state_and_queue.swap(self.set_state_on_drop_to, Ordering::AcqRel);
+
+        assert_eq!(state_and_queue & STATE_MASK, RUNNING);
+
+        unsafe {
+            let mut queue = (state_and_queue & !STATE_MASK) as *const Waiter;
+            while !queue.is_null() {
+                let next = (*queue).next;
+                let thread = (*queue).thread.replace(None).unwrap();
+                (*queue).signaled.store(true, Ordering::Release);
+                queue = next;
+                thread.unpark();
+            }
+        }
+    }
+}
+
+fn initialize_inner(my_state_and_queue: &AtomicUsize, init: &mut dyn FnMut() -> bool) -> bool {
+    let mut state_and_queue = my_state_and_queue.load(Ordering::Acquire);
+
+    loop {
+        match state_and_queue {
+            COMPLETE => return true,
+            INCOMPLETE => {
+                let old = my_state_and_queue.compare_and_swap(
+                    state_and_queue,
+                    RUNNING,
+                    Ordering::Acquire,
+                );
+                if old != state_and_queue {
+                    state_and_queue = old;
+                    continue;
+                }
+                let mut waiter_queue = WaiterQueue {
+                    state_and_queue: my_state_and_queue,
+                    set_state_on_drop_to: INCOMPLETE,
+                };
+                let success = init();
+
+                waiter_queue.set_state_on_drop_to = if success { COMPLETE } else { INCOMPLETE };
+                return success;
+            }
+            _ => {
+                assert!(state_and_queue & STATE_MASK == RUNNING);
+                wait(&my_state_and_queue, state_and_queue);
+                state_and_queue = my_state_and_queue.load(Ordering::Acquire);
+            }
+        }
+    }
+}
+
+fn wait(state_and_queue: &AtomicUsize, mut current_state: usize) {
+    loop {
+        if current_state & STATE_MASK != RUNNING {
+            return;
+        }
+
+        let node = Waiter {
+            thread: Cell::new(Some(thread::current())),
+            signaled: AtomicBool::new(false),
+            next: (current_state & !STATE_MASK) as *const Waiter,
+        };
+        let me = &node as *const Waiter as usize;
+
+        let old = state_and_queue.compare_and_swap(current_state, me | RUNNING, Ordering::Release);
+        if old != current_state {
+            current_state = old;
+            continue;
+        }
+
+        while !node.signaled.load(Ordering::Acquire) {
+            thread::park();
+        }
+        break;
+    }
+}
+
 /// A value which is initialized on the first access.
 ///
 /// This type is a thread-safe `Lazy`, and can be used in statics.
@@ -652,7 +763,6 @@ mod tests {
 
         let res = panic::catch_unwind(|| cell.get_or_try_init(|| -> Result<_, ()> { panic!() }));
         assert!(res.is_err());
-        assert!(!cell.is_initialized());
         assert!(cell.get().is_none());
 
         assert_eq!(cell.get_or_try_init(|| Err(())), Err(()));

src/libstd/sync/once.rs

@@ -132,7 +132,6 @@ unsafe impl Send for Once {}
 #[derive(Debug)]
 pub struct OnceState {
     poisoned: bool,
-    set_state_on_drop_to: Cell<usize>,
 }
 
 /// Initialization value for static [`Once`] values.
@@ -322,7 +321,7 @@ impl Once {
         }
 
         let mut f = Some(f);
-        self.call_inner(true, &mut |p| f.take().unwrap()(p));
+        self.call_inner(true, &mut |p| f.take().unwrap()(&OnceState { poisoned: p }));
     }
 
     /// Returns `true` if some `call_once` call has completed
@@ -386,7 +385,7 @@ impl Once {
     // currently no way to take an `FnOnce` and call it via virtual dispatch
     // without some allocation overhead.
     #[cold]
-    fn call_inner(&self, ignore_poisoning: bool, init: &mut dyn FnMut(&OnceState)) {
+    fn call_inner(&self, ignore_poisoning: bool, init: &mut dyn FnMut(bool)) {
         let mut state_and_queue = self.state_and_queue.load(Ordering::Acquire);
         loop {
             match state_and_queue {
@@ -414,9 +413,8 @@ impl Once {
                     };
                     // Run the initialization function, letting it know if we're
                     // poisoned or not.
-                    let init_state = OnceState { poisoned: state_and_queue == POISONED, set_state_on_drop_to: Cell::new(COMPLETE) };
-                    init(&init_state);
-                    waiter_queue.set_state_on_drop_to = init_state.set_state_on_drop_to.get();
+                    init(state_and_queue == POISONED);
+                    waiter_queue.set_state_on_drop_to = COMPLETE;
                     break;
                 }
                 _ => {
@@ -556,14 +554,6 @@ impl OnceState {
     pub fn poisoned(&self) -> bool {
         self.poisoned
     }
-
-    /// Poison the associated [`Once`] without explicitly panicking.
-    ///
-    /// [`Once`]: struct.Once.html
-    // NOTE: This is currently only exposed for the `lazy` module
-    pub(crate) fn poison(&self) {
-        self.set_state_on_drop_to.set(POISONED);
-    }
 }
 
 #[cfg(all(test, not(target_os = "emscripten")))]