pingap_core/
ttl_lru_limit.rs

1// Copyright 2024-2025 Tree xie.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use super::{now_ms, LOG_CATEGORY};
16use std::time::Duration;
17use tinyufo::TinyUfo;
18use tracing::debug;
19
20#[derive(Debug, Clone)]
21struct TtlLimit {
22    count: usize,
23    created_at: u64,
24}
25
26pub struct TtlLruLimit {
27    ttl: u64,
28    ufo: TinyUfo<String, TtlLimit>,
29    max: usize,
30}
31
32impl TtlLruLimit {
33    /// Creates a new TTL-based LRU limit with the specified parameters.
34    ///
35    /// # Arguments
36    ///
37    /// * `size` - The maximum number of entries to store in the LRU cache
38    /// * `ttl` - The time-to-live duration after which entries are considered expired
39    /// * `max` - The maximum count allowed per key within the TTL window
40    pub fn new(size: usize, ttl: Duration, max: usize) -> Self {
41        Self {
42            ttl: ttl.as_millis() as u64,
43            max,
44            ufo: TinyUfo::new(size, size),
45        }
46    }
47    /// Creates a new compact TTL-based LRU limit with the specified parameters.
48    ///
49    /// # Arguments
50    ///
51    /// * `size` - The maximum number of entries to store in the LRU cache
52    /// * `ttl` - The time-to-live duration after which entries are considered expired
53    pub fn new_compact(size: usize, ttl: Duration, max: usize) -> Self {
54        Self {
55            ttl: ttl.as_millis() as u64,
56            max,
57            ufo: TinyUfo::new_compact(size, size),
58        }
59    }
60    /// Validates whether a key has not exceeded its rate limit.
61    ///
62    /// # Arguments
63    ///
64    /// * `key` - The key to validate
65    ///
66    /// # Returns
67    ///
68    /// Returns `true` if the key is within its limit or has expired, `false` otherwise.
69    pub fn validate(&self, key: &str) -> bool {
70        let mut should_reset = false;
71        let mut valid = false;
72        let key = key.to_string();
73
74        if let Some(value) = self.ufo.get(&key) {
75            debug!(
76                category = LOG_CATEGORY,
77                key,
78                value = format!("{value:?}"),
79                "ttl lru limit"
80            );
81            // validate expired first
82            if now_ms() - value.created_at > self.ttl {
83                valid = true;
84                should_reset = true;
85            } else if value.count < self.max {
86                valid = true;
87            }
88        } else {
89            valid = true
90        }
91        if should_reset {
92            // ufo does not support remove
93            self.ufo.put(
94                key,
95                TtlLimit {
96                    count: 0,
97                    created_at: 0,
98                },
99                1,
100            );
101        }
102
103        valid
104    }
105    /// Increments the counter for the specified key.
106    /// If the key doesn't exist, creates a new entry with count 1.
107    /// If the key exists but was reset (count = 0), updates its creation timestamp.
108    ///
109    /// # Arguments
110    ///
111    /// * `key` - The key to increment
112    pub fn inc(&self, key: &str) {
113        let key = key.to_string();
114        let data = if let Some(mut value) = self.ufo.get(&key) {
115            // the reset value
116            if value.created_at == 0 {
117                value.created_at = now_ms();
118            }
119            value.count += 1;
120            value
121        } else {
122            TtlLimit {
123                count: 1,
124                created_at: now_ms(),
125            }
126        };
127        self.ufo.put(key, data, 1);
128    }
129}
130
131#[cfg(test)]
132mod test {
133    use super::TtlLruLimit;
134    use pretty_assertions::assert_eq;
135    use std::time::Duration;
136
137    #[test]
138    fn test_ttl_lru_limit() {
139        let limit = TtlLruLimit::new(5, Duration::from_millis(500), 3);
140
141        let key = "abc";
142        assert_eq!(true, limit.validate(key));
143        limit.inc(key);
144        limit.inc(key);
145        assert_eq!(true, limit.validate(key));
146        limit.inc(key);
147        assert_eq!(false, limit.validate(key));
148        std::thread::sleep(Duration::from_millis(600));
149        assert_eq!(true, limit.validate(key));
150    }
151}
pingap_core/ttl_lru_limit.rs

pingap_core/
ttl_lru_limit.rs
