Skip to content

Commit

Permalink
sync: drop wakers after unlocking in Notify
Browse files Browse the repository at this point in the history
  • Loading branch information
Darksonn committed Feb 19, 2023
1 parent 2e0372b commit 6ccb60e
Showing 1 changed file with 22 additions and 4 deletions.
26 changes: 22 additions & 4 deletions tokio/src/sync/notify.rs
Original file line number Diff line number Diff line change
Expand Up @@ -917,10 +917,17 @@ impl Notified<'_> {
}
}

let mut old_waker = None;
if waker.is_some() {
// Safety: called while locked.
//
// This code makes the following assignment, but it
// stores the old waker in `old_waker` so that we can
// drop it after releasing the mutex.
//
// (*waiter.get()).waker = waker;
unsafe {
(*waiter.get()).waker = waker;
old_waker = std::mem::replace(&mut (*waiter.get()).waker, waker);
}
}

Expand All @@ -931,6 +938,9 @@ impl Notified<'_> {

*state = Waiting;

drop(waiters);
drop(old_waker);

return Poll::Pending;
}
Waiting => {
Expand All @@ -945,12 +955,13 @@ impl Notified<'_> {

// Safety: called while locked
let w = unsafe { &mut *waiter.get() };
let mut old_waker = None;

if w.notified.is_some() {
// Our waker has been notified and our waiter is already removed from
// the list. Reset the notification and convert to `Done`.
old_waker = std::mem::take(&mut w.waker);
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
Expand All @@ -960,7 +971,7 @@ impl Notified<'_> {
// 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;
old_waker = std::mem::take(&mut w.waker);

// 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.
Expand All @@ -975,10 +986,14 @@ impl Notified<'_> {
None => true,
};
if should_update {
w.waker = Some(waker.clone());
old_waker = std::mem::replace(&mut w.waker, Some(waker.clone()));
}
}

// Drop the old waker after releasing the lock.
drop(waiters);
drop(old_waker);

return Poll::Pending;
}

Expand All @@ -988,6 +1003,9 @@ impl Notified<'_> {
// is helpful to visualize the scope of the critical
// section.
drop(waiters);

// Drop the old waker after releasing the lock.
drop(old_waker);
}
Done => {
return Poll::Ready(());
Expand Down

0 comments on commit 6ccb60e

Please sign in to comment.