sync: add SetOnce (#7418)

tokio/src/sync/mod.rs

@@ -488,6 +488,9 @@ cfg_sync! {
     mod once_cell;
     pub use self::once_cell::{OnceCell, SetError};
 
+    mod set_once;
+    pub use self::set_once::{SetOnce, SetOnceError};
+
     pub mod watch;
 }
 

tokio/src/sync/set_once.rs

@@ -0,0 +1,404 @@
+use super::Notify;
+
+use crate::loom::cell::UnsafeCell;
+use crate::loom::sync::{atomic::AtomicBool, Mutex};
+
+use std::error::Error;
+use std::fmt;
+use std::mem::MaybeUninit;
+use std::ops::Drop;
+use std::ptr;
+use std::sync::atomic::Ordering;
+
+// This file contains an implementation of an SetOnce. The value of SetOnce
+// can only be modified once during initialization.
+//
+//  1. When `value_set` is false, the `value` is not initialized and wait()
+//      future will keep on waiting.
+//  2. When `value_set` is true, the wait() future completes, get() will return
+//      Some(&T)
+//
+// The value cannot be changed after set() is called. Subsequent calls to set()
+// will return a `SetOnceError`.
+
+/// A thread-safe cell that can be written to only once.
+///
+/// A `SetOnce` is inspired from python's [`asyncio.Event`] type. It can be
+/// used to wait until the value of the `SetOnce` is set like a "Event" mechanism.
+///
+/// # Example
+///
+/// ```
+/// use tokio::sync::{SetOnce, SetOnceError};
+///
+/// static ONCE: SetOnce<u32> = SetOnce::const_new();
+///
+/// #[tokio::main]
+/// async fn main() -> Result<(), SetOnceError<u32>> {
+///
+///     // set the value inside a task somewhere...
+///     tokio::spawn(async move { ONCE.set(20) });
+///
+///     // checking with .get doesn't block main thread
+///     println!("{:?}", ONCE.get());
+///
+///     // wait until the value is set, blocks the thread
+///     println!("{:?}", ONCE.wait().await);
+///
+///     Ok(())
+/// }
+/// ```
+///
+/// A `SetOnce` is typically used for global variables that need to be
+/// initialized once on first use, but need no further changes. The `SetOnce`
+/// in Tokio allows the initialization procedure to be asynchronous.
+///
+/// # Example
+///
+/// ```
+/// use tokio::sync::{SetOnce, SetOnceError};
+/// use std::sync::Arc;
+///
+/// #[tokio::main]
+/// async fn main() -> Result<(), SetOnceError<u32>> {
+///     let once = SetOnce::new();
+///
+///     let arc = Arc::new(once);
+///     let first_cl = Arc::clone(&arc);
+///     let second_cl = Arc::clone(&arc);
+///
+///     // set the value inside a task
+///     tokio::spawn(async move { first_cl.set(20) }).await.unwrap()?;
+///
+///     // wait inside task to not block the main thread
+///     tokio::spawn(async move {
+///         // wait inside async context for the value to be set
+///         assert_eq!(*second_cl.wait().await, 20);
+///     }).await.unwrap();
+///
+///     // subsequent set calls will fail
+///     assert!(arc.set(30).is_err());
+///
+///     println!("{:?}", arc.get());
+///
+///     Ok(())
+/// }
+/// ```
+///
+/// [`asyncio.Event`]: https://docs.python.org/3/library/asyncio-event.html
+pub struct SetOnce<T> {
+    value_set: AtomicBool,
+    value: UnsafeCell<MaybeUninit<T>>,
+    notify: Notify,
+    // we lock the mutex inside set to ensure
+    // only one caller of set can run at a time
+    lock: Mutex<()>,
+}
+
+impl<T> Default for SetOnce<T> {
+    fn default() -> SetOnce<T> {
+        SetOnce::new()
+    }
+}
+
+impl<T: fmt::Debug> fmt::Debug for SetOnce<T> {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt.debug_struct("SetOnce")
+            .field("value", &self.get())
+            .finish()
+    }
+}
+
+impl<T: Clone> Clone for SetOnce<T> {
+    fn clone(&self) -> SetOnce<T> {
+        SetOnce::new_with(self.get().cloned())
+    }
+}
+
+impl<T: PartialEq> PartialEq for SetOnce<T> {
+    fn eq(&self, other: &SetOnce<T>) -> bool {
+        self.get() == other.get()
+    }
+}
+
+impl<T: Eq> Eq for SetOnce<T> {}
+
+impl<T> Drop for SetOnce<T> {
+    fn drop(&mut self) {
+        // TODO: Use get_mut()
+        if self.value_set.load(Ordering::Relaxed) {
+            // SAFETY: If the value_set is true, then the value is initialized
+            // then there is a value to be dropped and this is safe
+            unsafe { self.value.with_mut(|ptr| ptr::drop_in_place(ptr as *mut T)) }
+        }
+    }
+}
+
+impl<T> From<T> for SetOnce<T> {
+    fn from(value: T) -> Self {
+        SetOnce {
+            value_set: AtomicBool::new(true),
+            value: UnsafeCell::new(MaybeUninit::new(value)),
+            notify: Notify::new(),
+            lock: Mutex::new(()),
+        }
+    }
+}
+
+impl<T> SetOnce<T> {
+    /// Creates a new empty `SetOnce` instance.
+    pub fn new() -> Self {
+        Self {
+            value_set: AtomicBool::new(false),
+            value: UnsafeCell::new(MaybeUninit::uninit()),
+            notify: Notify::new(),
+            lock: Mutex::new(()),
+        }
+    }
+
+    /// Creates a new empty `SetOnce` instance.
+    ///
+    /// Equivalent to `SetOnce::new`, except that it can be used in static
+    /// variables.
+    ///
+    /// When using the `tracing` [unstable feature], a `SetOnce` created with
+    /// `const_new` will not be instrumented. As such, it will not be visible
+    /// in [`tokio-console`]. Instead, [`SetOnce::new`] should be used to
+    /// create an instrumented object if that is needed.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use tokio::sync::{SetOnce, SetOnceError};
+    ///
+    /// static ONCE: SetOnce<u32> = SetOnce::const_new();
+    ///
+    /// fn get_global_integer() -> Result<Option<&'static u32>, SetOnceError<u32>> {
+    ///     ONCE.set(2)?;
+    ///     Ok(ONCE.get())
+    /// }
+    ///
+    /// #[tokio::main]
+    /// async fn main() -> Result<(), SetOnceError<u32>> {
+    ///     let result = get_global_integer()?;
+    ///
+    ///     assert_eq!(result, Some(&2));
+    ///     Ok(())
+    /// }
+    /// ```
+    ///
+    /// [`tokio-console`]: https://github.com/tokio-rs/console
+    /// [unstable feature]: crate#unstable-features
+    #[cfg(not(all(loom, test)))]
+    pub const fn const_new() -> Self {
+        Self {
+            value_set: AtomicBool::new(false),
+            value: UnsafeCell::new(MaybeUninit::uninit()),
+            notify: Notify::const_new(),
+            lock: Mutex::const_new(()),
+        }
+    }
+
+    /// Creates a new `SetOnce` that contains the provided value, if any.
+    ///
+    /// If the `Option` is `None`, this is equivalent to `SetOnce::new`.
+    ///
+    /// [`SetOnce::new`]: crate::sync::SetOnce::new
+    pub fn new_with(value: Option<T>) -> Self {
+        if let Some(v) = value {
+            SetOnce::from(v)
+        } else {
+            SetOnce::new()
+        }
+    }
+
+    /// Creates a new `SetOnce` that contains the provided value.
+    ///
+    /// # Example
+    ///
+    /// When using the `tracing` [unstable feature], a `SetOnce` created with
+    /// `const_new_with` will not be instrumented. As such, it will not be
+    /// visible in [`tokio-console`]. Instead, [`SetOnce::new_with`] should be
+    /// used to create an instrumented object if that is needed.
+    ///
+    /// ```
+    /// use tokio::sync::SetOnce;
+    ///
+    /// static ONCE: SetOnce<u32> = SetOnce::const_new_with(1);
+    ///
+    /// fn get_global_integer() -> Option<&'static u32> {
+    ///     ONCE.get()
+    /// }
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let result = get_global_integer();
+    ///
+    ///     assert_eq!(result, Some(&1));
+    /// }
+    /// ```
+    ///
+    /// [`tokio-console`]: https://github.com/tokio-rs/console
+    /// [unstable feature]: crate#unstable-features
+    #[cfg(not(all(loom, test)))]
+    pub const fn const_new_with(value: T) -> Self {
+        Self {
+            value_set: AtomicBool::new(true),
+            value: UnsafeCell::new(MaybeUninit::new(value)),
+            notify: Notify::const_new(),
+            lock: Mutex::const_new(()),
+        }
+    }
+
+    /// Returns `true` if the `SetOnce` currently contains a value, and `false`
+    /// otherwise.
+    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)
+    }
+
+    // SAFETY: The SetOnce must not be empty.
+    unsafe fn get_unchecked(&self) -> &T {
+        &*self.value.with(|ptr| (*ptr).as_ptr())
+    }
+
+    /// Returns a reference to the value currently stored in the `SetOnce`, or
+    /// `None` if the `SetOnce` is empty.
+    pub fn get(&self) -> Option<&T> {
+        if self.initialized() {
+            // SAFETY: the SetOnce is initialized, so we can safely
+            // call get_unchecked and return the value
+            Some(unsafe { self.get_unchecked() })
+        } else {
+            None
+        }
+    }
+
+    /// Sets the value of the `SetOnce` to the given value if the `SetOnce` is
+    /// empty.
+    ///
+    /// If the `SetOnce` already has a value, this call will fail with an
+    /// [`SetOnceError`].
+    ///
+    /// [`SetOnceError`]: crate::sync::SetOnceError
+    pub fn set(&self, value: T) -> Result<(), SetOnceError<T>> {
+        if self.initialized() {
+            return Err(SetOnceError(value));
+        }
+
+        // SAFETY: lock the mutex to ensure only one caller of set
+        // can run at a time.
+        let guard = self.lock.lock();
+
+        if self.initialized() {
+            // If the value is already set, we return an error
+            drop(guard);
+
+            return Err(SetOnceError(value));
+        }
+
+        // SAFETY: We have locked the mutex and checked if the value is
+        // initalized or not, so we can safely write to the value
+        unsafe {
+            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);
+
+        drop(guard);
+
+        // notify the waiting wakers that the value is set
+        self.notify.notify_waiters();
+
+        Ok(())
+    }
+
+    /// Takes the value from the cell, destroying the cell in the process.
+    /// Returns `None` if the cell is empty.
+    pub fn into_inner(self) -> Option<T> {
+        // TODO: Use get_mut()
+        let value_set = self.value_set.load(Ordering::Relaxed);
+
+        if value_set {
+            // 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
+            // implementation does not try to drop the value again.
+            self.value_set.store(false, Ordering::Relaxed);
+
+            // SAFETY: The SetOnce is currently initialized, we can assume the
+            // value is initialized and return that, when we return the value
+            // we give the drop handler to the return scope.
+            Some(unsafe { self.value.with_mut(|ptr| ptr::read(ptr).assume_init()) })
+        } else {
+            None
+        }
+    }
+
+    /// Waits until set is called. The future returned will keep blocking until
+    /// the `SetOnce` is initialized.
+    ///
+    /// If the `SetOnce` is already initialized, it will return the value
+    /// immediately.
+    ///
+    /// # Note
+    ///
+    /// This will keep waiting until the `SetOnce` is initialized, so it
+    /// should be used with care to avoid blocking the current task
+    /// indefinitely.
+    pub async fn wait(&self) -> &T {
+        loop {
+            if let Some(val) = self.get() {
+                return val;
+            }
+
+            let notify_fut = self.notify.notified();
+            {
+                // Taking the lock here ensures that a concurrent call to `set`
+                // will see the creation of `notify_fut` in case the check
+                // fails.
+                let _guard = self.lock.lock();
+
+                if self.value_set.load(Ordering::Relaxed) {
+                    // SAFETY: the state is initialized
+                    return unsafe { self.get_unchecked() };
+                }
+            }
+
+            // wait until the value is set
+            notify_fut.await;
+        }
+    }
+}
+
+// Since `get` gives us access to immutable references of the SetOnce, SetOnce
+// can only be Sync if T is Sync, otherwise SetOnce would allow sharing
+// references of !Sync values across threads. We need T to be Send in order for
+// SetOnce to by Sync because we can use `set` on `&SetOnce<T>` to send values
+// (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
+// 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> {}
+
+/// Error that can be returned from [`SetOnce::set`].
+///
+/// This error means that the `SetOnce` was already initialized when
+/// set was called
+///
+/// [`SetOnce::set`]: crate::sync::SetOnce::set
+#[derive(Debug, PartialEq, Eq)]
+pub struct SetOnceError<T>(pub T);
+
+impl<T> fmt::Display for SetOnceError<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "SetOnceError")
+    }
+}
+
+impl<T: fmt::Debug> Error for SetOnceError<T> {}

