Update tokio::sync::SetOnce with AtomicU8

Update `tokio::sync::SetOnce` struct to internally use a `AtomicU8`
to keep track of `AtomicState`, I introduced a
`pub(crate) AtomicState` enum to keep readable track of the State.
It is a `#[repr(u8)]` enum to keep track of u8 values the state
might be in.

The drop implementation is changed to use ptr::drop_in_place. We cast the
 `MaybeUnint<T>` to `*mut T`.

Author: Daksh14

tokio/src/sync/set_once.rs

@@ -4,7 +4,7 @@ use std::fmt;
 use std::mem::MaybeUninit;
 use std::ops::Drop;
 use std::ptr;
-use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::atomic::{AtomicU8, Ordering};
 
 // This file contains an implementation of an SetOnce. The value of SetOnce
 // can only be modified once during initialization.
@@ -70,11 +70,22 @@ use std::sync::atomic::{AtomicBool, Ordering};
 /// }
 /// ```
 pub struct SetOnce<T> {
-    value_set: AtomicBool,
+    // AtomicU8 has 3 states, EMPTY, SETTING and DONE
+    // is described by the `AtomicState` enum.
+    value_set: AtomicU8,
     value: UnsafeCell<MaybeUninit<T>>,
     notify: Notify,
 }
 
