tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-24 00:00:11 +02:00
Code Issues Packages Projects Releases Wiki Activity

sync: make notify_waiters calls atomic (#5458)

Browse Source
This commit is contained in:
Tymoteusz Wiśniewski
2023-02-19 14:10:38 +01:00
committed by GitHub
parent 0f17d69303
commit 795754a846
4 changed files with 474 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/sync/notify.rs
+148 -42
View File
@@ -7,7 +7,7 @@
use crate::loom::sync::atomic::AtomicUsize;
use crate::loom::sync::Mutex;
use crate::util::linked_list::{self, LinkedList};
use crate::util::linked_list::{self, GuardedLinkedList, LinkedList};
use crate::util::WakeList;
use std::cell::UnsafeCell;
@@ -20,6 +20,7 @@ use std::sync::atomic::Ordering::SeqCst;
use std::task::{Context, Poll, Waker};
type WaitList = LinkedList<Waiter, <Waiter as linked_list::Link>::Target>;
type GuardedWaitList = GuardedLinkedList<Waiter, <Waiter as linked_list::Link>::Target>;
/// Notifies a single task to wake up.
///
@@ -198,10 +199,16 @@ type WaitList = LinkedList<Waiter, <Waiter as linked_list::Link>::Target>;
/// [`Semaphore`]: crate::sync::Semaphore
#[derive(Debug)]
pub struct Notify {
// This uses 2 bits to store one of `EMPTY`,
// `state` uses 2 bits to store one of `EMPTY`,
// `WAITING` or `NOTIFIED`. The rest of the bits
// are used to store the number of times `notify_waiters`
// was called.
//
// Throughout the code there are two assumptions:
// - state can be transitioned *from* `WAITING` only if
// `waiters` lock is held
// - number of times `notify_waiters` was called can
// be modified only if `waiters` lock is held
state: AtomicUsize,
waiters: Mutex<WaitList>,
}
@@ -229,6 +236,17 @@ struct Waiter {
_p: PhantomPinned,
}
impl Waiter {
fn new() -> Waiter {
Waiter {
pointers: linked_list::Pointers::new(),
waker: None,
notified: None,
_p: PhantomPinned,
}
}
}
generate_addr_of_methods! {
impl<> Waiter {
unsafe fn addr_of_pointers(self: NonNull<Self>) -> NonNull<linked_list::Pointers<Waiter>> {
@@ -237,6 +255,59 @@ generate_addr_of_methods! {
}
}
/// List used in `Notify::notify_waiters`. It wraps a guarded linked list
/// and gates the access to it on `notify.waiters` mutex. It also empties
/// the list on drop.
struct NotifyWaitersList<'a> {
list: GuardedWaitList,
is_empty: bool,
notify: &'a Notify,
}
impl<'a> NotifyWaitersList<'a> {
fn new(
unguarded_list: WaitList,
guard: Pin<&'a mut UnsafeCell<Waiter>>,
notify: &'a Notify,
) -> NotifyWaitersList<'a> {
// Safety: pointer to the guarding waiter is not null.
let guard_ptr = unsafe { NonNull::new_unchecked(guard.get()) };
let list = unguarded_list.into_guarded(guard_ptr);
NotifyWaitersList {
list,
is_empty: false,
notify,
}
}
/// Removes the last element from the guarded list. Modifying this list
/// requires an exclusive access to the main list in `Notify`.
fn pop_back_locked(&mut self, _waiters: &mut WaitList) -> Option<NonNull<Waiter>> {
let result = self.list.pop_back();
if result.is_none() {
// Save information about emptiness to avoid waiting for lock
// in the destructor.
self.is_empty = true;
}
result
}
}
impl Drop for NotifyWaitersList<'_> {
fn drop(&mut self) {
// If the list is not empty, we unlink all waiters from it.
// We do not wake the waiters to avoid double panics.
if !self.is_empty {
let _lock_guard = self.notify.waiters.lock();
while let Some(mut waiter) = self.list.pop_back() {
// Safety: we hold the lock.
let waiter = unsafe { waiter.as_mut() };
waiter.notified = Some(NotificationType::AllWaiters);
}
}
}
}
/// Future returned from [`Notify::notified()`].
///
/// This future is fused, so once it has completed, any future calls to poll
@@ -249,6 +320,9 @@ pub struct Notified<'a> {
/// The current state of the receiving process.
state: State,
/// Number of calls to `notify_waiters` at the time of creation.
notify_waiters_calls: usize,
/// Entry in the waiter `LinkedList`.
waiter: UnsafeCell<Waiter>,
}
@@ -258,7 +332,7 @@ unsafe impl<'a> Sync for Notified<'a> {}
#[derive(Debug)]
enum State {
Init(usize),
Init,
Waiting,
Done,
}
@@ -383,17 +457,13 @@ impl Notify {
/// ```
pub fn notified(&self) -> Notified<'_> {
// we load the number of times notify_waiters
// was called and store that in our initial state
// was called and store that in the future.
let state = self.state.load(SeqCst);
Notified {
notify: self,
state: State::Init(state >> NOTIFY_WAITERS_SHIFT),
waiter: UnsafeCell::new(Waiter {
pointers: linked_list::Pointers::new(),
waker: None,
notified: None,
_p: PhantomPinned,
}),
state: State::Init,
notify_waiters_calls: get_num_notify_waiters_calls(state),
waiter: UnsafeCell::new(Waiter::new()),
}
}
@@ -500,12 +570,9 @@ impl Notify {
/// }
/// ```
pub fn notify_waiters(&self) {
let mut wakers = WakeList::new();
// There are waiters, the lock must be acquired to notify.
let mut waiters = self.waiters.lock();
// The state must be reloaded while the lock is held. The state may only
// 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);
@@ -516,12 +583,30 @@ impl Notify {
return;
}
// At this point, it is guaranteed that the state will not
// concurrently change, as holding the lock is required to
// transition **out** of `WAITING`.
// Increment the number of times this method was called
// and transition to empty.
let new_state = set_state(inc_num_notify_waiters_calls(curr), EMPTY);
self.state.store(new_state, SeqCst);
// It is critical for `GuardedLinkedList` safety that the guard node is
// pinned in memory and is not dropped until the guarded list is dropped.
let guard = UnsafeCell::new(Waiter::new());
pin!(guard);
// We move all waiters to a secondary list. It uses a `GuardedLinkedList`
// underneath to allow every waiter to safely remove itself from it.
//
// * This list will be still guarded by the `waiters` lock.
// `NotifyWaitersList` wrapper makes sure we hold the lock to modify it.
// * This wrapper will empty the list on drop. It is critical for safety
// that we will not leave any list entry with a pointer to the local
// guard node after this function returns / panics.
let mut list = NotifyWaitersList::new(std::mem::take(&mut *waiters), guard, self);
let mut wakers = WakeList::new();
'outer: loop {
while wakers.can_push() {
match waiters.pop_back() {
match list.pop_back_locked(&mut waiters) {
Some(mut waiter) => {
// Safety: `waiters` lock is still held.
let waiter = unsafe { waiter.as_mut() };
@@ -540,20 +625,17 @@ impl Notify {
}
}
// Release the lock before notifying.
drop(waiters);
// One of the wakers may panic, but the remaining waiters will still
// be unlinked from the list in `NotifyWaitersList` destructor.
wakers.wake_all();
// Acquire the lock again.
waiters = self.waiters.lock();
}
// All waiters will be notified, the state must be transitioned to
// `EMPTY`. As transitioning **from** `WAITING` requires the lock to be
// held, a `store` is sufficient.
let new = set_state(inc_num_notify_waiters_calls(curr), EMPTY);
self.state.store(new, SeqCst);
// Release the lock before notifying
drop(waiters);
@@ -730,26 +812,32 @@ impl Notified<'_> {
/// A custom `project` implementation is used in place of `pin-project-lite`
/// as a custom drop implementation is needed.
fn project(self: Pin<&mut Self>) -> (&Notify, &mut State, &UnsafeCell<Waiter>) {
fn project(self: Pin<&mut Self>) -> (&Notify, &mut State, &usize, &UnsafeCell<Waiter>) {
unsafe {
// Safety: both `notify` and `state` are `Unpin`.
// Safety: `notify`, `state` and `notify_waiters_calls` are `Unpin`.
is_unpin::<&Notify>();
is_unpin::<AtomicUsize>();
is_unpin::<usize>();
let me = self.get_unchecked_mut();
(me.notify, &mut me.state, &me.waiter)
(
me.notify,
&mut me.state,
&me.notify_waiters_calls,
&me.waiter,
)
}
}
fn poll_notified(self: Pin<&mut Self>, waker: Option<&Waker>) -> Poll<()> {
use State::*;
let (notify, state, waiter) = self.project();
let (notify, state, notify_waiters_calls, waiter) = self.project();
loop {
match *state {
Init(initial_notify_waiters_calls) => {
Init => {
let curr = notify.state.load(SeqCst);
// Optimistically try acquiring a pending notification
@@ -779,7 +867,7 @@ impl Notified<'_> {
// if notify_waiters has been called after the future
// was created, then we are done
if get_num_notify_waiters_calls(curr) != initial_notify_waiters_calls {
if get_num_notify_waiters_calls(curr) != *notify_waiters_calls {
*state = Done;
return Poll::Ready(());
}
@@ -846,21 +934,37 @@ impl Notified<'_> {
return Poll::Pending;
}
Waiting => {
// Currently in the "Waiting" state, implying the caller has
// a waiter stored in the waiter list (guarded by
// `notify.waiters`). In order to access the waker fields,
// we must hold the lock.
// Currently in the "Waiting" state, implying the caller has a waiter stored in
// a waiter list (guarded by `notify.waiters`). In order to access the waker
// fields, we must acquire the lock.
let waiters = notify.waiters.lock();
let mut waiters = notify.waiters.lock();
// Load the state with the lock held.
let curr = notify.state.load(SeqCst);
// Safety: called while locked
let w = unsafe { &mut *waiter.get() };
if w.notified.is_some() {
// Our waker has been notified. Reset the fields and
// remove it from the list.
w.waker = None;
// Our waker has been notified and our waiter is already removed from
// the list. Reset the notification and convert to `Done`.
w.notified = None;
w.waker = None;
*state = Done;
} else if get_num_notify_waiters_calls(curr) != *notify_waiters_calls {
// Before we add a waiter to the list we check if these numbers are
// different while holding the lock. If these numbers are different now,
// it means that there is a call to `notify_waiters` in progress and this
// waiter must be contained by a guarded list used in `notify_waiters`.
// We can treat the waiter as notified and remove it from the list, as
// it would have been notified in the `notify_waiters` call anyways.
w.waker = None;
// Safety: we hold the lock, so we have an exclusive access to the list.
// The list is used in `notify_waiters`, so it must be guarded.
unsafe { waiters.remove(NonNull::new_unchecked(w)) };
*state = Done;
} else {
@@ -906,7 +1010,7 @@ impl Drop for Notified<'_> {
use State::*;
// Safety: The type only transitions to a "Waiting" state when pinned.
let (notify, state, waiter) = unsafe { Pin::new_unchecked(self).project() };
let (notify, state, _, waiter) = unsafe { Pin::new_unchecked(self).project() };
// This is where we ensure safety. The `Notified` value is being
// dropped, which means we must ensure that the waiter entry is no
@@ -917,8 +1021,10 @@ impl Drop for Notified<'_> {
// remove the entry from the list (if not already removed)
//
// safety: the waiter is only added to `waiters` by virtue of it
// being the only `LinkedList` available to the type.
// Safety: we hold the lock, so we have an exclusive access to every list the
// waiter may be contained in. If the node is not contained in the `waiters`
// list, then it is contained by a guarded list used by `notify_waiters` and
// in such case it must be a middle node.
unsafe { waiters.remove(NonNull::new_unchecked(waiter.get())) };
if waiters.is_empty() && get_state(notify_state) == WAITING {
tokio/src/sync/tests/loom_notify.rs
+191
View File
@@ -4,6 +4,11 @@ use loom::future::block_on;
use loom::sync::Arc;
use loom::thread;
use tokio_test::{assert_pending, assert_ready};
/// `util::wake_list::NUM_WAKERS`
const WAKE_LIST_SIZE: usize = 32;
#[test]
fn notify_one() {
loom::model(|| {
@@ -138,3 +143,189 @@ fn notify_drop() {
th2.join().unwrap();
});
}
/// Polls two `Notified` futures and checks if poll results are consistent
/// with each other. If the first future is notified by a `notify_waiters`
/// call, then the second one must be notified as well.
#[test]
fn notify_waiters_poll_consistency() {
fn notify_waiters_poll_consistency_variant(poll_setting: [bool; 2]) {
let notify = Arc::new(Notify::new());
let mut notified = [
tokio_test::task::spawn(notify.notified()),
tokio_test::task::spawn(notify.notified()),
];
for i in 0..2 {
if poll_setting[i] {
assert_pending!(notified[i].poll());
}
}
let tx = notify.clone();
let th = thread::spawn(move || {
tx.notify_waiters();
});
let res1 = notified[0].poll();
let res2 = notified[1].poll();
// If res1 is ready, then res2 must also be ready.
assert!(res1.is_pending() || res2.is_ready());
th.join().unwrap();
}
// We test different scenarios in which pending futures had or had not
// been polled before the call to `notify_waiters`.
loom::model(|| notify_waiters_poll_consistency_variant([false, false]));
loom::model(|| notify_waiters_poll_consistency_variant([true, false]));
loom::model(|| notify_waiters_poll_consistency_variant([false, true]));
loom::model(|| notify_waiters_poll_consistency_variant([true, true]));
}
/// Polls two `Notified` futures and checks if poll results are consistent
/// with each other. If the first future is notified by a `notify_waiters`
/// call, then the second one must be notified as well.
///
/// Here we also add other `Notified` futures in between to force the two
/// tested futures to end up in different chunks.
#[test]
fn notify_waiters_poll_consistency_many() {
fn notify_waiters_poll_consistency_many_variant(order: [usize; 2]) {
let notify = Arc::new(Notify::new());
let mut futs = (0..WAKE_LIST_SIZE + 1)
.map(|_| tokio_test::task::spawn(notify.notified()))
.collect::<Vec<_>>();
assert_pending!(futs[order[0]].poll());
for i in 2..futs.len() {
assert_pending!(futs[i].poll());
}
assert_pending!(futs[order[1]].poll());
let tx = notify.clone();
let th = thread::spawn(move || {
tx.notify_waiters();
});
let res1 = futs[0].poll();
let res2 = futs[1].poll();
// If res1 is ready, then res2 must also be ready.
assert!(res1.is_pending() || res2.is_ready());
th.join().unwrap();
}
// We test different scenarios in which futures are polled in different order.
loom::model(|| notify_waiters_poll_consistency_many_variant([0, 1]));
loom::model(|| notify_waiters_poll_consistency_many_variant([1, 0]));
}
/// Checks if a call to `notify_waiters` is observed as atomic when combined
/// with a concurrent call to `notify_one`.
#[test]
fn notify_waiters_is_atomic() {
fn notify_waiters_is_atomic_variant(tested_fut_index: usize) {
let notify = Arc::new(Notify::new());
let mut futs = (0..WAKE_LIST_SIZE + 1)
.map(|_| tokio_test::task::spawn(notify.notified()))
.collect::<Vec<_>>();
for fut in &mut futs {
assert_pending!(fut.poll());
}
let tx = notify.clone();
let th = thread::spawn(move || {
tx.notify_waiters();
});
block_on(async {
// If awaiting one of the futures completes, then we should be
// able to assume that all pending futures are notified. Therefore
// a notification from a subsequent `notify_one` call should not
// be consumed by an old future.
futs.remove(tested_fut_index).await;
let mut new_fut = tokio_test::task::spawn(notify.notified());
assert_pending!(new_fut.poll());
notify.notify_one();
// `new_fut` must consume the notification from `notify_one`.
assert_ready!(new_fut.poll());
});
th.join().unwrap();
}
// We test different scenarios in which the tested future is at the beginning
// or at the end of the waiters queue used by `Notify`.
loom::model(|| notify_waiters_is_atomic_variant(0));
loom::model(|| notify_waiters_is_atomic_variant(32));
}
/// Checks if a single call to `notify_waiters` does not get through two `Notified`
/// futures created and awaited sequentially like this:
/// ```ignore
/// notify.notified().await;
/// notify.notified().await;
/// ```
#[test]
fn notify_waiters_sequential_notified_await() {
use crate::sync::oneshot;
loom::model(|| {
let notify = Arc::new(Notify::new());
let (tx_fst, rx_fst) = oneshot::channel();
let (tx_snd, rx_snd) = oneshot::channel();
let receiver = thread::spawn({
let notify = notify.clone();
move || {
block_on(async {
// Poll the first `Notified` to put it as the first waiter
// in the queue.
let mut first_notified = tokio_test::task::spawn(notify.notified());
assert_pending!(first_notified.poll());
// Create additional waiters to force `notify_waiters` to
// release the lock at least once.
let _task_pile = (0..WAKE_LIST_SIZE + 1)
.map(|_| {
let mut fut = tokio_test::task::spawn(notify.notified());
assert_pending!(fut.poll());
fut
})
.collect::<Vec<_>>();
// We are ready for the notify_waiters call.
tx_fst.send(()).unwrap();
first_notified.await;
// Poll the second `Notified` future to try to insert
// it to the waiters queue.
let mut second_notified = tokio_test::task::spawn(notify.notified());
assert_pending!(second_notified.poll());
// Wait for the `notify_waiters` to end and check if we
// are woken up.
rx_snd.await.unwrap();
assert_pending!(second_notified.poll());
});
}
});
// Wait for the signal and call `notify_waiters`.
block_on(rx_fst).unwrap();
notify.notify_waiters();
tx_snd.send(()).unwrap();
receiver.join().unwrap();
});
}
tokio/src/sync/tests/notify.rs
+39
View File
@@ -46,6 +46,45 @@ fn notify_clones_waker_before_lock() {
let _ = future.poll(&mut cx);
}
#[cfg(panic = "unwind")]
#[test]
fn notify_waiters_handles_panicking_waker() {
use futures::task::ArcWake;
let notify = Arc::new(Notify::new());
struct PanickingWaker(Arc<Notify>);
impl ArcWake for PanickingWaker {
fn wake_by_ref(_arc_self: &Arc<Self>) {
panic!("waker panicked");
}
}
let bad_fut = notify.notified();
pin!(bad_fut);
let waker = futures::task::waker(Arc::new(PanickingWaker(notify.clone())));
let mut cx = Context::from_waker(&waker);
let _ = bad_fut.poll(&mut cx);
let mut futs = Vec::new();
for _ in 0..32 {
let mut fut = tokio_test::task::spawn(notify.notified());
assert!(fut.poll().is_pending());
futs.push(fut);
}
assert!(std::panic::catch_unwind(|| {
notify.notify_waiters();
})
.is_err());
for mut fut in futs {
assert!(fut.poll().is_ready());
}
}
#[test]
fn notify_simple() {
let notify = Notify::new();
tokio/src/util/linked_list.rs
+96 -2
View File
@@ -178,8 +178,12 @@ impl<L: Link> LinkedList<L, L::Target> {
///
/// # Safety
///
/// The caller **must** ensure that `node` is currently contained by
/// `self` or not contained by any other list.
/// The caller **must** ensure that exactly one of the following is true:
/// - `node` is currently contained by `self`,
/// - `node` is not contained by any list,
/// - `node` is currently contained by some other `GuardedLinkedList` **and**
/// the caller has an exclusive access to that list. This condition is
/// used by the linked list in `sync::Notify`.
pub(crate) unsafe fn remove(&mut self, node: NonNull<L::Target>) -> Option<L::Handle> {
if let Some(prev) = L::pointers(node).as_ref().get_prev() {
debug_assert_eq!(L::pointers(prev).as_ref().get_next(), Some(node));
@@ -290,6 +294,96 @@ cfg_io_readiness! {
}
}
// ===== impl GuardedLinkedList =====
feature! {
#![any(
feature = "process",
feature = "sync",
feature = "rt",
feature = "signal",
)]
/// An intrusive linked list, but instead of keeping pointers to the head
/// and tail nodes, it uses a special guard node linked with those nodes.
/// It means that the list is circular and every pointer of a node from
/// the list is not `None`, including pointers from the guard node.
///
/// If a list is empty, then both pointers of the guard node are pointing
/// at the guard node itself.
pub(crate) struct GuardedLinkedList<L, T> {
/// Pointer to the guard node.
guard: NonNull<T>,
/// Node type marker.
_marker: PhantomData<*const L>,
}
impl<U, L: Link<Handle = NonNull<U>>> LinkedList<L, L::Target> {
/// Turns a linked list into the guarded version by linking the guard node
/// with the head and tail nodes. Like with other nodes, you should guarantee
/// that the guard node is pinned in memory.
pub(crate) fn into_guarded(self, guard_handle: L::Handle) -> GuardedLinkedList<L, L::Target> {
// `guard_handle` is a NonNull pointer, we don't have to care about dropping it.
let guard = L::as_raw(&guard_handle);
unsafe {
if let Some(head) = self.head {
debug_assert!(L::pointers(head).as_ref().get_prev().is_none());
L::pointers(head).as_mut().set_prev(Some(guard));
L::pointers(guard).as_mut().set_next(Some(head));
// The list is not empty, so the tail cannot be `None`.
let tail = self.tail.unwrap();
debug_assert!(L::pointers(tail).as_ref().get_next().is_none());
L::pointers(tail).as_mut().set_next(Some(guard));
L::pointers(guard).as_mut().set_prev(Some(tail));
} else {
// The list is empty.
L::pointers(guard).as_mut().set_prev(Some(guard));
L::pointers(guard).as_mut().set_next(Some(guard));
}
}
GuardedLinkedList { guard, _marker: PhantomData }
}
}
impl<L: Link> GuardedLinkedList<L, L::Target> {
fn tail(&self) -> Option<NonNull<L::Target>> {
let tail_ptr = unsafe {
L::pointers(self.guard).as_ref().get_prev().unwrap()
};
// Compare the tail pointer with the address of the guard node itself.
// If the guard points at itself, then there are no other nodes and
// the list is considered empty.
if tail_ptr != self.guard {
Some(tail_ptr)
} else {
None
}
}
/// Removes the last element from a list and returns it, or None if it is
/// empty.
pub(crate) fn pop_back(&mut self) -> Option<L::Handle> {
unsafe {
let last = self.tail()?;
let before_last = L::pointers(last).as_ref().get_prev().unwrap();
L::pointers(self.guard).as_mut().set_prev(Some(before_last));
L::pointers(before_last).as_mut().set_next(Some(self.guard));
L::pointers(last).as_mut().set_prev(None);
L::pointers(last).as_mut().set_next(None);
Some(L::from_raw(last))
}
}
}
}
// ===== impl Pointers =====
impl<T> Pointers<T> {