use std::cell::Cell; use std::rc::Rc; use futures::task::AtomicTask; #[derive(Clone)] /// Simple counter with ability to notify task on reaching specific number /// /// Counter could be cloned, total ncount is shared across all clones. pub struct Counter(Rc); struct CounterInner { count: Cell, max: usize, task: AtomicTask, } impl Counter { /// Create `Counter` instance and set max value. pub fn new(max: usize) -> Self { Counter(Rc::new(CounterInner { max, count: Cell::new(0), task: AtomicTask::new(), })) } pub fn get(&self) -> CounterGuard { CounterGuard::new(self.0.clone()) } pub fn check(&self) -> bool { self.0.check() } pub fn total(&self) -> usize { self.0.count.get() } } pub struct CounterGuard(Rc); impl CounterGuard { fn new(inner: Rc) -> Self { inner.inc(); CounterGuard(inner) } } impl Drop for CounterGuard { fn drop(&mut self) { self.0.dec(); } } impl CounterInner { fn inc(&self) { let num = self.count.get() + 1; self.count.set(num); if num == self.max { self.task.register(); } } fn dec(&self) { let num = self.count.get(); self.count.set(num - 1); if num == self.max { self.task.notify(); } } fn check(&self) -> bool { self.count.get() < self.max } }