apply limits

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);
+    }
+}