use crate::sync::watch;
use std::sync::Mutex;
#[derive(Debug)]
pub struct Barrier {
state: Mutex<BarrierState>,
wait: watch::Receiver<usize>,
n: usize,
}
#[derive(Debug)]
struct BarrierState {
waker: watch::Sender<usize>,
arrived: usize,
generation: usize,
}
impl Barrier {
pub fn new(mut n: usize) -> Barrier {
let (waker, wait) = crate::sync::watch::channel(0);
if n == 0 {
n = 1;
}
Barrier {
state: Mutex::new(BarrierState {
waker,
arrived: 0,
generation: 1,
}),
n,
wait,
}
}
pub async fn wait(&self) -> BarrierWaitResult {
let generation = {
let mut state = self.state.lock().unwrap();
let generation = state.generation;
state.arrived += 1;
if state.arrived == self.n {
state
.waker
.broadcast(state.generation)
.expect("there is at least one receiver");
state.arrived = 0;
state.generation += 1;
return BarrierWaitResult(true);
}
generation
};
let mut wait = self.wait.clone();
loop {
if wait.recv().await.expect("sender hasn't been closed") >= generation {
break;
}
}
BarrierWaitResult(false)
}
}
#[derive(Debug, Clone)]
pub struct BarrierWaitResult(bool);
impl BarrierWaitResult {
pub fn is_leader(&self) -> bool {
self.0
}
}