sync: fix panic in Notified::enable (#4747)

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());
                             }
                         }

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()));
+}