+#[repr(u8)]
+#[derive(PartialEq, Eq)]
+pub(crate) enum AtomicState {
+    Empty = 0,
+    // SETTING means "set is currently running"
+    Setting = 1,
+    Done = 2,
+}
+
 impl<T> Default for SetOnce<T> {
     fn default() -> SetOnce<T> {
         SetOnce::new()
@@ -106,22 +117,22 @@ impl<T: Eq> Eq for SetOnce<T> {}
 impl<T> Drop for SetOnce<T> {
     fn drop(&mut self) {
         if self.initialized() {
-            // SAFETY: We're inside the drop implementation of SetOnce
-            // AND we're also initalized. This is the best way to ensure
-            // out data gets dropped
+            // SAFETY: Inside the drop implementation of SetOnce also
+            // initalized. we cast `MaybeUninit<T>` to *mut T
+            // because `MaybeUninit<T>` doesnt have a destructor
             unsafe {
-                let _ = self.value.with_mut(|ptr| ptr::read(ptr).assume_init());
+                self.value.with_mut(|ptr| ptr::drop_in_place(ptr as *mut T));
             }
-            // no need to set the flag to false as this set once is being
-            // dropped
+            // no need to change value_set here as SetOnce is being dropped
         }
     }
 }
 
 impl<T> From<T> for SetOnce<T> {
     fn from(value: T) -> Self {
         SetOnce {
-            value_set: AtomicBool::new(true),
+            // value has been initialized
+            value_set: AtomicU8::new(AtomicState::Done as _),
             value: UnsafeCell::new(MaybeUninit::new(value)),
             notify: Notify::new(),
         }
@@ -132,7 +143,7 @@ impl<T> SetOnce<T> {
     /// Creates a new empty `SetOnce` instance.
     pub fn new() -> Self {
         Self {
-            value_set: AtomicBool::new(false),
+            value_set: AtomicU8::new(AtomicState::Empty as _),
             value: UnsafeCell::new(MaybeUninit::uninit()),
             notify: Notify::new(),
         }
@@ -174,7 +185,7 @@ impl<T> SetOnce<T> {
     #[cfg(not(all(loom, test)))]
     pub const fn const_new() -> Self {
         Self {
-            value_set: AtomicBool::new(false),
+            value_set: AtomicU8::new(AtomicState::Empty as _),
             value: UnsafeCell::new(MaybeUninit::uninit()),
             notify: Notify::const_new(),
         }
@@ -224,7 +235,7 @@ impl<T> SetOnce<T> {
     #[cfg(not(all(loom, test)))]
     pub const fn const_new_with(value: T) -> Self {
         Self {
-            value_set: AtomicBool::new(true),
+            value_set: AtomicU8::new(AtomicState::Done as _),
             value: UnsafeCell::new(MaybeUninit::new(value)),
             notify: Notify::const_new(),
         }
@@ -235,7 +246,7 @@ impl<T> SetOnce<T> {
     pub fn initialized(&self) -> bool {
         // Using acquire ordering so we're able to read/catch any writes that
         // are done with `Ordering::Release`
-        self.value_set.load(Ordering::Acquire)
+        self.value_set.load(Ordering::Acquire) == AtomicState::Done as _
     }
 
     // SAFETY: The SetOnce must not be empty.
@@ -254,13 +265,19 @@ impl<T> SetOnce<T> {
     }
 
     // SAFETY: The caller of this function needs to ensure that this function is
-    // called only when the value_set AtomicBool is flipped from FALSE to TRUE
-    // meaning that the value is being set from uinitialized to initialized via
-    // this function
+    // called only when the self.value_set AtomicUsize is flipped from
+    // AtomicState::Empty to AtomicState::Setting meaning that the value is
+    // being set from uinitialized to initialized via this function.
+    //
+    // It needs to be done in the order to ensure the safety of this function
+    // across threads, this function will mutate the value_set to be
+    // AtomicState::Done.
     unsafe fn set_value(&self, value: T) {
-        unsafe {
-            self.value.with_mut(|ptr| (*ptr).as_mut_ptr().write(value));
-        }
+        self.value.with_mut(|ptr| (*ptr).as_mut_ptr().write(value));
+
+        // Store that we're done in Release so any `Acquire` can see
+        self.value_set
+            .store(AtomicState::Done as _, Ordering::Release);
 
         // notify the waiting wakers that the value is set
         self.notify.notify_waiters();
@@ -282,19 +299,35 @@ impl<T> SetOnce<T> {
             return Err(SetError::AlreadyInitializedError(value));
         }
 
-        // 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);
-            }
+        // to capture any writes that are done with `Ordering::Release`
+        let ordering = Ordering::Acquire;
 
-            Ok(())
-        } else {
-            Err(SetError::InitializingError(value))
+        match self.value_set.compare_exchange(
+            // we expect the set to be uninitialized
+            AtomicState::Empty as _,
+            // we want to lock the Set
+            AtomicState::Setting as _,
+            ordering,
+            ordering,
+        ) {
+            Ok(_) => {
+                // SAFETY: We changed the state from EMPTY to SETTING,
+                // this declares the set to be currently running
+                unsafe {
+                    self.set_value(value);
+                }
+
+                Ok(())
+            }
+            Err(true_val) => {
+                // The value is being currently set in another thread
+                if true_val == AtomicState::Setting as _ {
+                    Err(SetError::InitializingError(value))
+                } else {
+                    // here means true_val is either Done or higher
+                    Err(SetError::AlreadyInitializedError(value))
+                }
+            }
         }
     }
 
@@ -304,9 +337,9 @@ impl<T> SetOnce<T> {
         if self.initialized() {
             // Since we have taken ownership of self, its drop implementation
             // will be called by the end of this function, to prevent a double
-            // free we will set the value_set to false so that the drop
+            // free we will set the value_set to EMPTY so that the drop
             // implementation does not try to drop the value again.
-            *self.value_set.get_mut() = false;
+            *self.value_set.get_mut() = AtomicState::Empty as _;
 
             // SAFETY: The SetOnce is currently initialized, we can assume the
             // value is initialized and return that, when we return the value
@@ -337,7 +370,7 @@ impl<T> SetOnce<T> {
 // (of type T) across threads.
 unsafe impl<T: Sync + Send> Sync for SetOnce<T> {}
 
-// Access to SetOnce's value is guarded by the Atomic boolean flag
+// Access to SetOnce's value is guarded by the AtomicUsize state
 // and atomic operations on `value_set`, so as long as T itself is Send
 // it's safe to send it to another thread
 unsafe impl<T: Send> Send for SetOnce<T> {}