5 files changed, 94 insertions, 40 deletions

diff --git a/library/std/src/sync/lazy_lock.rs b/library/std/src/sync/lazy_lock.rs
index 4a1530530..7e85d6a06 100644
--- a/library/std/src/sync/lazy_lock.rs
+++ b/library/std/src/sync/lazy_lock.rs
@@ -1,8 +1,21 @@
-use crate::cell::Cell;
+use crate::cell::UnsafeCell;
 use crate::fmt;
+use crate::mem::ManuallyDrop;
 use crate::ops::Deref;
 use crate::panic::{RefUnwindSafe, UnwindSafe};
-use crate::sync::OnceLock;
+use crate::sync::Once;
+
+use super::once::ExclusiveState;
+
+// We use the state of a Once as discriminant value. Upon creation, the state is
+// "incomplete" and `f` contains the initialization closure. In the first call to
+// `call_once`, `f` is taken and run. If it succeeds, `value` is set and the state
+// is changed to "complete". If it panics, the Once is poisoned, so none of the
+// two fields is initialized.
+union Data<T, F> {
+    value: ManuallyDrop<T>,
+    f: ManuallyDrop<F>,
+}
 
 /// A value which is initialized on the first access.
 ///
@@ -43,16 +56,17 @@ use crate::sync::OnceLock;
 /// ```
 #[unstable(feature = "once_cell", issue = "74465")]
 pub struct LazyLock<T, F = fn() -> T> {
-    cell: OnceLock<T>,
-    init: Cell<Option<F>>,
+    once: Once,
+    data: UnsafeCell<Data<T, F>>,
 }
+
 impl<T, F: FnOnce() -> T> LazyLock<T, F> {
     /// Creates a new lazy value with the given initializing
     /// function.
     #[inline]
     #[unstable(feature = "once_cell", issue = "74465")]
     pub const fn new(f: F) -> LazyLock<T, F> {
-        LazyLock { cell: OnceLock::new(), init: Cell::new(Some(f)) }
+        LazyLock { once: Once::new(), data: UnsafeCell::new(Data { f: ManuallyDrop::new(f) }) }
     }
 
     /// Forces the evaluation of this lazy value and
@@ -74,10 +88,50 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
     #[inline]
     #[unstable(feature = "once_cell", issue = "74465")]
     pub fn force(this: &LazyLock<T, F>) -> &T {
-        this.cell.get_or_init(|| match this.init.take() {
-            Some(f) => f(),
-            None => panic!("Lazy instance has previously been poisoned"),
-        })
+        this.once.call_once(|| {
+            // SAFETY: `call_once` only runs this closure once, ever.
+            let data = unsafe { &mut *this.data.get() };
+            let f = unsafe { ManuallyDrop::take(&mut data.f) };
+            let value = f();
+            data.value = ManuallyDrop::new(value);
+        });
+
+        // SAFETY:
+        // There are four possible scenarios:
+        // * the closure was called and initialized `value`.
+        // * the closure was called and panicked, so this point is never reached.
+        // * the closure was not called, but a previous call initialized `value`.
+        // * the closure was not called because the Once is poisoned, so this point
+        //   is never reached.
+        // So `value` has definitely been initialized and will not be modified again.
+        unsafe { &*(*this.data.get()).value }
+    }
+}
+
+impl<T, F> LazyLock<T, F> {
+    /// Get the inner value if it has already been initialized.
+    fn get(&self) -> Option<&T> {
+        if self.once.is_completed() {
+            // SAFETY:
+            // The closure has been run successfully, so `value` has been initialized
+            // and will not be modified again.
+            Some(unsafe { &*(*self.data.get()).value })
+        } else {
+            None
+        }
+    }
+}
+
+#[unstable(feature = "once_cell", issue = "74465")]
+impl<T, F> Drop for LazyLock<T, F> {
+    fn drop(&mut self) {
+        match self.once.state() {
+            ExclusiveState::Incomplete => unsafe { ManuallyDrop::drop(&mut self.data.get_mut().f) },
+            ExclusiveState::Complete => unsafe {
+                ManuallyDrop::drop(&mut self.data.get_mut().value)
+            },
+            ExclusiveState::Poisoned => {}
+        }
     }
 }
 
@@ -103,23 +157,23 @@ impl<T: Default> Default for LazyLock<T> {
 #[unstable(feature = "once_cell", issue = "74465")]
 impl<T: fmt::Debug, F> fmt::Debug for LazyLock<T, F> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_struct("Lazy").field("cell", &self.cell).finish_non_exhaustive()
+        match self.get() {
+            Some(v) => f.debug_tuple("LazyLock").field(v).finish(),
+            None => f.write_str("LazyLock(Uninit)"),
+        }
     }
 }
 
 // We never create a `&F` from a `&LazyLock<T, F>` so it is fine
 // to not impl `Sync` for `F`
-// we do create a `&mut Option<F>` in `force`, but this is
-// properly synchronized, so it only happens once
-// so it also does not contribute to this impl.
 #[unstable(feature = "once_cell", issue = "74465")]
