src/libstd/sync/barrier.rs

   1 use crate::fmt;
   2 use crate::sync::{Condvar, Mutex};
   3
   4 /// A barrier enables multiple threads to synchronize the beginning
   5 /// of some computation.
   6 ///
   7 /// # Examples
   8 ///
   9 /// ```
  10 /// use std::sync::{Arc, Barrier};
  11 /// use std::thread;
  12 ///
  13 /// let mut handles = Vec::with_capacity(10);
  14 /// let barrier = Arc::new(Barrier::new(10));
  15 /// for _ in 0..10 {
  16 ///     let c = barrier.clone();
  17 ///     // The same messages will be printed together.
  18 ///     // You will NOT see any interleaving.
  19 ///     handles.push(thread::spawn(move|| {
  20 ///         println!("before wait");
  21 ///         c.wait();
  22 ///         println!("after wait");
  23 ///     }));
  24 /// }
  25 /// // Wait for other threads to finish.
  26 /// for handle in handles {
  27 ///     handle.join().unwrap();
  28 /// }
  29 /// ```
  30 #[stable(feature = "rust1", since = "1.0.0")]
  31 pub struct Barrier {
  32     lock: Mutex<BarrierState>,
  33     cvar: Condvar,
  34     num_threads: usize,
  35 }
  36
  37 // The inner state of a double barrier
  38 struct BarrierState {
  39     count: usize,
  40     generation_id: usize,
  41 }
  42
  43 /// A `BarrierWaitResult` is returned by [`wait`] when all threads in the [`Barrier`]
  44 /// have rendezvoused.
  45 ///
  46 /// [`wait`]: struct.Barrier.html#method.wait
  47 /// [`Barrier`]: struct.Barrier.html
  48 ///
  49 /// # Examples
  50 ///
  51 /// ```
  52 /// use std::sync::Barrier;
  53 ///
  54 /// let barrier = Barrier::new(1);
  55 /// let barrier_wait_result = barrier.wait();
  56 /// ```
  57 #[stable(feature = "rust1", since = "1.0.0")]
  58 pub struct BarrierWaitResult(bool);
  59
  60 #[stable(feature = "std_debug", since = "1.16.0")]
  61 impl fmt::Debug for Barrier {
  62     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  63         f.pad("Barrier { .. }")
  64     }
  65 }
  66
  67 impl Barrier {
  68     /// Creates a new barrier that can block a given number of threads.
  69     ///
  70     /// A barrier will block `n`-1 threads which call [`wait`] and then wake up
  71     /// all threads at once when the `n`th thread calls [`wait`].
  72     ///
  73     /// [`wait`]: #method.wait
  74     ///
  75     /// # Examples
  76     ///
  77     /// ```
  78     /// use std::sync::Barrier;
  79     ///
  80     /// let barrier = Barrier::new(10);
  81     /// ```
  82     #[stable(feature = "rust1", since = "1.0.0")]
  83     pub fn new(n: usize) -> Barrier {
  84         Barrier {
  85             lock: Mutex::new(BarrierState { count: 0, generation_id: 0 }),
  86             cvar: Condvar::new(),
  87             num_threads: n,
  88         }
  89     }
  90
  91     /// Blocks the current thread until all threads have rendezvoused here.
  92     ///
  93     /// Barriers are re-usable after all threads have rendezvoused once, and can
  94     /// be used continuously.
  95     ///
  96     /// A single (arbitrary) thread will receive a [`BarrierWaitResult`] that
  97     /// returns `true` from [`is_leader`] when returning from this function, and
  98     /// all other threads will receive a result that will return `false` from
  99     /// [`is_leader`].
 100     ///
 101     /// [`BarrierWaitResult`]: struct.BarrierWaitResult.html
 102     /// [`is_leader`]: struct.BarrierWaitResult.html#method.is_leader
 103     ///
 104     /// # Examples
 105     ///
 106     /// ```
 107     /// use std::sync::{Arc, Barrier};
 108     /// use std::thread;
 109     ///
 110     /// let mut handles = Vec::with_capacity(10);
 111     /// let barrier = Arc::new(Barrier::new(10));
 112     /// for _ in 0..10 {
 113     ///     let c = barrier.clone();
 114     ///     // The same messages will be printed together.
 115     ///     // You will NOT see any interleaving.
 116     ///     handles.push(thread::spawn(move|| {
 117     ///         println!("before wait");
 118     ///         c.wait();
 119     ///         println!("after wait");
 120     ///     }));
 121     /// }
 122     /// // Wait for other threads to finish.
 123     /// for handle in handles {
 124     ///     handle.join().unwrap();
 125     /// }
 126     /// ```
 127     #[stable(feature = "rust1", since = "1.0.0")]
 128     pub fn wait(&self) -> BarrierWaitResult {
 129         let mut lock = self.lock.lock().unwrap();
 130         let local_gen = lock.generation_id;
 131         lock.count += 1;
 132         if lock.count < self.num_threads {
 133             // We need a while loop to guard against spurious wakeups.
 134             // http://en.wikipedia.org/wiki/Spurious_wakeup
 135             while local_gen == lock.generation_id && lock.count < self.num_threads {
 136                 lock = self.cvar.wait(lock).unwrap();
 137             }
 138             BarrierWaitResult(false)
 139         } else {
 140             lock.count = 0;
 141             lock.generation_id = lock.generation_id.wrapping_add(1);
 142             self.cvar.notify_all();
 143             BarrierWaitResult(true)
 144         }
 145     }
 146 }
 147
 148 #[stable(feature = "std_debug", since = "1.16.0")]
 149 impl fmt::Debug for BarrierWaitResult {
 150     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 151         f.debug_struct("BarrierWaitResult").field("is_leader", &self.is_leader()).finish()
 152     }
 153 }
 154
 155 impl BarrierWaitResult {
 156     /// Returns `true` if this thread from [`wait`] is the "leader thread".
 157     ///
 158     /// Only one thread will have `true` returned from their result, all other
 159     /// threads will have `false` returned.
 160     ///
 161     /// [`wait`]: struct.Barrier.html#method.wait
 162     ///
 163     /// # Examples
 164     ///
 165     /// ```
 166     /// use std::sync::Barrier;
 167     ///
 168     /// let barrier = Barrier::new(1);
 169     /// let barrier_wait_result = barrier.wait();
 170     /// println!("{:?}", barrier_wait_result.is_leader());
 171     /// ```
 172     #[stable(feature = "rust1", since = "1.0.0")]
 173     pub fn is_leader(&self) -> bool {
 174         self.0
 175     }
 176 }
 177
 178 #[cfg(test)]
 179 mod tests {
 180     use crate::sync::mpsc::{channel, TryRecvError};
 181     use crate::sync::{Arc, Barrier};
 182     use crate::thread;
 183
 184     #[test]
 185     #[cfg_attr(target_os = "emscripten", ignore)]
 186     fn test_barrier() {
 187         const N: usize = 10;
 188
 189         let barrier = Arc::new(Barrier::new(N));
 190         let (tx, rx) = channel();
 191
 192         for _ in 0..N - 1 {
 193             let c = barrier.clone();
 194             let tx = tx.clone();
 195             thread::spawn(move || {
 196                 tx.send(c.wait().is_leader()).unwrap();
 197             });
 198         }
 199
 200         // At this point, all spawned threads should be blocked,
 201         // so we shouldn't get anything from the port
 202         assert!(matches!(rx.try_recv(), Err(TryRecvError::Empty)));
 203
 204         let mut leader_found = barrier.wait().is_leader();
 205
 206         // Now, the barrier is cleared and we should get data.
 207         for _ in 0..N - 1 {
 208             if rx.recv().unwrap() {
 209                 assert!(!leader_found);
 210                 leader_found = true;
 211             }
 212         }
 213         assert!(leader_found);
 214     }
 215 }