feat: implement sync::Barrier

Based on the implementation in tokio-rs/tokio#1571

Author: dignifiedquire

Cargo.toml

@@ -42,6 +42,7 @@ num_cpus = "1.10.1"
 pin-utils = "0.1.0-alpha.4"
 slab = "0.4.2"
 kv-log-macro = "1.0.4"
+broadcaster = "0.2.4"
 
 [dev-dependencies]
 femme = "1.2.0"

src/sync/barrier.rs

@@ -0,0 +1,130 @@
+use crate::sync::Mutex;
+use broadcaster::BroadcastChannel;
+
+/// A barrier enables multiple threads to synchronize the beginning of some computation.
+///
+/// ```
+/// # fn main() { async_std::task::block_on(async {
+/// #
+/// use std::sync::Arc;
+/// use async_std::sync::Barrier;
+/// use futures_util::future::join_all;
+///
+/// let mut handles = Vec::with_capacity(10);
+/// let barrier = Arc::new(Barrier::new(10));
+/// for _ in 0..10 {
+///     let c = barrier.clone();
+///     // The same messages will be printed together.
+///     // You will NOT see any interleaving.
+///     handles.push(async move {
+///         println!("before wait");
+///         let wr = c.wait().await;
+///         println!("after wait");
+///         wr
+///     });
+/// }
+/// // Will not resolve until all "before wait" messages have been printed
+/// let wrs = join_all(handles).await;
+/// // Exactly one barrier will resolve as the "leader"
+/// assert_eq!(wrs.into_iter().filter(|wr| wr.is_leader()).count(), 1);
+/// # });
+/// # }
+/// ```
+// #[derive(Debug)]
+pub struct Barrier {
+    state: Mutex<BarrierState>,
+    wait: BroadcastChannel<usize>,
+    n: usize,
+}
+
+// #[derive(Debug)]
+struct BarrierState {
+    waker: BroadcastChannel<usize>,
+    arrived: usize,
+    generation: usize,
+}
+
+impl Barrier {
+    /// Creates a new barrier that can block a given number of threads.
+    ///
+    /// A barrier will block `n`-1 threads which call [`Barrier::wait`] and then wake up all
+    /// threads at once when the `n`th thread calls `wait`.
+    pub fn new(mut n: usize) -> Barrier {
+        let waker = BroadcastChannel::new();
+        let wait = waker.clone();
+
+        if n == 0 {
+            // if n is 0, it's not clear what behavior the user wants.
+            // in std::sync::Barrier, an n of 0 exhibits the same behavior as n == 1, where every
+            // .wait() immediately unblocks, so we adopt that here as well.
+            n = 1;
+        }
+
+        Barrier {
+            state: Mutex::new(BarrierState {
+                waker,
+                arrived: 0,
+                generation: 1,
+            }),
+            n,
+            wait,
+        }
+    }
+
+    /// Does not resolve until all tasks have rendezvoused here.
+    ///
+    /// Barriers are re-usable after all threads have rendezvoused once, and can
+    /// be used continuously.
+    ///
+    /// A single (arbitrary) future will receive a [`BarrierWaitResult`] that returns `true` from
+    /// [`BarrierWaitResult::is_leader`] when returning from this function, and all other threads
+    /// will receive a result that will return `false` from `is_leader`.
+    pub async fn wait(&self) -> BarrierWaitResult {
+        // NOTE: The implementation in tokio use a sync mutex, decide if we agree with their reasoning or
+        // not.
+        let mut state = self.state.lock().await;
+        let generation = state.generation;
+        state.arrived += 1;
+        if state.arrived == self.n {
+            // we are the leader for this generation
+            // wake everyone, increment the generation, and return
+            state
+                .waker
+                .send(&state.generation)
+                .await
+                .expect("there is at least one receiver");
+            state.arrived = 0;
+            state.generation += 1;
+            return BarrierWaitResult(true);
+        }
+
+        drop(state);
+
+        // we're going to have to wait for the last of the generation to arrive
+        let mut wait = self.wait.clone();
+
+        loop {
+            // 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.expect("sender hasn't been closed") >= generation {
+                break;
+            }
+        }
+
+        BarrierWaitResult(false)
+    }
+}
+
+/// A `BarrierWaitResult` is returned by `wait` when all threads in the `Barrier` have rendezvoused.
+#[derive(Debug, Clone)]
+pub struct BarrierWaitResult(bool);
+
+impl BarrierWaitResult {
+    /// Returns true if this thread from wait is the "leader thread".
+    ///
+    /// Only one thread will have `true` returned from their result, all other threads will have
+    /// `false` returned.
+    pub fn is_leader(&self) -> bool {
+        self.0
+    }
+}

