ntex_util/services/
counter.rs

1use std::{cell::Cell, cell::RefCell, future::poll_fn, rc::Rc, task::Context, task::Poll};
2
3use crate::task::LocalWaker;
4
5/// A shared count with an asynchronous capacity notification.
6///
7/// Clones use the same count and capacity, but each clone registers its own
8/// waiting task.
9#[derive(Debug)]
10pub struct Counter(usize, Rc<CounterInner>);
11
12#[derive(Debug)]
13struct CounterInner {
14    count: Cell<usize>,
15    capacity: Cell<usize>,
16    tasks: RefCell<slab::Slab<LocalWaker>>,
17}
18
19impl Counter {
20    /// Creates a counter with the specified capacity.
21    pub fn new(capacity: usize) -> Self {
22        let mut tasks = slab::Slab::new();
23        let idx = tasks.insert(LocalWaker::new());
24
25        Counter(
26            idx,
27            Rc::new(CounterInner {
28                count: Cell::new(0),
29                capacity: Cell::new(capacity),
30                tasks: RefCell::new(tasks),
31            }),
32        )
33    }
34
35    /// Acquires one count and returns a guard that releases it on drop.
36    ///
37    /// This does not wait for capacity; call [`available`](Self::available)
38    /// first when exceeding the configured capacity is not acceptable.
39    pub fn get(&self) -> CounterGuard {
40        CounterGuard::new(self.1.clone())
41    }
42
43    /// Changes the capacity and wakes tasks waiting for availability.
44    pub fn set_capacity(&self, cap: usize) {
45        self.1.capacity.set(cap);
46        self.1.notify();
47    }
48
49    /// Returns `true` if another count can be acquired without exceeding the capacity.
50    pub fn is_available(&self) -> bool {
51        self.1.count.get() < self.1.capacity.get()
52    }
53
54    /// Waits until the counter has free capacity.
55    ///
56    /// Returns immediately if there is capacity available. Otherwise,
57    /// registers the current task for wakeup and waits until a slot is freed.
58    pub async fn available(&self) {
59        poll_fn(|cx| {
60            if self.poll_available(cx) {
61                Poll::Ready(())
62            } else {
63                Poll::Pending
64            }
65        })
66        .await;
67    }
68
69    /// Waits until the counter reaches its capacity (i.e., becomes unavailable).
70    pub async fn unavailable(&self) {
71        poll_fn(|cx| {
72            if self.poll_available(cx) {
73                Poll::Pending
74            } else {
75                Poll::Ready(())
76            }
77        })
78        .await;
79    }
80
81    /// Check if counter is not at capacity. If counter at capacity
82    /// it registers notification for current task.
83    fn poll_available(&self, cx: &mut Context<'_>) -> bool {
84        if self.1.count.get() < self.1.capacity.get() {
85            true
86        } else {
87            let tasks = self.1.tasks.borrow();
88            tasks[self.0].register(cx.waker());
89            false
90        }
91    }
92
93    /// Returns the number of currently held guards.
94    pub fn total(&self) -> usize {
95        self.1.count.get()
96    }
97}
98
99impl Clone for Counter {
100    fn clone(&self) -> Self {
101        let idx = self.1.tasks.borrow_mut().insert(LocalWaker::new());
102        Self(idx, self.1.clone())
103    }
104}
105
106impl Drop for Counter {
107    fn drop(&mut self) {
108        self.1.tasks.borrow_mut().remove(self.0);
109    }
110}
111
112#[derive(Debug)]
113/// An acquired counter slot.
114///
115/// Dropping the guard releases the slot and wakes availability waiters.
116pub struct CounterGuard(Rc<CounterInner>);
117
118impl CounterGuard {
119    fn new(inner: Rc<CounterInner>) -> Self {
120        inner.inc();
121        CounterGuard(inner)
122    }
123}
124
125impl Unpin for CounterGuard {}
126
127impl Drop for CounterGuard {
128    fn drop(&mut self) {
129        self.0.dec();
130    }
131}
132
133impl CounterInner {
134    fn inc(&self) {
135        self.count.set(self.count.get() + 1);
136    }
137
138    fn dec(&self) {
139        let num = self.count.get();
140        self.count.set(num - 1);
141        if num == self.capacity.get() {
142            self.notify();
143        }
144    }
145
146    fn notify(&self) {
147        let tasks = self.tasks.borrow();
148        for (_, task) in &*tasks {
149            task.wake();
150        }
151    }
152}
ntex_util/services/counter.rs

ntex_util/services/
counter.rs
