arcium-primitives 0.4.3

Arcium primitives
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

//! Utility module for rate limiting mechanisms

pub mod token_bucket;

use std::sync::{
    atomic::{AtomicU64, Ordering},
    Arc,
};

pub use token_bucket::{TokenBucket, TokenBucketConfig};

pub trait RateLimiter: Send + Sync {
    /// The numerical type used to count tokens
    type TokenType;

    /// Try to consume a specified number of tokens
    /// Returns true if successful, false if not enough tokens available
    /// (thus not consuming any tokens)
    fn try_consume(&self, tokens: Self::TokenType) -> bool;

    /// Get the number of available tokens
    fn available_tokens(&self) -> Self::TokenType;
}

impl RateLimiter for Arc<AtomicU64> {
    type TokenType = u64;

    fn try_consume(&self, tokens: u64) -> bool {
        self.fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| {
            if current >= tokens {
                Some(current - tokens)
            } else {
                None
            }
        })
        .is_ok()
    }

    fn available_tokens(&self) -> u64 {
        self.load(Ordering::Acquire)
    }
}