From 911ab21d7012a50e53971ad1292a9f18de22d4c8 Mon Sep 17 00:00:00 2001 From: Daksh <41485688+Daksh14@users.noreply.github.com> Date: Fri, 25 Jul 2025 04:29:14 -0400 Subject: [PATCH] sync: add `SetOnce` (#7418) --- tokio/src/sync/mod.rs | 3 + tokio/src/sync/set_once.rs | 404 ++++++++++++++++++++++++++ tokio/src/sync/tests/loom_set_once.rs | 72 +++++ tokio/src/sync/tests/mod.rs | 1 + tokio/tests/async_send_sync.rs | 6 + tokio/tests/sync_set_once.rs | 180 ++++++++++++ 6 files changed, 666 insertions(+) create mode 100644 tokio/src/sync/set_once.rs create mode 100644 tokio/src/sync/tests/loom_set_once.rs create mode 100644 tokio/tests/sync_set_once.rs diff --git a/tokio/src/sync/mod.rs b/tokio/src/sync/mod.rs index 3a6c8b827..a2502a76e 100644 --- a/tokio/src/sync/mod.rs +++ b/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; } diff --git a/tokio/src/sync/set_once.rs b/tokio/src/sync/set_once.rs new file mode 100644 index 000000000..017bef547 --- /dev/null +++ b/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 = SetOnce::const_new(); +/// +/// #[tokio::main] +/// async fn main() -> Result<(), SetOnceError> { +/// +/// // 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> { +/// 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 { + value_set: AtomicBool, + value: UnsafeCell>, + notify: Notify, + // we lock the mutex inside set to ensure + // only one caller of set can run at a time + lock: Mutex<()>, +} + +impl Default for SetOnce { + fn default() -> SetOnce { + SetOnce::new() + } +} + +impl fmt::Debug for SetOnce { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("SetOnce") + .field("value", &self.get()) + .finish() + } +} + +impl Clone for SetOnce { + fn clone(&self) -> SetOnce { + SetOnce::new_with(self.get().cloned()) + } +} + +impl PartialEq for SetOnce { + fn eq(&self, other: &SetOnce) -> bool { + self.get() == other.get() + } +} + +impl Eq for SetOnce {} + +impl Drop for SetOnce { + 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 From for SetOnce { + fn from(value: T) -> Self { + SetOnce { + value_set: AtomicBool::new(true), + value: UnsafeCell::new(MaybeUninit::new(value)), + notify: Notify::new(), + lock: Mutex::new(()), + } + } +} + +impl SetOnce { + /// 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 = SetOnce::const_new(); + /// + /// fn get_global_integer() -> Result, SetOnceError> { + /// ONCE.set(2)?; + /// Ok(ONCE.get()) + /// } + /// + /// #[tokio::main] + /// async fn main() -> Result<(), SetOnceError> { + /// 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) -> 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 = 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> { + 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 { + // 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` to send values +// (of type T) across threads. +unsafe impl Sync for SetOnce {} + +// 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 Send for SetOnce {} + +/// 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(pub T); + +impl fmt::Display for SetOnceError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "SetOnceError") + } +} + +impl Error for SetOnceError {} diff --git a/tokio/src/sync/tests/loom_set_once.rs b/tokio/src/sync/tests/loom_set_once.rs new file mode 100644 index 000000000..957d6f92f --- /dev/null +++ b/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, +} + +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(); + }); +} diff --git a/tokio/src/sync/tests/mod.rs b/tokio/src/sync/tests/mod.rs index ee76418ac..9f082b42d 100644 --- a/tokio/src/sync/tests/mod.rs +++ b/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; } diff --git a/tokio/tests/async_send_sync.rs b/tokio/tests/async_send_sync.rs index c9cedc38b..7d4bb9bcd 100644 --- a/tokio/tests/async_send_sync.rs +++ b/tokio/tests/async_send_sync.rs @@ -350,6 +350,9 @@ assert_value!(tokio::sync::Notify: Send & Sync & Unpin); assert_value!(tokio::sync::OnceCell: !Send & !Sync & Unpin); assert_value!(tokio::sync::OnceCell: Send & !Sync & Unpin); assert_value!(tokio::sync::OnceCell: Send & Sync & Unpin); +assert_value!(tokio::sync::SetOnce: !Send & !Sync & Unpin); +assert_value!(tokio::sync::SetOnce: Send & !Sync & Unpin); +assert_value!(tokio::sync::SetOnce: Send & Sync & Unpin); assert_value!(tokio::sync::OwnedMutexGuard: !Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedMutexGuard: Send & !Sync & Unpin); assert_value!(tokio::sync::OwnedMutexGuard: Send & Sync & Unpin); @@ -481,6 +484,9 @@ async_assert_fn!(tokio::sync::OnceCell::get_or_init( _, fn() -> Pin::get_or_try_init( _, fn() -> Pin> + Send + Sync>>): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_try_init( _, fn() -> Pin> + Send>>): Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::OnceCell::get_or_try_init( _, fn() -> Pin>>>): !Send & !Sync & !Unpin); +async_assert_fn!(tokio::sync::SetOnce::wait(_): !Send & !Sync & !Unpin); +async_assert_fn!(tokio::sync::SetOnce::wait(_): !Send & !Sync & !Unpin); +async_assert_fn!(tokio::sync::SetOnce::wait(_): Send & Sync & !Unpin); async_assert_fn!(tokio::sync::RwLock::read(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::RwLock::write(_): !Send & !Sync & !Unpin); async_assert_fn!(tokio::sync::RwLock::read(_): !Send & !Sync & !Unpin); diff --git a/tokio/tests/sync_set_once.rs b/tokio/tests/sync_set_once.rs new file mode 100644 index 000000000..5b6d88a9f --- /dev/null +++ b/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, +} + +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 = 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 = 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 = 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 = 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 = SetOnce::const_new(); + let uninit = ONCE.get(); + assert!(uninit.is_none()); +} + +#[test] +fn set_twice() { + static ONCE: SetOnce = 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 = 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 = 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::::new(); + + let val = once.into_inner(); + + assert!(val.is_none()); +}