From 37ca2f049c552fde50a2535de6dfc61ee6e96eed Mon Sep 17 00:00:00 2001 From: Sam <130320493+srxg@users.noreply.github.com> Date: Wed, 3 Sep 2025 16:37:50 +0100 Subject: [PATCH] sync: remove inner mutex in `SetOnce` (#7554) --- tokio/src/sync/notify.rs | 43 +++++++++++++++++++++++++++++++++----- tokio/src/sync/set_once.rs | 43 +++++++++++++------------------------- 2 files changed, 53 insertions(+), 33 deletions(-) diff --git a/tokio/src/sync/notify.rs b/tokio/src/sync/notify.rs index d46079793..b79c14d8e 100644 --- a/tokio/src/sync/notify.rs +++ b/tokio/src/sync/notify.rs @@ -741,12 +741,14 @@ impl Notify { /// } /// ``` pub fn notify_waiters(&self) { - let mut waiters = self.waiters.lock(); - - // The state must be loaded while the lock is held. The state may only - // transition out of WAITING while the lock is held. - let curr = self.state.load(SeqCst); + self.lock_waiter_list().notify_waiters(); + } + fn inner_notify_waiters<'a>( + &'a self, + curr: usize, + mut waiters: crate::loom::sync::MutexGuard<'a, LinkedList>, + ) { if matches!(get_state(curr), EMPTY | NOTIFIED) { // There are no waiting tasks. All we need to do is increment the // number of times this method was called. @@ -814,6 +816,20 @@ impl Notify { wakers.wake_all(); } + + pub(crate) fn lock_waiter_list(&self) -> NotifyGuard<'_> { + let guarded_waiters = self.waiters.lock(); + + // The state must be loaded while the lock is held. The state may only + // transition out of WAITING while the lock is held. + let current_state = self.state.load(SeqCst); + + NotifyGuard { + guarded_notify: self, + guarded_waiters, + current_state, + } + } } impl Default for Notify { @@ -1374,3 +1390,20 @@ unsafe impl linked_list::Link for Waiter { } fn is_unpin() {} + +/// A guard that provides exclusive access to a `Notify`'s internal +/// waiters list. +/// +/// While this guard is held, the `Notify` instance's waiter list is locked. +pub(crate) struct NotifyGuard<'a> { + guarded_notify: &'a Notify, + guarded_waiters: crate::loom::sync::MutexGuard<'a, WaitList>, + current_state: usize, +} + +impl NotifyGuard<'_> { + pub(crate) fn notify_waiters(self) { + self.guarded_notify + .inner_notify_waiters(self.current_state, self.guarded_waiters); + } +} diff --git a/tokio/src/sync/set_once.rs b/tokio/src/sync/set_once.rs index e4adbf593..3170a9cda 100644 --- a/tokio/src/sync/set_once.rs +++ b/tokio/src/sync/set_once.rs @@ -1,14 +1,16 @@ use super::Notify; use crate::loom::cell::UnsafeCell; -use crate::loom::sync::{atomic::AtomicBool, Mutex}; +use crate::loom::sync::atomic::AtomicBool; use std::error::Error; use std::fmt; +use std::future::{poll_fn, Future}; use std::mem::MaybeUninit; use std::ops::Drop; use std::ptr; use std::sync::atomic::Ordering; +use std::task::Poll; // This file contains an implementation of an SetOnce. The value of SetOnce // can only be modified once during initialization. @@ -90,9 +92,6 @@ 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 { @@ -140,7 +139,6 @@ impl From for SetOnce { value_set: AtomicBool::new(true), value: UnsafeCell::new(MaybeUninit::new(value)), notify: Notify::new(), - lock: Mutex::new(()), } } } @@ -152,7 +150,6 @@ impl SetOnce { value_set: AtomicBool::new(false), value: UnsafeCell::new(MaybeUninit::uninit()), notify: Notify::new(), - lock: Mutex::new(()), } } @@ -195,7 +192,6 @@ impl SetOnce { value_set: AtomicBool::new(false), value: UnsafeCell::new(MaybeUninit::uninit()), notify: Notify::const_new(), - lock: Mutex::const_new(()), } } @@ -246,7 +242,6 @@ impl SetOnce { value_set: AtomicBool::new(true), value: UnsafeCell::new(MaybeUninit::new(value)), notify: Notify::const_new(), - lock: Mutex::const_new(()), } } @@ -287,19 +282,16 @@ impl SetOnce { return Err(SetOnceError(value)); } - // SAFETY: lock the mutex to ensure only one caller of set + // SAFETY: lock notify to ensure only one caller of set // can run at a time. - let guard = self.lock.lock(); + let guard = self.notify.lock_waiter_list(); 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 + // initialized or not, so we can safely write to the value unsafe { self.value.with_mut(|ptr| (*ptr).as_mut_ptr().write(value)); } @@ -308,10 +300,8 @@ impl SetOnce { // 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(); + guard.notify_waiters(); Ok(()) } @@ -353,20 +343,17 @@ impl SetOnce { } 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(); + pin!(notify_fut); + poll_fn(|cx| { + // Register under the notify's internal lock. + let ret = notify_fut.as_mut().poll(cx); if self.value_set.load(Ordering::Relaxed) { - // SAFETY: the state is initialized - return unsafe { self.get_unchecked() }; + return Poll::Ready(()); } - } - - // wait until the value is set - notify_fut.await; + ret + }) + .await; } } }