tokio/src/sync/tests/loom_set_once.rs

@@ -0,0 +1,72 @@
+use crate::sync::SetOnce;
+
+use loom::future::block_on;
+use loom::sync::atomic::AtomicU32;
+use loom::thread;
+use std::sync::atomic::Ordering;
+use std::sync::Arc;
+
+#[derive(Clone)]
+struct DropCounter {
+    pub drops: Arc<AtomicU32>,
+}
+
+impl DropCounter {
+    pub fn new() -> Self {
+        Self {
+            drops: Arc::new(AtomicU32::new(0)),
+        }
+    }
+
+    fn assert_num_drops(&self, value: u32) {
+        assert_eq!(value, self.drops.load(Ordering::Relaxed));
+    }
+}
+
+impl Drop for DropCounter {
+    fn drop(&mut self) {
+        self.drops.fetch_add(1, Ordering::Relaxed);
+    }
+}
+
+#[test]
+fn set_once_drop_test() {
+    loom::model(|| {
+        let set_once = Arc::new(SetOnce::new());
+        let set_once_clone = Arc::clone(&set_once);
+
+        let drop_counter = DropCounter::new();
+        let counter_cl = drop_counter.clone();
+
+        let thread = thread::spawn(move || set_once_clone.set(counter_cl).is_ok());
+
+        let foo = drop_counter.clone();
+
+        let set = set_once.set(foo).is_ok();
+        let res = thread.join().unwrap();
+
+        drop(set_once);
+
+        drop_counter.assert_num_drops(2);
+        assert!(res != set);
+    });
+}
+
+#[test]
+fn set_once_wait_test() {
+    loom::model(|| {
+        let tx = Arc::new(SetOnce::new());
+        let rx_one = tx.clone();
+        let rx_two = tx.clone();
+
+        let thread = thread::spawn(move || {
+            assert!(rx_one.set(2).is_ok());
+        });
+
+        block_on(async {
+            assert_eq!(*rx_two.wait().await, 2);
+        });
+
+        thread.join().unwrap();
+    });
+}

