kraken_api_client/rate_limit/
ttl_cache.rs

1//! Time-to-live cache for tracking order lifetimes.
2//!
3//! This cache automatically expires entries after a configurable duration.
4//! It's used for tracking order creation times to calculate rate limit penalties
5//! when orders are cancelled.
6//!
7//! # Example
8//!
9//! ```rust
10//! use std::time::Duration;
11//! use kraken_api_client::rate_limit::TtlCache;
12//!
13//! let mut cache: TtlCache<String, i64> = TtlCache::new(Duration::from_secs(300));
14//!
15//! // Insert an order
16//! cache.insert("O123".to_string(), 1234567890);
17//!
18//! // Check if it exists
19//! assert!(cache.get(&"O123".to_string()).is_some());
20//!
21//! // Remove an order
22//! cache.remove(&"O123".to_string());
23//! assert!(cache.get(&"O123".to_string()).is_none());
24//! ```
25
26use std::collections::HashMap;
27use std::hash::Hash;
28use std::time::{Duration, Instant};
29
30/// A cache that automatically expires entries after a configurable TTL.
31///
32/// This is useful for tracking order lifetimes in rate limiting, where
33/// orders cancelled within certain time windows incur different penalties.
34#[derive(Debug)]
35pub struct TtlCache<K, V> {
36    cache: HashMap<K, (V, Instant)>,
37    ttl: Duration,
38}
39
40impl<K, V> TtlCache<K, V>
41where
42    K: Hash + Eq,
43{
44    /// Create a new TTL cache with the specified time-to-live duration.
45    ///
46    /// Entries will be considered expired after this duration.
47    pub fn new(ttl: Duration) -> Self {
48        Self {
49            cache: HashMap::new(),
50            ttl,
51        }
52    }
53
54    /// Create a new TTL cache with a specific initial capacity.
55    pub fn with_capacity(ttl: Duration, capacity: usize) -> Self {
56        Self {
57            cache: HashMap::with_capacity(capacity),
58            ttl,
59        }
60    }
61
62    /// Insert a key-value pair into the cache.
63    ///
64    /// The entry will be timestamped with the current time.
65    pub fn insert(&mut self, key: K, value: V) {
66        self.cache.insert(key, (value, Instant::now()));
67    }
68
69    /// Get a reference to a value if it exists and hasn't expired.
70    pub fn get(&self, key: &K) -> Option<&V> {
71        self.cache.get(key).and_then(|(value, timestamp)| {
72            if timestamp.elapsed() < self.ttl {
73                Some(value)
74            } else {
75                None
76            }
77        })
78    }
79
80    /// Get a mutable reference to a value if it exists and hasn't expired.
81    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
82        let ttl = self.ttl;
83        self.cache.get_mut(key).and_then(|(value, timestamp)| {
84            if timestamp.elapsed() < ttl {
85                Some(value)
86            } else {
87                None
88            }
89        })
90    }
91
92    /// Get the timestamp when the entry was inserted.
93    ///
94    /// Returns `None` if the key doesn't exist or has expired.
95    pub fn get_timestamp(&self, key: &K) -> Option<Instant> {
96        self.cache.get(key).and_then(|(_, timestamp)| {
97            if timestamp.elapsed() < self.ttl {
98                Some(*timestamp)
99            } else {
100                None
101            }
102        })
103    }
104
105    /// Get the age of an entry in the cache.
106    ///
107    /// Returns `None` if the key doesn't exist or has expired.
108    pub fn get_age(&self, key: &K) -> Option<Duration> {
109        self.cache.get(key).and_then(|(_, timestamp)| {
110            let age = timestamp.elapsed();
111            if age < self.ttl {
112                Some(age)
113            } else {
114                None
115            }
116        })
117    }
118
119    /// Remove an entry from the cache.
120    ///
121    /// Returns the value if it existed and hadn't expired, `None` otherwise.
122    pub fn remove(&mut self, key: &K) -> Option<V> {
123        self.cache.remove(key).and_then(|(value, timestamp)| {
124            if timestamp.elapsed() < self.ttl {
125                Some(value)
126            } else {
127                None
128            }
129        })
130    }
131
132    /// Remove an entry and return both the value and its age.
133    ///
134    /// Useful for calculating rate limit penalties based on order age.
135    pub fn remove_with_age(&mut self, key: &K) -> Option<(V, Duration)> {
136        self.cache.remove(key).and_then(|(value, timestamp)| {
137            let age = timestamp.elapsed();
138            if age < self.ttl {
139                Some((value, age))
140            } else {
141                None
142            }
143        })
144    }
145
146    /// Check if a key exists and hasn't expired.
147    pub fn contains(&self, key: &K) -> bool {
148        self.get(key).is_some()
149    }
150
151    /// Remove all expired entries from the cache.
152    ///
153    /// Call this periodically to free memory from expired entries.
154    pub fn cleanup(&mut self) {
155        let ttl = self.ttl;
156        self.cache.retain(|_, (_, timestamp)| timestamp.elapsed() < ttl);
157    }
158
159    /// Get the number of entries in the cache (including expired ones).
160    pub fn len(&self) -> usize {
161        self.cache.len()
162    }
163
164    /// Check if the cache is empty.
165    pub fn is_empty(&self) -> bool {
166        self.cache.is_empty()
167    }
168
169    /// Get the number of non-expired entries.
170    pub fn active_count(&self) -> usize {
171        let ttl = self.ttl;
172        self.cache
173            .values()
174            .filter(|(_, timestamp)| timestamp.elapsed() < ttl)
175            .count()
176    }
177
178    /// Clear all entries from the cache.
179    pub fn clear(&mut self) {
180        self.cache.clear();
181    }
182
183    /// Get the TTL duration for this cache.
184    pub fn ttl(&self) -> Duration {
185        self.ttl
186    }
187
188    /// Set a new TTL duration.
189    ///
190    /// This affects all future checks but doesn't modify existing timestamps.
191    pub fn set_ttl(&mut self, ttl: Duration) {
192        self.ttl = ttl;
193    }
194}
195
196impl<K, V> Default for TtlCache<K, V>
197where
198    K: Hash + Eq,
199{
200    fn default() -> Self {
201        // Default TTL of 5 minutes (300 seconds) as per Kraken's order penalty window
202        Self::new(Duration::from_secs(300))
203    }
204}
205
206#[cfg(test)]
207mod tests {
208    use super::*;
209    use std::thread;
210
211    #[test]
212    fn test_insert_and_get() {
213        let mut cache: TtlCache<String, i32> = TtlCache::new(Duration::from_secs(60));
214
215        cache.insert("key1".to_string(), 100);
216        assert_eq!(cache.get(&"key1".to_string()), Some(&100));
217        assert_eq!(cache.get(&"key2".to_string()), None);
218    }
219
220    #[test]
221    fn test_remove() {
222        let mut cache: TtlCache<String, i32> = TtlCache::new(Duration::from_secs(60));
223
224        cache.insert("key1".to_string(), 100);
225        assert_eq!(cache.remove(&"key1".to_string()), Some(100));
226        assert_eq!(cache.get(&"key1".to_string()), None);
227    }
228
229    #[test]
230    fn test_expiration() {
231        let mut cache: TtlCache<String, i32> = TtlCache::new(Duration::from_millis(50));
232
233        cache.insert("key1".to_string(), 100);
234        assert!(cache.get(&"key1".to_string()).is_some());
235
236        // Wait for expiration
237        thread::sleep(Duration::from_millis(60));
238        assert!(cache.get(&"key1".to_string()).is_none());
239    }
240
241    #[test]
242    fn test_cleanup() {
243        let mut cache: TtlCache<String, i32> = TtlCache::new(Duration::from_millis(50));
244
245        cache.insert("key1".to_string(), 100);
246        cache.insert("key2".to_string(), 200);
247        assert_eq!(cache.len(), 2);
248
249        // Wait for expiration
250        thread::sleep(Duration::from_millis(60));
251
252        // Entry still in HashMap but expired
253        assert_eq!(cache.len(), 2);
254
255        // Cleanup removes expired entries
256        cache.cleanup();
257        assert_eq!(cache.len(), 0);
258    }
259
260    #[test]
261    fn test_get_age() {
262        let mut cache: TtlCache<String, i32> = TtlCache::new(Duration::from_secs(60));
263
264        cache.insert("key1".to_string(), 100);
265        thread::sleep(Duration::from_millis(50));
266
267        let age = cache.get_age(&"key1".to_string()).unwrap();
268        assert!(age >= Duration::from_millis(50));
269        assert!(age < Duration::from_millis(100));
270    }
271
272    #[test]
273    fn test_remove_with_age() {
274        let mut cache: TtlCache<String, i32> = TtlCache::new(Duration::from_secs(60));
275
276        cache.insert("key1".to_string(), 100);
277        thread::sleep(Duration::from_millis(50));
278
279        let (value, age) = cache.remove_with_age(&"key1".to_string()).unwrap();
280        assert_eq!(value, 100);
281        assert!(age >= Duration::from_millis(50));
282    }
283
284    #[test]
285    fn test_contains() {
286        let mut cache: TtlCache<String, i32> = TtlCache::new(Duration::from_secs(60));
287
288        cache.insert("key1".to_string(), 100);
289        assert!(cache.contains(&"key1".to_string()));
290        assert!(!cache.contains(&"key2".to_string()));
291    }
292
293    #[test]
294    fn test_active_count() {
295        let mut cache: TtlCache<String, i32> = TtlCache::new(Duration::from_millis(50));
296
297        cache.insert("key1".to_string(), 100);
298        cache.insert("key2".to_string(), 200);
299        assert_eq!(cache.active_count(), 2);
300
301        thread::sleep(Duration::from_millis(60));
302        assert_eq!(cache.active_count(), 0);
303    }
304}
kraken_api_client/rate_limit/ttl_cache.rs

kraken_api_client/rate_limit/
ttl_cache.rs
