use crate::{Event, EventListener};
use core::cell::Cell;
use core::pin::Pin;
pub struct Barrier {
n: usize,
count: Cell<usize>,
generation: Cell<u64>,
event: Event<()>,
}
impl Barrier {
pub fn new(n: usize) -> Barrier {
Barrier {
n,
count: Cell::new(0),
generation: Cell::new(0),
event: Event::new(),
}
}
pub async fn wait(&self) -> BarrierWaitResult {
let local_gen = self.generation.get();
self.count.set(self.count.get() + 1);
let mut listener = EventListener::new(&self.event);
let mut listener = unsafe { Pin::new_unchecked(&mut listener) };
if self.count.get() < self.n {
while local_gen == self.generation.get() && self.count.get() < self.n {
listener.as_mut().await;
}
BarrierWaitResult { is_leader: false }
} else {
self.count.set(0);
self.generation.set(local_gen + 1);
self.event.notify(core::usize::MAX);
BarrierWaitResult { is_leader: true }
}
}
}
#[derive(Debug, Clone)]
pub struct BarrierWaitResult {
is_leader: bool,
}
impl BarrierWaitResult {
pub fn is_leader(&self) -> bool {
self.is_leader
}
}