tokio/src/sync/tests/mod.rs

@@ -14,4 +14,5 @@ cfg_loom! {
     mod loom_semaphore_batch;
     mod loom_watch;
     mod loom_rwlock;
+    mod loom_set_once;
 }

tokio/tests/async_send_sync.rs

@@ -350,6 +350,9 @@ assert_value!(tokio::sync::Notify: Send & Sync & Unpin);
 assert_value!(tokio::sync::OnceCell<NN>: !Send & !Sync & Unpin);
 assert_value!(tokio::sync::OnceCell<YN>: Send & !Sync & Unpin);
 assert_value!(tokio::sync::OnceCell<YY>: Send & Sync & Unpin);
+assert_value!(tokio::sync::SetOnce<NN>: !Send & !Sync & Unpin);
+assert_value!(tokio::sync::SetOnce<YN>: Send & !Sync & Unpin);
+assert_value!(tokio::sync::SetOnce<YY>: Send & Sync & Unpin);
 assert_value!(tokio::sync::OwnedMutexGuard<NN>: !Send & !Sync & Unpin);
 assert_value!(tokio::sync::OwnedMutexGuard<YN>: Send & !Sync & Unpin);
 assert_value!(tokio::sync::OwnedMutexGuard<YY>: Send & Sync & Unpin);
