scrappy_utils/
counter.rs

1use std::cell::Cell;
2use std::rc::Rc;
3use std::task;
4
5use crate::task::LocalWaker;
6
7#[derive(Clone)]
8/// Simple counter with ability to notify task on reaching specific number
9///
10/// Counter could be cloned, total ncount is shared across all clones.
11pub struct Counter(Rc<CounterInner>);
12
13struct CounterInner {
14    count: Cell<usize>,
15    capacity: usize,
16    task: LocalWaker,
17}
18
19impl Counter {
20    /// Create `Counter` instance and set max value.
21    pub fn new(capacity: usize) -> Self {
22        Counter(Rc::new(CounterInner {
23            capacity,
24            count: Cell::new(0),
25            task: LocalWaker::new(),
26        }))
27    }
28
29    /// Get counter guard.
30    pub fn get(&self) -> CounterGuard {
31        CounterGuard::new(self.0.clone())
32    }
33
34    /// Check if counter is not at capacity. If counter at capacity
35    /// it registers notification for current task.
36    pub fn available(&self, cx: &mut task::Context<'_>) -> bool {
37        self.0.available(cx)
38    }
39
40    /// Get total number of acquired counts
41    pub fn total(&self) -> usize {
42        self.0.count.get()
43    }
44}
45
46pub struct CounterGuard(Rc<CounterInner>);
47
48impl CounterGuard {
49    fn new(inner: Rc<CounterInner>) -> Self {
50        inner.inc();
51        CounterGuard(inner)
52    }
53}
54
55impl Unpin for CounterGuard {}
56
57impl Drop for CounterGuard {
58    fn drop(&mut self) {
59        self.0.dec();
60    }
61}
62
63impl CounterInner {
64    fn inc(&self) {
65        self.count.set(self.count.get() + 1);
66    }
67
68    fn dec(&self) {
69        let num = self.count.get();
70        self.count.set(num - 1);
71        if num == self.capacity {
72            self.task.wake();
73        }
74    }
75
76    fn available(&self, cx: &mut task::Context<'_>) -> bool {
77        if self.count.get() < self.capacity {
78            true
79        } else {
80            self.task.register(cx.waker());
81            false
82        }
83    }
84}
scrappy_utils/counter.rs

scrappy_utils/
counter.rs
