diff options
Diffstat (limited to 'src/limiter.rs')
-rw-r--r-- | src/limiter.rs | 37 |
1 files changed, 37 insertions, 0 deletions
diff --git a/src/limiter.rs b/src/limiter.rs new file mode 100644 index 0000000..97bdb5c --- /dev/null +++ b/src/limiter.rs @@ -0,0 +1,37 @@ +use std::sync::{atomic::AtomicUsize, Arc}; + +pub struct Limiter { + limit: usize, + counter: Arc<AtomicUsize>, +} + +impl Limiter { + pub fn new(limit: usize) -> Self { + Limiter { + counter: Default::default(), + limit, + } + } + pub fn obtain(&self) -> Option<LimitLock> { + let c = self + .counter + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + if c < self.limit { + Some(LimitLock { + counter: self.counter.clone(), + }) + } else { + None + } + } +} + +pub struct LimitLock { + counter: Arc<AtomicUsize>, +} +impl Drop for LimitLock { + fn drop(&mut self) { + self.counter + .fetch_sub(1, std::sync::atomic::Ordering::Relaxed); + } +} |