actix_utils/
counter.rs

1//! Task-notifying counter.
2
3use core::{cell::Cell, fmt, task};
4use std::rc::Rc;
5
6use local_waker::LocalWaker;
7
8/// Simple counter with ability to notify task on reaching specific number
9///
10/// Counter could be cloned, total n-count is shared across all clones.
11#[derive(Debug, Clone)]
12pub struct Counter(Rc<CounterInner>);
13
14impl Counter {
15    /// Create `Counter` instance with max value.
16    pub fn new(capacity: usize) -> Self {
17        Counter(Rc::new(CounterInner {
18            capacity,
19            count: Cell::new(0),
20            task: LocalWaker::new(),
21        }))
22    }
23
24    /// Create new counter guard, incrementing the counter.
25    #[inline]
26    pub fn get(&self) -> CounterGuard {
27        CounterGuard::new(self.0.clone())
28    }
29
30    /// Returns true if counter is below capacity. Otherwise, register to wake task when it is.
31    #[inline]
32    pub fn available(&self, cx: &mut task::Context<'_>) -> bool {
33        self.0.available(cx)
34    }
35
36    /// Get total number of acquired guards.
37    #[inline]
38    pub fn total(&self) -> usize {
39        self.0.count.get()
40    }
41}
42
43struct CounterInner {
44    count: Cell<usize>,
45    capacity: usize,
46    task: LocalWaker,
47}
48
49impl CounterInner {
50    fn inc(&self) {
51        self.count.set(self.count.get() + 1);
52    }
53
54    fn dec(&self) {
55        let num = self.count.get();
56        self.count.set(num - 1);
57        if num == self.capacity {
58            self.task.wake();
59        }
60    }
61
62    fn available(&self, cx: &mut task::Context<'_>) -> bool {
63        if self.count.get() < self.capacity {
64            true
65        } else {
66            self.task.register(cx.waker());
67            false
68        }
69    }
70}
71
72impl fmt::Debug for CounterInner {
73    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
74        f.debug_struct("Counter")
75            .field("count", &self.count.get())
76            .field("capacity", &self.capacity)
77            .field("task", &self.task)
78            .finish()
79    }
80}
81
82/// An RAII structure that keeps the underlying counter incremented until this guard is dropped.
83#[derive(Debug)]
84pub struct CounterGuard(Rc<CounterInner>);
85
86impl CounterGuard {
87    fn new(inner: Rc<CounterInner>) -> Self {
88        inner.inc();
89        CounterGuard(inner)
90    }
91}
92
93impl Unpin for CounterGuard {}
94
95impl Drop for CounterGuard {
96    fn drop(&mut self) {
97        self.0.dec();
98    }
99}
actix_utils/counter.rs

actix_utils/
counter.rs