@@ -481,6 +484,9 @@ async_assert_fn!(tokio::sync::OnceCell<YY>::get_or_init( _, fn() -> Pin<Box<dyn
 async_assert_fn!(tokio::sync::OnceCell<YY>::get_or_try_init( _, fn() -> Pin<Box<dyn Future<Output = std::io::Result<YY>> + Send + Sync>>): Send & Sync & !Unpin);
 async_assert_fn!(tokio::sync::OnceCell<YY>::get_or_try_init( _, fn() -> Pin<Box<dyn Future<Output = std::io::Result<YY>> + Send>>): Send & !Sync & !Unpin);
 async_assert_fn!(tokio::sync::OnceCell<YY>::get_or_try_init( _, fn() -> Pin<Box<dyn Future<Output = std::io::Result<YY>>>>): !Send & !Sync & !Unpin);
+async_assert_fn!(tokio::sync::SetOnce<NN>::wait(_): !Send & !Sync & !Unpin);
+async_assert_fn!(tokio::sync::SetOnce<YN>::wait(_): !Send & !Sync & !Unpin);
+async_assert_fn!(tokio::sync::SetOnce<YY>::wait(_): Send & Sync & !Unpin);
 async_assert_fn!(tokio::sync::RwLock<NN>::read(_): !Send & !Sync & !Unpin);
 async_assert_fn!(tokio::sync::RwLock<NN>::write(_): !Send & !Sync & !Unpin);
 async_assert_fn!(tokio::sync::RwLock<YN>::read(_): !Send & !Sync & !Unpin);

tokio/tests/sync_set_once.rs

@@ -0,0 +1,180 @@
+#![warn(rust_2018_idioms)]
+#![cfg(feature = "full")]
+
+use std::sync::{
+    atomic::{AtomicU32, Ordering},
+    Arc,
+};
+use tokio::sync::SetOnce;
+
+#[derive(Clone)]
+struct DropCounter {
+    drops: Arc<AtomicU32>,
+}
+
+impl DropCounter {
+    fn new() -> Self {
+        DropCounter {
+            drops: Arc::new(AtomicU32::new(0)),
+        }
+    }
+
+    fn assert_num_drops(&self, value: u32) {
+        assert_eq!(value, self.drops.load(Ordering::Relaxed));
+    }
+}
+
+impl Drop for DropCounter {
+    fn drop(&mut self) {
+        self.drops.fetch_add(1, Ordering::Relaxed);
+    }
+}
+
+#[test]
+fn drop_cell() {
+    let fooer = DropCounter::new();
+    let fooer_cl = fooer.clone();
+
+    {
+        let once_cell = SetOnce::new();
+        let prev = once_cell.set(fooer_cl);
+        assert!(prev.is_ok())
+    }
+
+    fooer.assert_num_drops(1);
+}
+
+#[test]
+fn drop_cell_new_with() {
+    let fooer = DropCounter::new();
+
+    {
+        let once_cell = SetOnce::new_with(Some(fooer.clone()));
+        assert!(once_cell.initialized());
+    }
+
+    fooer.assert_num_drops(1);
+}
+
+#[test]
+fn drop_into_inner() {
+    let fooer = DropCounter::new();
+
+    let once_cell = SetOnce::new();
+    assert!(once_cell.set(fooer.clone()).is_ok());
+    let val = once_cell.into_inner();
+    fooer.assert_num_drops(0);
+    drop(val);
+    fooer.assert_num_drops(1);
+}
+
+#[test]
+fn drop_into_inner_new_with() {
+    let fooer = DropCounter::new();
+
+    let once_cell = SetOnce::new_with(Some(fooer.clone()));
+    let val = once_cell.into_inner();
+    fooer.assert_num_drops(0);
+    drop(val);
+    fooer.assert_num_drops(1);
+}
+
+#[test]
+fn from() {
+    let cell = SetOnce::from(2);
+    assert_eq!(*cell.get().unwrap(), 2);
+}
+
+#[test]
+fn set_and_get() {
+    static ONCE: SetOnce<u32> = SetOnce::const_new();
+
+    ONCE.set(5).unwrap();
+    let value = ONCE.get().unwrap();
+    assert_eq!(*value, 5);
+}
+
+#[tokio::test]
+async fn set_and_wait() {
+    static ONCE: SetOnce<u32> = SetOnce::const_new();
+
+    tokio::spawn(async { ONCE.set(5) });
+
+    let value = ONCE.wait().await;
+    assert_eq!(*value, 5);
+}
+
+#[test]
+#[cfg_attr(target_family = "wasm", ignore)]
+fn set_and_wait_multiple_threads() {
+    static ONCE: SetOnce<u32> = SetOnce::const_new();
+
+    let res1 = std::thread::spawn(|| ONCE.set(4));
+
+    let res2 = std::thread::spawn(|| ONCE.set(3));
+
+    let result_first = res1.join().unwrap().is_err();
+    let result_two = res2.join().unwrap().is_err();
+
+    assert!(result_first != result_two);
+}
+
+#[tokio::test]
+#[cfg_attr(target_family = "wasm", ignore)]
+async fn set_and_wait_threads() {
+    static ONCE: SetOnce<u32> = SetOnce::const_new();
+
+    let thread = std::thread::spawn(|| {
+        ONCE.set(4).unwrap();
+    });
+
+    let value = ONCE.wait().await;
+    thread.join().unwrap();
+    assert_eq!(*value, 4);
+}
+
+#[test]
+fn get_uninit() {
+    static ONCE: SetOnce<u32> = SetOnce::const_new();
+    let uninit = ONCE.get();
+    assert!(uninit.is_none());
+}
+
+#[test]
+fn set_twice() {
+    static ONCE: SetOnce<u32> = SetOnce::const_new();
+
+    let first = ONCE.set(5);
+    assert_eq!(first, Ok(()));
+    let second = ONCE.set(6);
+    assert!(second.is_err());
+}
+
+#[test]
+fn is_none_initializing() {
+    static ONCE: SetOnce<u32> = SetOnce::const_new();
+
+    assert_eq!(ONCE.get(), None);
+
+    ONCE.set(20).unwrap();
+
+    assert!(ONCE.set(10).is_err());
+}
+
+#[tokio::test]
+async fn is_some_initializing() {
+    static ONCE: SetOnce<u32> = SetOnce::const_new();
+
+    tokio::spawn(async { ONCE.set(20) });
+
+    assert_eq!(*ONCE.wait().await, 20);
+}
+
+#[test]
+fn into_inner_int_empty_setonce() {
+    let once = SetOnce::<u32>::new();
+
+    let val = once.into_inner();
+
+    assert!(val.is_none());
+}