senax_common/linker/
stream.rs

1use fxhash::FxHashMap;
2use std::{
3    borrow::Borrow,
4    hash::Hash,
5    time::{Duration, SystemTime, UNIX_EPOCH},
6};
7use tokio::sync::RwLock;
8
9pub struct StreamCounter<K> {
10    tick: u64,
11    divisor: u64,
12    counter: RwLock<FxHashMap<u64, FxHashMap<K, usize>>>,
13}
14
15impl<K> StreamCounter<K>
16where
17    K: Hash + Eq,
18{
19    pub fn new(span: Duration, divisor: u64) -> Self {
20        let span = span.as_millis() as u64;
21        let tick = std::cmp::max(1, span / divisor);
22        Self {
23            tick,
24            divisor,
25            counter: Default::default(),
26        }
27    }
28
29    pub async fn add(&self, key: K) {
30        self.add_with_time(key, SystemTime::now()).await;
31    }
32
33    pub async fn add_with_time(&self, key: K, time: SystemTime) {
34        let mut counter = self.counter.write().await;
35        let time = time.duration_since(UNIX_EPOCH).unwrap().as_millis() as u64 / self.tick;
36        let c = counter.entry(time).or_default().entry(key).or_default();
37        *c += 1;
38        let limit = time - self.divisor;
39        counter.retain(|t, _| *t > limit);
40    }
41
42    pub async fn count<Q>(&self, key: &Q)
43    where
44        K: Borrow<Q>,
45        Q: Hash + Eq + ?Sized,
46    {
47        self.count_with_time(key, SystemTime::now()).await;
48    }
49
50    pub async fn count_with_time<Q>(&self, key: &Q, time: SystemTime) -> usize
51    where
52        K: Borrow<Q>,
53        Q: Hash + Eq + ?Sized,
54    {
55        let counter = self.counter.read().await;
56        let time =
57            time.duration_since(UNIX_EPOCH).unwrap().as_millis() as u64 / self.tick - self.divisor;
58        let mut count = 0;
59        for (t, map) in counter.iter() {
60            if *t > time
61                && let Some(c) = map.get(key.borrow())
62            {
63                count += c;
64            }
65        }
66        count
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73    #[tokio::test]
74    async fn test() {
75        let counter = StreamCounter::new(Duration::from_secs(60), 10);
76        let now = UNIX_EPOCH + Duration::from_secs(600);
77        counter.add_with_time("key".to_owned(), now).await;
78        assert_eq!(counter.count_with_time("key", now).await, 1);
79        assert_eq!(counter.count_with_time("dummy", now).await, 0);
80        let now = now + Duration::from_secs(10);
81        counter.add_with_time("key".to_owned(), now).await;
82        assert_eq!(counter.count_with_time("key", now).await, 2);
83        let now = now + Duration::from_secs(10);
84        counter.add_with_time("key".to_owned(), now).await;
85        assert_eq!(counter.count_with_time("key", now).await, 3);
86        let now = now + Duration::from_secs(10);
87        counter.add_with_time("key".to_owned(), now).await;
88        assert_eq!(counter.count_with_time("key", now).await, 4);
89        let now = now + Duration::from_secs(10);
90        counter.add_with_time("key".to_owned(), now).await;
91        assert_eq!(counter.count_with_time("key", now).await, 5);
92        let now = now + Duration::from_secs(10);
93        counter.add_with_time("key".to_owned(), now).await;
94        assert_eq!(counter.count_with_time("key", now).await, 6);
95        let now = now + Duration::from_secs(9);
96        counter.add_with_time("key".to_owned(), now).await;
97        assert_eq!(counter.count_with_time("key", now).await, 7);
98        let now = now + Duration::from_secs(2);
99        counter.add_with_time("key".to_owned(), now).await;
100        counter.add_with_time("key".to_owned(), now).await;
101        assert_eq!(counter.count_with_time("key", now).await, 8);
102        let now = now + Duration::from_secs(10);
103        counter.add_with_time("key".to_owned(), now).await;
104        counter.add_with_time("key".to_owned(), now).await;
105        assert_eq!(counter.count_with_time("key", now).await, 9);
106        let now = now + Duration::from_secs(10);
107        counter.add_with_time("key".to_owned(), now).await;
108        counter.add_with_time("key".to_owned(), now).await;
109        assert_eq!(counter.count_with_time("key", now).await, 10);
110    }
111}
senax_common/linker/stream.rs

senax_common/linker/
stream.rs
