crabka-broker 0.3.6

Single-node Apache Kafka-compatible broker (MVP)
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
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

//! Per-broker cache of `TokenBucket`s, one per (`quota_key`, `entity_key`) pair.

use std::sync::Arc;

use crabka_metadata::EntityKey;
use dashmap::DashMap;

use crate::throttle::TokenBucket;

#[derive(Debug, Default)]
pub struct QuotaBuckets {
    /// Keyed by (`quota_key`, canonical entity key). One bucket per
    /// (`quota_type`, entity) pair, lazy-allocated on first lookup.
    buckets: DashMap<(String, EntityKey), Arc<TokenBucket>>,
}

impl QuotaBuckets {
    #[must_use]
    pub fn new() -> Self {
        Self {
            buckets: DashMap::new(),
        }
    }

    /// Get or lazily create a bucket for `(quota_key, entity_key)`,
    /// initializing it to `initial_rate` if newly created.
    #[must_use]
    pub fn get_or_create(
        &self,
        quota_key: &str,
        entity_key: &EntityKey,
        initial_rate: u64,
    ) -> Arc<TokenBucket> {
        if let Some(b) = self
            .buckets
            .get(&(quota_key.to_string(), entity_key.clone()))
        {
            return b.clone();
        }
        let b = Arc::new(TokenBucket::new());
        b.set_rate(initial_rate);
        let entry = self
            .buckets
            .entry((quota_key.to_string(), entity_key.clone()))
            .or_insert_with(|| b.clone());
        entry.clone()
    }

    /// Iterate every (`quota_key`, `entity_key`, bucket) — used by the
    /// refresh task to push new rates after an image change.
    pub fn iter(&self) -> impl Iterator<Item = ((String, EntityKey), Arc<TokenBucket>)> + '_ {
        self.buckets
            .iter()
            .map(|r| (r.key().clone(), r.value().clone()))
    }

    #[cfg(test)]
    #[must_use]
    pub fn len(&self) -> usize {
        self.buckets.len()
    }

    #[cfg(test)]
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.buckets.is_empty()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use assert2::assert;

    fn key(user: &str) -> EntityKey {
        vec![("user".into(), Some(user.into()))]
    }

    #[test]
    fn get_or_create_returns_new_bucket_first_time() {
        let buckets = QuotaBuckets::new();
        let b = buckets.get_or_create("producer_byte_rate", &key("alice"), 1024);
        assert!(b.rate() == 1024);
        assert!(buckets.len() == 1);
    }

    #[test]
    fn get_or_create_returns_existing_bucket_second_time() {
        let buckets = QuotaBuckets::new();
        let b1 = buckets.get_or_create("producer_byte_rate", &key("alice"), 1024);
        let b2 = buckets.get_or_create("producer_byte_rate", &key("alice"), 4096);
        // Same Arc — initial_rate on second call is ignored.
        assert!(Arc::ptr_eq(&b1, &b2));
        assert!(b1.rate() == 1024);
        assert!(buckets.len() == 1);
    }

    #[test]
    fn different_quota_keys_get_different_buckets() {
        let buckets = QuotaBuckets::new();
        let _ = buckets.get_or_create("producer_byte_rate", &key("alice"), 1024);
        let _ = buckets.get_or_create("consumer_byte_rate", &key("alice"), 2048);
        assert!(buckets.len() == 2);
    }

    #[test]
    fn different_entities_get_different_buckets() {
        let buckets = QuotaBuckets::new();
        let _ = buckets.get_or_create("producer_byte_rate", &key("alice"), 1024);
        let _ = buckets.get_or_create("producer_byte_rate", &key("bob"), 2048);
        assert!(buckets.len() == 2);
    }
}