sync: preserve permit state in notify_waiters (#3660)

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
+            // There are no waiting tasks. All we need to do is increment the
-            // established between `notify` and `notified().await`.
+            // number of times this method was called.
-            // All we need to do is increment the number of times this
+            atomic_inc_num_notify_waiters_calls(&self.state);
-            // method was called.
-            self.state.store(inc_num_notify_waiters_calls(curr), SeqCst);
             return;
         }
 

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