do_memory_core/retry/
limiter.rs

1use std::sync::Arc;
2
3use tokio::sync::{OwnedSemaphorePermit, Semaphore};
4
5/// Concurrency limiter for retries using tokio semaphore permits.
6///
7/// Wraps a [`Semaphore`] so that at most `max_concurrent` retry
8/// operations run simultaneously. Callers acquire a permit before
9/// starting the retryable work and the permit is automatically released
10/// when the returned guard is dropped.
11pub struct ConcurrencyLimiter {
12    semaphore: Arc<Semaphore>,
13    max_concurrent: usize,
14}
15
16impl ConcurrencyLimiter {
17    /// Create a new limiter.
18    pub fn new(max_concurrent: usize) -> Arc<Self> {
19        Arc::new(Self {
20            semaphore: Arc::new(Semaphore::new(max_concurrent)),
21            max_concurrent,
22        })
23    }
24
25    /// Acquire a concurrency permit.
26    ///
27    /// Returns an owned permit that releases when dropped. This method
28    /// waits (async) until a permit is available.
29    pub async fn acquire(&self) -> OwnedSemaphorePermit {
30        Arc::clone(&self.semaphore)
31            .acquire_owned()
32            .await
33            .expect("semaphore closed unexpectedly")
34    }
35
36    /// Number of permits currently available (not held).
37    pub fn available_permits(&self) -> usize {
38        self.semaphore.available_permits()
39    }
40
41    /// Maximum concurrent retries this limiter allows.
42    pub fn max_concurrent(&self) -> usize {
43        self.max_concurrent
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50    use std::sync::atomic::{AtomicUsize, Ordering};
51    use std::time::Duration;
52
53    #[tokio::test]
54    async fn limits_concurrency() {
55        let limiter = ConcurrencyLimiter::new(2);
56        let running = Arc::new(AtomicUsize::new(0));
57        let max_observed = Arc::new(AtomicUsize::new(0));
58
59        let mut handles = Vec::new();
60        for _ in 0..5 {
61            let limiter = Arc::clone(&limiter);
62            let running = Arc::clone(&running);
63            let max_observed = Arc::clone(&max_observed);
64            handles.push(tokio::spawn(async move {
65                let _permit = limiter.acquire().await;
66                let cur = running.fetch_add(1, Ordering::SeqCst) + 1;
67                max_observed.fetch_max(cur, Ordering::SeqCst);
68                tokio::time::sleep(Duration::from_millis(10)).await;
69                running.fetch_sub(1, Ordering::SeqCst);
70            }));
71        }
72
73        for h in handles {
74            h.await.unwrap();
75        }
76        assert!(
77            max_observed.load(Ordering::SeqCst) <= 2,
78            "max observed concurrency was {}",
79            max_observed.load(Ordering::SeqCst)
80        );
81    }
82
83    #[tokio::test]
84    async fn available_permits_decrease() {
85        let limiter = ConcurrencyLimiter::new(3);
86        assert_eq!(limiter.available_permits(), 3);
87
88        let p1 = limiter.acquire().await;
89        assert_eq!(limiter.available_permits(), 2);
90
91        let p2 = limiter.acquire().await;
92        assert_eq!(limiter.available_permits(), 1);
93
94        drop(p1);
95        assert_eq!(limiter.available_permits(), 2);
96
97        drop(p2);
98        assert_eq!(limiter.available_permits(), 3);
99    }
100}
do_memory_core/retry/limiter.rs

do_memory_core/retry/
limiter.rs
