commonware_utils/
concurrency.rs

1//! Utilities for managing concurrency.
2
3use crate::sync::Mutex;
4use core::{
5    hash::Hash,
6    num::NonZeroU32,
7    sync::atomic::{AtomicU32, Ordering},
8};
9use std::{collections::HashSet, sync::Arc};
10
11/// Limit the concurrency of some operation without blocking.
12pub struct Limiter {
13    max: u32,
14    current: Arc<AtomicU32>,
15}
16
17impl Limiter {
18    /// Create a limiter that allows up to `max` concurrent reservations.
19    pub fn new(max: NonZeroU32) -> Self {
20        Self {
21            max: max.get(),
22            current: Arc::new(AtomicU32::new(0)),
23        }
24    }
25
26    /// Attempt to reserve a slot. Returns `None` when the limiter is saturated.
27    pub fn try_acquire(&self) -> Option<Reservation> {
28        // TODO: migrate to `try_update` once MSRV is >= 1.95
29        #[allow(deprecated)]
30        self.current
31            .fetch_update(Ordering::AcqRel, Ordering::Relaxed, |current| {
32                (current < self.max).then_some(current + 1)
33            })
34            .map(|_| Reservation {
35                current: self.current.clone(),
36            })
37            .ok()
38    }
39}
40
41/// A reservation for a slot in the [Limiter].
42pub struct Reservation {
43    current: Arc<AtomicU32>,
44}
45
46impl Drop for Reservation {
47    fn drop(&mut self) {
48        self.current.fetch_sub(1, Ordering::AcqRel);
49    }
50}
51
52/// Limit the concurrency of some keyed operation without blocking.
53pub struct KeyedLimiter<K: Eq + Hash + Clone> {
54    max: u32,
55    current: Arc<Mutex<HashSet<K>>>,
56}
57
58impl<K: Eq + Hash + Clone> KeyedLimiter<K> {
59    /// Create a limiter that allows up to `max` concurrent reservations.
60    pub fn new(max: NonZeroU32) -> Self {
61        Self {
62            max: max.get(),
63            current: Arc::new(Mutex::new(HashSet::new())),
64        }
65    }
66
67    /// Attempt to reserve a slot for a given key. Returns `None` when the limiter is saturated or
68    /// the key is already reserved.
69    pub fn try_acquire(&self, key: K) -> Option<KeyedReservation<K>> {
70        let mut current = self.current.lock();
71        if current.len() >= self.max as usize {
72            return None;
73        }
74        if !current.insert(key.clone()) {
75            return None;
76        }
77        drop(current);
78
79        Some(KeyedReservation {
80            key,
81            current: self.current.clone(),
82        })
83    }
84}
85
86/// A reservation for a slot in the [KeyedLimiter].
87pub struct KeyedReservation<K: Eq + Hash + Clone> {
88    key: K,
89    current: Arc<Mutex<HashSet<K>>>,
90}
91
92impl<K: Eq + Hash + Clone> Drop for KeyedReservation<K> {
93    fn drop(&mut self) {
94        self.current.lock().remove(&self.key);
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101    use crate::NZU32;
102
103    #[test]
104    fn allows_reservations_up_to_max() {
105        let limiter = Limiter::new(NZU32!(2));
106
107        let first = limiter
108            .try_acquire()
109            .expect("first reservation should succeed");
110        let second = limiter
111            .try_acquire()
112            .expect("second reservation should succeed");
113
114        assert!(limiter.try_acquire().is_none());
115
116        drop(second);
117        let third = limiter
118            .try_acquire()
119            .expect("reservation after drop should succeed");
120
121        drop(third);
122        drop(first);
123    }
124
125    #[test]
126    fn allows_reservations_up_to_max_for_key() {
127        let limiter = KeyedLimiter::new(NZU32!(2));
128
129        let first = limiter
130            .try_acquire(0)
131            .expect("first reservation should succeed");
132        let second = limiter
133            .try_acquire(1)
134            .expect("second reservation should succeed");
135        assert!(limiter.try_acquire(2).is_none());
136
137        drop(second);
138        let third = limiter
139            .try_acquire(2)
140            .expect("third reservation should succeed");
141
142        drop(third);
143        drop(first);
144    }
145
146    #[test]
147    fn blocks_conflicting_reservations_for_key() {
148        let limiter = KeyedLimiter::new(NZU32!(2));
149
150        let first = limiter
151            .try_acquire(0)
152            .expect("first reservation should succeed");
153        assert!(limiter.try_acquire(0).is_none());
154
155        drop(first);
156        let second = limiter
157            .try_acquire(0)
158            .expect("second reservation should succeed");
159
160        drop(second);
161    }
162}
commonware_utils/concurrency.rs

commonware_utils/
concurrency.rs
