diff --git a/tokio/src/sync/notify.rs b/tokio/src/sync/notify.rs index 35d0804a13a..31a1df96653 100644 --- a/tokio/src/sync/notify.rs +++ b/tokio/src/sync/notify.rs @@ -855,7 +855,11 @@ impl Notified<'_> { } else { // Update the waker, if necessary. if let Some(waker) = waker { - if !w.waker.as_ref().unwrap().will_wake(waker) { + let should_update = match w.waker.as_ref() { + Some(current_waker) => current_waker.will_wake(waker), + None => true, + }; + if should_update { w.waker = Some(waker.clone()); } } diff --git a/tokio/tests/sync_notify.rs b/tokio/tests/sync_notify.rs index 5318d131cf8..160a568dc40 100644 --- a/tokio/tests/sync_notify.rs +++ b/tokio/tests/sync_notify.rs @@ -154,3 +154,56 @@ fn notify_one_after_dropped_all() { assert_ready!(notified2.poll()); } + +#[test] +fn test_notify_one_not_enabled() { + let notify = Notify::new(); + let mut future = spawn(notify.notified()); + + notify.notify_one(); + assert_ready!(future.poll()); +} + +#[test] +fn test_notify_one_after_enable() { + let notify = Notify::new(); + let mut future = spawn(notify.notified()); + + future.enter(|_, fut| assert!(!fut.enable())); + + notify.notify_one(); + assert_ready!(future.poll()); + future.enter(|_, fut| assert!(fut.enable())); +} + +#[test] +fn test_poll_after_enable() { + let notify = Notify::new(); + let mut future = spawn(notify.notified()); + + future.enter(|_, fut| assert!(!fut.enable())); + assert_pending!(future.poll()); +} + +#[test] +fn test_enable_after_poll() { + let notify = Notify::new(); + let mut future = spawn(notify.notified()); + + assert_pending!(future.poll()); + future.enter(|_, fut| assert!(!fut.enable())); +} + +#[test] +fn test_enable_consumes_permit() { + let notify = Notify::new(); + + // Add a permit. + notify.notify_one(); + + let mut future1 = spawn(notify.notified()); + future1.enter(|_, fut| assert!(fut.enable())); + + let mut future2 = spawn(notify.notified()); + future2.enter(|_, fut| assert!(!fut.enable())); +}