From 6400934bebb4c04899c64a45c3403c05da49ead0 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Sat, 4 Jun 2022 22:33:37 +0200 Subject: [PATCH 1/3] sync: fix panic in `Notified::enable` --- tokio/src/sync/notify.rs | 6 +++++- tokio/tests/sync_notify.rs | 20 ++++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) 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..b77a653b50f 100644 --- a/tokio/tests/sync_notify.rs +++ b/tokio/tests/sync_notify.rs @@ -154,3 +154,23 @@ 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| fut.enable()); + + notify.notify_one(); + assert_ready!(future.poll()); +} From bcf5160beb99aa8a86eea3ee084ce5265d19775f Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Sat, 4 Jun 2022 22:51:39 +0200 Subject: [PATCH 2/3] Add more tests --- tokio/tests/sync_notify.rs | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/tokio/tests/sync_notify.rs b/tokio/tests/sync_notify.rs index b77a653b50f..71b3ea6fb3e 100644 --- a/tokio/tests/sync_notify.rs +++ b/tokio/tests/sync_notify.rs @@ -169,8 +169,27 @@ fn test_notify_one_after_enable() { let notify = Notify::new(); let mut future = spawn(notify.notified()); - future.enter(|_, fut| fut.enable()); + 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())); } From 98b28862c78d4bbcd996a0ad423310a7709b9973 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Sat, 4 Jun 2022 22:56:15 +0200 Subject: [PATCH 3/3] Add another test --- tokio/tests/sync_notify.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tokio/tests/sync_notify.rs b/tokio/tests/sync_notify.rs index 71b3ea6fb3e..160a568dc40 100644 --- a/tokio/tests/sync_notify.rs +++ b/tokio/tests/sync_notify.rs @@ -193,3 +193,17 @@ fn test_enable_after_poll() { 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())); +}