shared_state/
ttl.rs

1//! Time-to-live cache backed by DashMap.
2
3use dashmap::DashMap;
4use std::hash::Hash;
5use std::time::{Duration, Instant};
6
7/// A cache entry with expiration time.
8struct CacheEntry<V> {
9    value: V,
10    expires_at: Instant,
11}
12
13/// A concurrent cache with per-entry TTL eviction.
14///
15/// Entries are automatically considered expired after their TTL elapses.
16/// Expired entries are lazily evicted on access.
17pub struct TtlCache<K, V> {
18    inner: DashMap<K, CacheEntry<V>>,
19    ttl: Duration,
20}
21
22impl<K: Eq + Hash + Clone, V: Clone> TtlCache<K, V> {
23    /// Create a new cache with the given TTL for all entries.
24    pub fn new(ttl: Duration) -> Self {
25        Self {
26            inner: DashMap::new(),
27            ttl,
28        }
29    }
30
31    /// Get a value if it exists and has not expired.
32    pub fn get(&self, key: &K) -> Option<V> {
33        let entry = self.inner.get(key)?;
34        if entry.expires_at > Instant::now() {
35            Some(entry.value.clone())
36        } else {
37            drop(entry);
38            self.inner.remove(key);
39            None
40        }
41    }
42
43    /// Insert a value with the configured TTL.
44    pub fn insert(&self, key: K, value: V) {
45        self.inner.insert(
46            key,
47            CacheEntry {
48                value,
49                expires_at: Instant::now() + self.ttl,
50            },
51        );
52    }
53
54    /// Remove an entry.
55    pub fn remove(&self, key: &K) -> Option<V> {
56        self.inner.remove(key).map(|e| e.1.value)
57    }
58
59    /// Remove an entry only if it has not expired, returning its value.
60    ///
61    /// Unlike [`TtlCache::remove`] (which returns the value even when
62    /// expired), this is the single-use-consume primitive: an expired entry
63    /// is dropped and reported as absent — exactly the semantics a CSRF
64    /// state store or one-time token cache needs.
65    pub fn take_fresh(&self, key: &K) -> Option<V> {
66        let (_, entry) = self.inner.remove(key)?;
67        if entry.expires_at > Instant::now() {
68            Some(entry.value)
69        } else {
70            None
71        }
72    }
73
74    /// Remove all expired entries.
75    pub fn cleanup(&self) {
76        let now = Instant::now();
77        self.inner.retain(|_, entry| entry.expires_at > now);
78    }
79
80    /// Return the number of entries (including expired ones).
81    pub fn len(&self) -> usize {
82        self.inner.len()
83    }
84
85    /// Return true if the cache is empty.
86    pub fn is_empty(&self) -> bool {
87        self.inner.is_empty()
88    }
89}
90
91impl<K: Eq + Hash + Clone, V: Clone> Default for TtlCache<K, V> {
92    fn default() -> Self {
93        Self::new(Duration::from_secs(300))
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100    use std::time::Duration;
101
102    #[test]
103    fn insert_and_get() {
104        let cache = TtlCache::new(Duration::from_secs(60));
105        cache.insert("key1", 42);
106        assert_eq!(cache.get(&"key1"), Some(42));
107    }
108
109    #[test]
110    fn expired_entry_returns_none() {
111        let cache = TtlCache::new(Duration::from_millis(1));
112        cache.insert("key1", 42);
113        std::thread::sleep(Duration::from_millis(5));
114        assert_eq!(cache.get(&"key1"), None);
115    }
116
117    #[test]
118    fn remove_entry() {
119        let cache = TtlCache::new(Duration::from_secs(60));
120        cache.insert("key1", 42);
121        assert_eq!(cache.remove(&"key1"), Some(42));
122        assert_eq!(cache.get(&"key1"), None);
123    }
124
125    #[test]
126    fn take_fresh_returns_unexpired_and_consumes() {
127        let cache = TtlCache::new(Duration::from_secs(60));
128        cache.insert("key1", 42);
129        assert_eq!(cache.take_fresh(&"key1"), Some(42));
130        // Single use: gone after the first take.
131        assert_eq!(cache.take_fresh(&"key1"), None);
132        assert_eq!(cache.get(&"key1"), None);
133    }
134
135    #[test]
136    fn take_fresh_drops_expired_entries() {
137        let cache = TtlCache::new(Duration::from_millis(1));
138        cache.insert("key1", 42);
139        std::thread::sleep(Duration::from_millis(5));
140        assert_eq!(
141            cache.take_fresh(&"key1"),
142            None,
143            "expired entries must not be handed out by take_fresh"
144        );
145        assert_eq!(cache.len(), 0, "expired entry must be consumed");
146    }
147
148    #[test]
149    fn cleanup_removes_expired() {
150        let cache = TtlCache::new(Duration::from_millis(1));
151        cache.insert("key1", 1);
152        cache.insert("key2", 2);
153        std::thread::sleep(Duration::from_millis(5));
154        cache.cleanup();
155        assert_eq!(cache.len(), 0);
156    }
157
158    #[test]
159    fn len_and_is_empty() {
160        let cache = TtlCache::new(Duration::from_secs(60));
161        assert!(cache.is_empty());
162        cache.insert("key1", 42);
163        assert_eq!(cache.len(), 1);
164        assert!(!cache.is_empty());
165    }
166}
shared_state/ttl.rs

shared_state/
ttl.rs