-unsafe impl<T, F: Send> Sync for LazyLock<T, F> where OnceLock<T>: Sync {}
+unsafe impl<T: Sync + Send, F: Send> Sync for LazyLock<T, F> {}
 // auto-derived `Send` impl is OK.
 
 #[unstable(feature = "once_cell", issue = "74465")]
-impl<T, F: UnwindSafe> RefUnwindSafe for LazyLock<T, F> where OnceLock<T>: RefUnwindSafe {}
+impl<T: RefUnwindSafe + UnwindSafe, F: UnwindSafe> RefUnwindSafe for LazyLock<T, F> {}
 #[unstable(feature = "once_cell", issue = "74465")]
-impl<T, F: UnwindSafe> UnwindSafe for LazyLock<T, F> where OnceLock<T>: UnwindSafe {}
+impl<T: UnwindSafe, F: UnwindSafe> UnwindSafe for LazyLock<T, F> {}
 
 #[cfg(test)]
 mod tests;
diff --git a/library/std/src/sync/mod.rs b/library/std/src/sync/mod.rs
index ba20bab87..4edc95617 100644
--- a/library/std/src/sync/mod.rs
+++ b/library/std/src/sync/mod.rs
@@ -186,7 +186,7 @@ mod condvar;
 mod lazy_lock;
 mod mpmc;
 mod mutex;
-mod once;
+pub(crate) mod once;
 mod once_lock;
 mod poison;
 mod remutex;
diff --git a/library/std/src/sync/mpmc/array.rs b/library/std/src/sync/mpmc/array.rs
index c1e3e48b0..c6bb09b04 100644
--- a/library/std/src/sync/mpmc/array.rs
+++ b/library/std/src/sync/mpmc/array.rs
@@ -319,19 +319,10 @@ impl<T> Channel<T> {
     ) -> Result<(), SendTimeoutError<T>> {
         let token = &mut Token::default();
         loop {
-            // Try sending a message several times.
-            let backoff = Backoff::new();
-            loop {
-                if self.start_send(token) {
-                    let res = unsafe { self.write(token, msg) };
-                    return res.map_err(SendTimeoutError::Disconnected);
-                }
-
-                if backoff.is_completed() {
-                    break;
-                } else {
-                    backoff.spin_light();
-                }
+            // Try sending a message.
+            if self.start_send(token) {
+                let res = unsafe { self.write(token, msg) };
+                return res.map_err(SendTimeoutError::Disconnected);
             }
 
             if let Some(d) = deadline {
@@ -379,6 +370,7 @@ impl<T> Channel<T> {
     pub(crate) fn recv(&self, deadline: Option<Instant>) -> Result<T, RecvTimeoutError> {
         let token = &mut Token::default();
         loop {
+            // Try receiving a message.
             if self.start_recv(token) {
                 let res = unsafe { self.read(token) };
                 return res.map_err(|_| RecvTimeoutError::Disconnected);
diff --git a/library/std/src/sync/mpmc/utils.rs b/library/std/src/sync/mpmc/utils.rs
index cfe42750d..d053d69e2 100644
--- a/library/std/src/sync/mpmc/utils.rs
+++ b/library/std/src/sync/mpmc/utils.rs
@@ -105,10 +105,8 @@ impl Backoff {
 
     /// Backs off using lightweight spinning.
     ///
-    /// This method should be used for:
-    ///     - Retrying an operation because another thread made progress. i.e. on CAS failure.
-    ///     - Waiting for an operation to complete by spinning optimistically for a few iterations
-    ///     before falling back to parking the thread (see `Backoff::is_completed`).
+    /// This method should be used for retrying an operation because another thread made
+    /// progress. i.e. on CAS failure.
     #[inline]
     pub fn spin_light(&self) {
         let step = self.step.get().min(SPIN_LIMIT);
@@ -134,10 +132,4 @@ impl Backoff {
 
         self.step.set(self.step.get() + 1);
     }
-
-    /// Returns `true` if quadratic backoff has completed and parking the thread is advised.
-    #[inline]
-    pub fn is_completed(&self) -> bool {
-        self.step.get() > SPIN_LIMIT
-    }
 }
diff --git a/library/std/src/sync/once.rs b/library/std/src/sync/once.rs
index 0f25417d6..1b17c3108 100644
--- a/library/std/src/sync/once.rs
+++ b/library/std/src/sync/once.rs
@@ -43,6 +43,12 @@ pub struct OnceState {
     pub(crate) inner: sys::OnceState,
 }
 
+pub(crate) enum ExclusiveState {
+    Incomplete,
+    Poisoned,
+    Complete,
+}
+
 /// Initialization value for static [`Once`] values.
 ///
 /// # Examples
@@ -248,6 +254,16 @@ impl Once {
     pub fn is_completed(&self) -> bool {
         self.inner.is_completed()
     }
+
+    /// Returns the current state of the `Once` instance.
+    ///
+    /// Since this takes a mutable reference, no initialization can currently
+    /// be running, so the state must be either "incomplete", "poisoned" or
+    /// "complete".
+    #[inline]
+    pub(crate) fn state(&mut self) -> ExclusiveState {
+        self.inner.state()
+    }
 }
 
 #[stable(feature = "std_debug", since = "1.16.0")]