From 6845a93cbf52ca8b0709918db0c8c558d6d8720e Mon Sep 17 00:00:00 2001
From: Simon Lindholm <simon.lindholm10@gmail.com>
Date: Wed, 5 May 2021 15:20:24 +0200
Subject: [PATCH] sync: preserve permit state in notify_waiters (#3660)

---
 tokio/src/sync/notify.rs            | 12 ++++++-----
 tokio/src/sync/tests/loom_notify.rs | 33 +++++++++++++++++++++++++++--
 2 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/tokio/src/sync/notify.rs b/tokio/src/sync/notify.rs
index 2d30da9c6..5d2132f79 100644
--- a/tokio/src/sync/notify.rs
+++ b/tokio/src/sync/notify.rs
@@ -192,6 +192,10 @@ fn inc_num_notify_waiters_calls(data: usize) -> usize {
     data + (1 << NOTIFY_WAITERS_SHIFT)
 }
 
+fn atomic_inc_num_notify_waiters_calls(data: &AtomicUsize) {
+    data.fetch_add(1 << NOTIFY_WAITERS_SHIFT, SeqCst);
+}
+
 impl Notify {
     /// Create a new `Notify`, initialized without a permit.
     ///
@@ -394,11 +398,9 @@ impl Notify {
         let curr = self.state.load(SeqCst);
 
         if let EMPTY | NOTIFIED = get_state(curr) {
-            // There are no waiting tasks. In this case, no synchronization is
-            // established between `notify` and `notified().await`.
-            // All we need to do is increment the number of times this
-            // method was called.
-            self.state.store(inc_num_notify_waiters_calls(curr), SeqCst);
+            // There are no waiting tasks. All we need to do is increment the
+            // number of times this method was called.
+            atomic_inc_num_notify_waiters_calls(&self.state);
             return;
         }
 
diff --git a/tokio/src/sync/tests/loom_notify.rs b/tokio/src/sync/tests/loom_notify.rs
index 4be949a37..d484a7581 100644
--- a/tokio/src/sync/tests/loom_notify.rs
+++ b/tokio/src/sync/tests/loom_notify.rs
@@ -33,12 +33,41 @@ fn notify_waiters() {
             tx.notify_waiters();
         });
 
-        th.join().unwrap();
-
         block_on(async {
             notified1.await;
             notified2.await;
         });
+
+        th.join().unwrap();
+    });
+}
+
+#[test]
+fn notify_waiters_and_one() {
+    loom::model(|| {
+        let notify = Arc::new(Notify::new());
+        let tx1 = notify.clone();
+        let tx2 = notify.clone();
+
+        let th1 = thread::spawn(move || {
+            tx1.notify_waiters();
+        });
+
+        let th2 = thread::spawn(move || {
+            tx2.notify_one();
+        });
+
+        let th3 = thread::spawn(move || {
+            let notified = notify.notified();
+
+            block_on(async {
+                notified.await;
+            });
+        });
+
+        th1.join().unwrap();
+        th2.join().unwrap();
+        th3.join().unwrap();
     });
 }