sync: tweak watch API

tokio/src/sync/barrier.rs

@@ -110,9 +110,11 @@ impl Barrier {
         let mut wait = self.wait.clone();
 
         loop {
+            let _ = wait.changed().await;
+
             // note that the first time through the loop, this _will_ yield a generation
             // immediately, since we cloned a receiver that has never seen any values.
-            if wait.recv().await >= generation {
+            if *wait.borrow() >= generation {
                 break;
             }
         }

tokio/src/sync/mod.rs

@@ -355,10 +355,8 @@
 //!             let op = my_async_operation();
 //!             tokio::pin!(op);
 //!
-//!             // Receive the **initial** configuration value. As this is the
-//!             // first time the config is received from the watch, it will
-//!             // always complete immediatedly.
-//!             let mut conf = rx.recv().await;
+//!             // Get the initial config value
+//!             let mut conf = rx.borrow().clone();
 //!
 //!             let mut op_start = Instant::now();
 //!             let mut delay = time::delay_until(op_start + conf.timeout);
@@ -375,8 +373,8 @@
 //!                         // Restart the timeout
 //!                         delay = time::delay_until(op_start + conf.timeout);
 //!                     }
-//!                     new_conf = rx.recv() => {
-//!                         conf = new_conf;
+//!                     _ = rx.changed() => {
+//!                         conf = rx.borrow().clone();
 //!
 //!                         // The configuration has been updated. Update the
 //!                         // `delay` using the new `timeout` value.

tokio/src/sync/notify.rs

@@ -121,7 +121,7 @@ struct Waiter {
 
 /// Future returned from `notified()`
 #[derive(Debug)]
-struct Notified<'a> {
+pub struct Notified<'a> {
     /// The `Notify` being received on.
     notify: &'a Notify,
 
@@ -170,6 +170,12 @@ impl Notify {
 
     /// Wait for a notification.
     ///
+    /// Equivalent to:
+    ///
+    /// ```ignore
+    /// async fn notified(&self);
+    /// ```
+    ///
     /// Each `Notify` value holds a single permit. If a permit is available from
     /// an earlier call to [`notify()`], then `notified().await` will complete
     /// immediately, consuming that permit. Otherwise, `notified().await` waits
@@ -197,7 +203,7 @@ impl Notify {
     ///     notify.notify();
     /// }
     /// ```
-    pub async fn notified(&self) {
+    pub fn notified(&self) -> Notified<'_> {
         Notified {
             notify: self,
             state: State::Init,
@@ -208,7 +214,6 @@ impl Notify {
                 _p: PhantomPinned,
             }),
         }
-        .await
     }
 
     /// Notifies a waiting task
@@ -277,6 +282,45 @@ impl Notify {
             waker.wake();
         }
     }
+
+    /// Notifies all waiting tasks
+    pub(crate) fn notify_waiters(&self) {
+        // There are waiters, the lock must be acquired to notify.
+        let mut waiters = self.waiters.lock().unwrap();
+
+        // The state must be reloaded while the lock is held. The state may only
+        // transition out of WAITING while the lock is held.
+        let curr = self.state.load(SeqCst);
+
+        if let EMPTY | NOTIFIED = curr {
+            // There are no waiting tasks. In this case, no synchronization is
+            // established between `notify` and `notified().await`.
+            return;
+        }
+
+        // At this point, it is guaranteed that the state will not
+        // concurrently change, as holding the lock is required to
+        // transition **out** of `WAITING`.
+        //
+        // Get pending waiters
+        while let Some(mut waiter) = waiters.pop_back() {
+            // Safety: `waiters` lock is still held.
+            let waiter = unsafe { waiter.as_mut() };
+
+            assert!(!waiter.notified);
+
+            waiter.notified = true;
+
+            if let Some(waker) = waiter.waker.take() {
+                waker.wake();
+            }
+        }
+
+        // All waiters have been notified, the state must be transitioned to
+        // `EMPTY`. As transitioning **from** `WAITING` requires the lock to be
+        // held, a `store` is sufficient.
+        self.state.store(EMPTY, SeqCst);
+    }
 }
 
 impl Default for Notify {
@@ -428,6 +472,8 @@ impl Future for Notified<'_> {
                     waiters.push_front(unsafe { NonNull::new_unchecked(waiter.get()) });
 
                     *state = Waiting;
+
+                    return Poll::Pending;
                 }
                 Waiting => {
                     // Currently in the "Waiting" state, implying the caller has

tokio/src/sync/tests/loom_watch.rs

@@ -0,0 +1,20 @@
+use crate::sync::watch;
+
+use loom::future::block_on;
+use loom::thread;
+
+#[test]
+fn smoke() {
+    loom::model(|| {
+        let (tx, mut rx) = watch::channel(1);
+
+        let th = thread::spawn(move || {
+            tx.send(2).unwrap();
+        });
+
+        block_on(rx.changed()).unwrap();
+        assert_eq!(*rx.borrow(), 2);
+
+        th.join().unwrap();
+    })
+}

tokio/src/sync/tests/mod.rs

@@ -13,4 +13,5 @@ cfg_loom! {
     mod loom_oneshot;
     mod loom_semaphore_batch;
     mod loom_semaphore_ll;
+    mod loom_watch;
 }