Drop wakers after unlocking the mutex in Notify

tokio/src/sync/notify.rs

@@ -917,10 +917,17 @@ impl Notified<'_> {
                         }
                     }
 
+                    let mut old_waker = None;
                     if waker.is_some() {
                         // Safety: called while locked.
+                        //
+                        // This code makes the following assignment, but it
+                        // stores the old waker in `old_waker` so that we can
+                        // drop it after releasing the mutex.
+                        //
+                        // (*waiter.get()).waker = waker;
                         unsafe {
-                            (*waiter.get()).waker = waker;
+                            old_waker = std::mem::replace(&mut (*waiter.get()).waker, waker);
                         }
                     }
 
@@ -931,6 +938,9 @@ impl Notified<'_> {
 
                     *state = Waiting;
 
+                    drop(waiters);
+                    drop(old_waker);
+
                     return Poll::Pending;
                 }
                 Waiting => {
@@ -945,12 +955,13 @@ impl Notified<'_> {
 
                     // Safety: called while locked
                     let w = unsafe { &mut *waiter.get() };
+                    let mut old_waker = None;
 
                     if w.notified.is_some() {
                         // Our waker has been notified and our waiter is already removed from
                         // the list. Reset the notification and convert to `Done`.
+                        old_waker = std::mem::take(&mut w.waker);
                         w.notified = None;
-                        w.waker = None;
                         *state = Done;
                     } else if get_num_notify_waiters_calls(curr) != *notify_waiters_calls {
                         // Before we add a waiter to the list we check if these numbers are
@@ -960,7 +971,7 @@ impl Notified<'_> {
                         // We can treat the waiter as notified and remove it from the list, as
                         // it would have been notified in the `notify_waiters` call anyways.
 
-                        w.waker = None;
+                        old_waker = std::mem::take(&mut w.waker);
 
                         // Safety: we hold the lock, so we have an exclusive access to the list.
                         // The list is used in `notify_waiters`, so it must be guarded.
@@ -975,10 +986,14 @@ impl Notified<'_> {
                                 None => true,
                             };
                             if should_update {
-                                w.waker = Some(waker.clone());
+                                old_waker = std::mem::replace(&mut w.waker, Some(waker.clone()));
                             }
                         }
 
+                        // Drop the old waker after releasing the lock.
+                        drop(waiters);
+                        drop(old_waker);
+
                         return Poll::Pending;
                     }
 
@@ -988,6 +1003,9 @@ impl Notified<'_> {
                     // is helpful to visualize the scope of the critical
                     // section.
                     drop(waiters);
+
+                    // Drop the old waker after releasing the lock.
+                    drop(old_waker);
                 }
                 Done => {
                     return Poll::Ready(());