Use loom Mutex, remove check in get_wait

Author: Daksh14

tokio/src/sync/set_once.rs

@@ -1,11 +1,13 @@
 use super::{Notify, SetError};
-use crate::{loom::cell::UnsafeCell, pin};
+use crate::{
+    loom::{cell::UnsafeCell, sync::Mutex},
+    pin,
+};
 use std::fmt;
 use std::mem::MaybeUninit;
 use std::ops::Drop;
 use std::ptr;
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::Mutex;
 
 // This file contains an implementation of an SetOnce. The value of SetOnce
 // can only be modified once during initialization.
@@ -181,7 +183,7 @@ impl<T> SetOnce<T> {
             value_set: AtomicBool::new(false),
             value: UnsafeCell::new(MaybeUninit::uninit()),
             notify: Notify::const_new(),
-            lock: Mutex::new(()),
+            lock: Mutex::const_new(()),
         }
     }
 
@@ -232,7 +234,7 @@ impl<T> SetOnce<T> {
             value_set: AtomicBool::new(true),
             value: UnsafeCell::new(MaybeUninit::new(value)),
             notify: Notify::const_new(),
-            lock: Mutex::new(()),
+            lock: Mutex::const_new(()),
         }
     }
 
@@ -271,6 +273,10 @@ impl<T> SetOnce<T> {
             self.value.with_mut(|ptr| (*ptr).as_mut_ptr().write(value));
         }
 
+        // Using release ordering so any threads that read a true from this
+        // atomic is able to read the value we just stored.
+        self.value_set.store(true, Ordering::Release);
+
         // notify the waiting wakers that the value is set
         self.notify.notify_waiters();
     }
@@ -293,29 +299,13 @@ impl<T> SetOnce<T> {
 
         // SAFETY: lock the mutex to ensure only one caller of set
         // can run at a time.
-        match self.lock.lock() {
-            Ok(_) => {
-                // Using release ordering so any threads that read a true from this
-                // atomic is able to read the value we just stored.
-                if !self.value_set.swap(true, Ordering::Release) {
-                    // SAFETY: We are swapping the value_set AtomicBool from FALSE to
-                    // TRUE with it being previously false and followed by that we are
-                    // initializing the unsafe Cell field with the value
-                    unsafe {
-                        self.set_value(value);
-                    }
+        let _ = self.lock.lock();
 
-                    Ok(())
-                } else {
-                    Err(SetError::AlreadyInitializedError(value))
-                }
-            }
-            Err(_) => {
-                // If we failed to lock the mutex, it means some other task is
-                // trying to set the value, so we return an error.
-                Err(SetError::InitializingError(value))
-            }
+        unsafe {
+            self.set_value(value);
         }
+
+        Ok(())
     }
 
     /// Takes the value from the cell, destroying the cell in the process.
@@ -341,41 +331,28 @@ impl<T> SetOnce<T> {
     /// the `SetOnce` is initialized.
     ///
     /// If the `SetOnce` is already initialized, it will return the value
-    // immediately.
+    /// immediately.
     ///
     /// # Panics
     ///
     /// If the `SetOnce` is not initialized after waiting, it will panic. To
     /// avoid this, use `get_wait()` which returns an `Option<&T>` instead of
     /// `&T`.
     pub async fn wait(&self) -> &T {
-        match self.get_wait().await {
-            Some(val) => val,
-            _ => panic!("SetOnce::wait called but the SetOnce is not initialized"),
-        }
-    }
-
-    /// Waits until set is called.
-    ///
-    /// If the state failed to initalize it will return `None`.
-    pub async fn get_wait(&self) -> Option<&T> {
         let notify_fut = self.notify.notified();
         pin!(notify_fut);
 
         if self.value_set.load(Ordering::Acquire) {
             // SAFETY: the state is initialized
-            return Some(unsafe { self.get_unchecked() });
+            return unsafe { self.get_unchecked() };
         }
+
         // wait until the value is set
         (&mut notify_fut).await;
 
-        // look at the state again
-        if self.value_set.load(Ordering::Acquire) {
-            // SAFETY: the state is initialized
-            return Some(unsafe { self.get_unchecked() });
-        }
-
-        None
+        // SAFETY: Its not possible for the state to initialize after
+        // waker is notified
+        unsafe { self.get_unchecked() }
     }
 }