diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 000000000..32a9786fa --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1 @@ +edition = "2018" diff --git a/tokio-sync/src/barrier.rs b/tokio-sync/src/barrier.rs new file mode 100644 index 000000000..f71b4de9b --- /dev/null +++ b/tokio-sync/src/barrier.rs @@ -0,0 +1,130 @@ +use crate::watch; +use std::sync::Mutex; + +/// A barrier enables multiple threads to synchronize the beginning of some computation. +/// +/// ``` +/// # #[tokio::main] +/// # async fn main() { +/// use std::sync::Arc; +/// use tokio_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, + wait: watch::Receiver, + n: usize, +} + +#[derive(Debug)] +struct BarrierState { + waker: watch::Sender, + 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, wait) = crate::watch::channel(0); + + 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: we are taking a _synchronous_ lock here. + // It is okay to do so because the critical section is fast and never yields, so it cannot + // deadlock even if another future is concurrently holding the lock. + // It is _desireable_ to do so as synchronous Mutexes are, at least in theory, faster than + // the asynchronous counter-parts, so we should use them where possible [citation needed]. + let mut state = self.state.lock().unwrap(); + 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 + .broadcast(state.generation) + .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 + } +} diff --git a/tokio-sync/src/lib.rs b/tokio-sync/src/lib.rs index 868ed71a9..4b472b0d4 100644 --- a/tokio-sync/src/lib.rs +++ b/tokio-sync/src/lib.rs @@ -29,6 +29,7 @@ macro_rules! if_fuzz { }} } +mod barrier; mod loom; pub mod mpsc; mod mutex; @@ -37,5 +38,6 @@ pub mod semaphore; mod task; pub mod watch; +pub use barrier::{Barrier, BarrierWaitResult}; pub use mutex::{Mutex, MutexGuard}; pub use task::AtomicWaker; diff --git a/tokio-sync/tests/barrier.rs b/tokio-sync/tests/barrier.rs new file mode 100644 index 000000000..b3c40e8ff --- /dev/null +++ b/tokio-sync/tests/barrier.rs @@ -0,0 +1,86 @@ +#![warn(rust_2018_idioms)] + +use tokio_sync::Barrier; +use tokio_test::task::spawn; +use tokio_test::{assert_pending, assert_ready}; + +#[test] +fn zero_does_not_block() { + let b = Barrier::new(0); + + { + let mut w = spawn(b.wait()); + let wr = assert_ready!(w.poll()); + assert!(wr.is_leader()); + } + { + let mut w = spawn(b.wait()); + let wr = assert_ready!(w.poll()); + assert!(wr.is_leader()); + } +} + +#[test] +fn single() { + let b = Barrier::new(1); + + { + let mut w = spawn(b.wait()); + let wr = assert_ready!(w.poll()); + assert!(wr.is_leader()); + } + { + let mut w = spawn(b.wait()); + let wr = assert_ready!(w.poll()); + assert!(wr.is_leader()); + } + { + let mut w = spawn(b.wait()); + let wr = assert_ready!(w.poll()); + assert!(wr.is_leader()); + } +} + +#[test] +fn tango() { + let b = Barrier::new(2); + + let mut w1 = spawn(b.wait()); + assert_pending!(w1.poll()); + + let mut w2 = spawn(b.wait()); + let wr2 = assert_ready!(w2.poll()); + let wr1 = assert_ready!(w1.poll()); + + assert!(wr1.is_leader() || wr2.is_leader()); + assert!(!(wr1.is_leader() && wr2.is_leader())); +} + +#[test] +fn lots() { + let b = Barrier::new(100); + + for _ in 0..10 { + let mut wait = Vec::new(); + for _ in 0..99 { + let mut w = spawn(b.wait()); + assert_pending!(w.poll()); + wait.push(w); + } + for w in &mut wait { + assert_pending!(w.poll()); + } + + // pass the barrier + let mut w = spawn(b.wait()); + let mut found_leader = assert_ready!(w.poll()).is_leader(); + for mut w in wait { + let wr = assert_ready!(w.poll()); + if wr.is_leader() { + assert!(!found_leader); + found_leader = true; + } + } + assert!(found_leader); + } +}