Repository: tokio-rs/tokio (Rust)

Issue:
`sync/Notify::enable()` causes panic for unbound multi-producer multi-consumer (mpmc) channel
**Version**
tokio v1.19.0

**Platform**
Linux 62bead782d68 5.10.76-linuxkit #1 SMP PREEMPT Mon Nov 8 11:22:26 UTC 2021 aarch64 aarch64 aarch64 GNU/Linux

**Description**
There's an example for unbound multi-producer multi-consumer (mpmc) channel by using `Notify`: https://github.com/tokio-rs/tokio/blob/master/tokio/src/sync/notify.rs#L116

But the example code fails with a panic.

I tried this code as below, it's almost the same with the example.

```
    fn try_recv(&self) -> Option<AsyncRequestMessage> {
        self.requests.lock().unwrap().pop_front()
    }

    async fn recv(&self) -> Result<AsyncRequestMessage> {
        let future = self.notifier.notified();
        tokio::pin!(future);

        loop {
            // Make sure that no wakeup is lost if we get `None` from `try_recv`.
            future.as_mut().enable();

            if let Some(msg) = self.try_recv() {
                return Ok(msg);
            } else if self.closed.load(Ordering::Acquire) {
                return Err(Error::new(ErrorKind::BrokenPipe, "channel has been closed"));
            }

            // Wait for a call to `notify_one`.
            //
            // This uses `.as_mut()` to avoid consuming the future,
            // which lets us call `Pin::set` below.
            future.as_mut().await;

            // Reset the future in case another call to `try_recv` got the message before us.
            future.set(self.notifier.notified());
        }
    }
```

The above code cause a panic as below: 
thread 'nydus_storage_worker_0' panicked at 'called `Option::unwrap()` on a `None` value', /root/.cargo/registry/src/github.com-1ecc6299db9ec823/tokio-1.19.0/src/sync/notify.rs:858:50
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
thread 'nydus_storage_worker_1' panicked at 'called `Option::unwrap()` on a `None` value', /root/.cargo/registry/src/github.com-1ecc6299db9ec823/tokio-1.19.0/src/sync/notify.rs:858:50

It seems like a bug at https://github.com/tokio-rs/tokio/blob/master/tokio/src/sync/notify.rs#L858 
When `Notified::enable()` is called for the first time, waker is still None and the state is set to [`Waiting`](https://github.com/tokio-rs/tokio/blob/master/tokio/src/sync/notify.rs#L833)
Then `Notifiled::poll()` is called by `.await`, it triggers the panic because `waiter` is still `None`.

So we should check `waiter` is not `None` before calling `as_ref().unwrap()`.
```
                       if let Some(waker) = waker {
                            if w.waker.is_none() || !w.waker.as_ref().unwrap().will_wake(waker) {
                                w.waker = Some(waker.clone());
                            }
                        }
```





Discussion:
Thank you for the bug report. I've opened a PR to fix this.



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


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