src/sync/mod.rs

@@ -32,8 +32,10 @@
 #[doc(inline)]
 pub use std::sync::{Arc, Weak};
 
+pub use barrier::{Barrier, BarrierWaitResult};
 pub use mutex::{Mutex, MutexGuard};
 pub use rwlock::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 
+mod barrier;
 mod mutex;
 mod rwlock;

tests/barrier.rs

@@ -0,0 +1,97 @@
+use std::sync::Arc;
+
+use async_std::sync::Barrier;
+use async_std::task;
+
+#[test]
+fn barrier_zero_does_not_block() {
+    let b = Arc::new(Barrier::new(0));
+
+    let c = b.clone();
+    task::block_on(async move {
+        let w = task::spawn(async move { c.wait().await });
+        let wr = w.await;
+        assert!(wr.is_leader());
+    });
+
+    let c = b.clone();
+    task::block_on(async move {
+        let w = task::spawn(async move { c.wait().await });
+        let wr = w.await;
+        assert!(wr.is_leader());
+    });
+}
+
+#[test]
+fn barrier_single() {
+    let b = Arc::new(Barrier::new(1));
+
+    let c = b.clone();
+    task::block_on(async move {
+        let w = task::spawn(async move { c.wait().await });
+        let wr = w.await;
+        assert!(wr.is_leader());
+    });
+    let c = b.clone();
+    task::block_on(async move {
+        let w = task::spawn(async move { c.wait().await });
+        let wr = w.await;
+        assert!(wr.is_leader());
+    });
+    let c = b.clone();
+    task::block_on(async move {
+        let w = task::spawn(async move { c.wait().await });
+        let wr = w.await;
+        assert!(wr.is_leader());
+    });
+}
+
+#[test]
+fn barrier_tango() {
+    let b = Arc::new(Barrier::new(2));
+
+    task::block_on(async move {
+        let c = b.clone();
+        let w1 = task::spawn(async move { c.wait().await });
+
+        let c = b.clone();
+        let w2 = task::spawn(async move { c.wait().await });
+
+        let ws = futures_util::future::join_all(vec![w1, w2]).await;
+        let wr1 = &ws[0];
+        let wr2 = &ws[1];
+
+        assert!(wr1.is_leader() || wr2.is_leader());
+        assert!(!(wr1.is_leader() && wr2.is_leader()));
+    });
+}
+
+#[test]
+fn barrier_lots() {
+    let b = Arc::new(Barrier::new(100));
+
+    task::block_on(async move {
+        for _ in 0..10 {
+            let mut wait = Vec::new();
+            for _ in 0..99 {
+                let c = b.clone();
+                let w = task::spawn(async move { c.wait().await });
+                wait.push(w);
+            }
+
+            // pass the barrier
+            let c = b.clone();
+            let w = task::spawn(async move { c.wait().await });
+
+            let mut found_leader = w.await.is_leader();
+            for w in wait {
+                let wr = w.await;
+                if wr.is_leader() {
+                    assert!(!found_leader);
+                    found_leader = true;
+                }
+            }
+            assert!(found_leader);
+        }
+    });
+}