From 9f132172dbb00b1f6b6bda64bb9b194382079499 Mon Sep 17 00:00:00 2001 From: Hegui Dai Date: Tue, 31 Mar 2026 19:52:02 +0800 Subject: [PATCH] sync: fix `notify_waiters` priority in `Notify` (#7996) Previously, if a `notify_waiters()` was followed by a `notify_one()`, an unpolled `Notified` future created before the `notify_waiters()` call would consume the `NOTIFIED` permit created by the `notify_one()` call. This commit fixes this by verifying the `notify_waiters_calls` count before optimistically attempting to acquire the `NOTIFIED` permit. If the count indicates a `notify_waiters()` call has already happened, the future transitions directly to `State::Done` and leaves the permit intact for other waiters. Fixes: #7965 --- tokio/src/sync/notify.rs | 8 ++++++++ tokio/tests/sync_notify.rs | 12 ++++++++++++ 2 files changed, 20 insertions(+) diff --git a/tokio/src/sync/notify.rs b/tokio/src/sync/notify.rs index 33bce0e4f..7fb856c98 100644 --- a/tokio/src/sync/notify.rs +++ b/tokio/src/sync/notify.rs @@ -1118,6 +1118,14 @@ impl NotifiedProject<'_> { State::Init => { let curr = notify.state.load(SeqCst); + // Check if `notify_waiters` was called before attempting to acquire + // the `NOTIFIED` state. If a broadcast occurred, we will be woken by it, + // leaving the `notify_one` permit for other waiters. + if get_num_notify_waiters_calls(curr) != *notify_waiters_calls { + *state = State::Done; + continue 'outer_loop; + } + // Optimistically try acquiring a pending notification let res = notify.state.compare_exchange( set_state(curr, NOTIFIED), diff --git a/tokio/tests/sync_notify.rs b/tokio/tests/sync_notify.rs index 13b3f921e..ee7a9ecf0 100644 --- a/tokio/tests/sync_notify.rs +++ b/tokio/tests/sync_notify.rs @@ -301,3 +301,15 @@ fn test_waker_update() { assert!(future.is_woken()); } + +#[test] +fn unpolled_future_completed_by_notify_waiters_preserves_notify_one_permit() { + use futures::FutureExt; + let notify = Notify::new(); + let notified1 = notify.notified(); + notify.notify_waiters(); + notify.notify_one(); + assert!(notified1.now_or_never().is_some()); + let notified2 = notify.notified(); + assert!(notified2.now_or_never().is_some()); +}