cached/stores/
timed.rs

1use std::cmp::Eq;
2use std::hash::Hash;
3use web_time::Instant;
4
5#[cfg(feature = "ahash")]
6use hashbrown::{hash_map::Entry, HashMap};
7
8#[cfg(not(feature = "ahash"))]
9use std::collections::{hash_map::Entry, HashMap};
10
11#[cfg(feature = "async")]
12use {super::CachedAsync, async_trait::async_trait, futures::Future};
13
14use crate::CloneCached;
15
16use super::Cached;
17
18/// Enum used for defining the status of time-cached values
19#[derive(Debug)]
20pub(super) enum Status {
21    NotFound,
22    Found,
23    Expired,
24}
25
26/// Cache store bound by time
27///
28/// Values are timestamped when inserted and are
29/// evicted if expired at time of retrieval.
30///
31/// Note: This cache is in-memory only
32#[derive(Clone, Debug)]
33pub struct TimedCache<K, V> {
34    pub(super) store: HashMap<K, (Instant, V)>,
35    pub(super) seconds: u64,
36    pub(super) hits: u64,
37    pub(super) misses: u64,
38    pub(super) initial_capacity: Option<usize>,
39    pub(super) refresh: bool,
40}
41
42impl<K: Hash + Eq, V> TimedCache<K, V> {
43    /// Creates a new `TimedCache` with a specified lifespan
44    #[must_use]
45    pub fn with_lifespan(seconds: u64) -> TimedCache<K, V> {
46        Self::with_lifespan_and_refresh(seconds, false)
47    }
48
49    /// Creates a new `TimedCache` with a specified lifespan and
50    /// cache-store with the specified pre-allocated capacity
51    #[must_use]
52    pub fn with_lifespan_and_capacity(seconds: u64, size: usize) -> TimedCache<K, V> {
53        TimedCache {
54            store: Self::new_store(Some(size)),
55            seconds,
56            hits: 0,
57            misses: 0,
58            initial_capacity: Some(size),
59            refresh: false,
60        }
61    }
62
63    /// Creates a new `TimedCache` with a specified lifespan which
64    /// refreshes the ttl when the entry is retrieved
65    #[must_use]
66    pub fn with_lifespan_and_refresh(seconds: u64, refresh: bool) -> TimedCache<K, V> {
67        TimedCache {
68            store: Self::new_store(None),
69            seconds,
70            hits: 0,
71            misses: 0,
72            initial_capacity: None,
73            refresh,
74        }
75    }
76
77    /// Returns if the lifetime is refreshed when the value is retrieved
78    #[must_use]
79    pub fn refresh(&self) -> bool {
80        self.refresh
81    }
82
83    /// Sets if the lifetime is refreshed when the value is retrieved
84    pub fn set_refresh(&mut self, refresh: bool) {
85        self.refresh = refresh;
86    }
87
88    fn new_store(capacity: Option<usize>) -> HashMap<K, (Instant, V)> {
89        capacity.map_or_else(HashMap::new, HashMap::with_capacity)
90    }
91
92    /// Returns a reference to the cache's `store`
93    #[must_use]
94    pub fn get_store(&self) -> &HashMap<K, (Instant, V)> {
95        &self.store
96    }
97
98    /// Remove any expired values from the cache
99    pub fn flush(&mut self) {
100        let seconds = self.seconds;
101        self.store
102            .retain(|_, (instant, _)| instant.elapsed().as_secs() < seconds);
103    }
104
105    fn status<Q>(&mut self, key: &Q) -> Status
106    where
107        K: std::borrow::Borrow<Q>,
108        Q: std::hash::Hash + Eq + ?Sized,
109    {
110        let mut val = self.store.get_mut(key);
111        if let Some(&mut (instant, _)) = val.as_mut() {
112            if instant.elapsed().as_secs() < self.seconds {
113                if self.refresh {
114                    *instant = Instant::now();
115                }
116                Status::Found
117            } else {
118                Status::Expired
119            }
120        } else {
121            Status::NotFound
122        }
123    }
124}
125
126impl<K: Hash + Eq, V> Cached<K, V> for TimedCache<K, V> {
127    fn cache_get<Q>(&mut self, key: &Q) -> Option<&V>
128    where
129        K: std::borrow::Borrow<Q>,
130        Q: std::hash::Hash + Eq + ?Sized,
131    {
132        match self.status(key) {
133            Status::NotFound => {
134                self.misses += 1;
135                None
136            }
137            Status::Found => {
138                self.hits += 1;
139                self.store.get(key).map(|stamped| &stamped.1)
140            }
141            Status::Expired => {
142                self.misses += 1;
143                self.store.remove(key).unwrap();
144                None
145            }
146        }
147    }
148
149    fn cache_get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
150    where
151        K: std::borrow::Borrow<Q>,
152        Q: std::hash::Hash + Eq + ?Sized,
153    {
154        match self.status(key) {
155            Status::NotFound => {
156                self.misses += 1;
157                None
158            }
159            Status::Found => {
160                self.hits += 1;
161                self.store.get_mut(key).map(|stamped| &mut stamped.1)
162            }
163            Status::Expired => {
164                self.misses += 1;
165                self.store.remove(key).unwrap();
166                None
167            }
168        }
169    }
170
171    fn cache_get_or_set_with<F: FnOnce() -> V>(&mut self, key: K, f: F) -> &mut V {
172        match self.store.entry(key) {
173            Entry::Occupied(mut occupied) => {
174                if occupied.get().0.elapsed().as_secs() < self.seconds {
175                    if self.refresh {
176                        occupied.get_mut().0 = Instant::now();
177                    }
178                    self.hits += 1;
179                } else {
180                    self.misses += 1;
181                    let val = f();
182                    occupied.insert((Instant::now(), val));
183                }
184                &mut occupied.into_mut().1
185            }
186            Entry::Vacant(vacant) => {
187                self.misses += 1;
188                let val = f();
189                &mut vacant.insert((Instant::now(), val)).1
190            }
191        }
192    }
193
194    fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, E>(
195        &mut self,
196        key: K,
197        f: F,
198    ) -> Result<&mut V, E> {
199        match self.store.entry(key) {
200            Entry::Occupied(mut occupied) => {
201                if occupied.get().0.elapsed().as_secs() < self.seconds {
202                    if self.refresh {
203                        occupied.get_mut().0 = Instant::now();
204                    }
205                    self.hits += 1;
206                } else {
207                    self.misses += 1;
208                    let val = f()?;
209                    occupied.insert((Instant::now(), val));
210                }
211                Ok(&mut occupied.into_mut().1)
212            }
213            Entry::Vacant(vacant) => {
214                self.misses += 1;
215                let val = f()?;
216                Ok(&mut vacant.insert((Instant::now(), val)).1)
217            }
218        }
219    }
220
221    fn cache_set(&mut self, key: K, val: V) -> Option<V> {
222        let stamped = (Instant::now(), val);
223        self.store.insert(key, stamped).and_then(|(instant, v)| {
224            if instant.elapsed().as_secs() < self.seconds {
225                Some(v)
226            } else {
227                None
228            }
229        })
230    }
231    fn cache_remove<Q>(&mut self, k: &Q) -> Option<V>
232    where
233        K: std::borrow::Borrow<Q>,
234        Q: std::hash::Hash + Eq + ?Sized,
235    {
236        self.store.remove(k).and_then(|(instant, v)| {
237            if instant.elapsed().as_secs() < self.seconds {
238                Some(v)
239            } else {
240                None
241            }
242        })
243    }
244    fn cache_clear(&mut self) {
245        self.store.clear();
246    }
247    fn cache_reset_metrics(&mut self) {
248        self.misses = 0;
249        self.hits = 0;
250    }
251    fn cache_reset(&mut self) {
252        self.store = Self::new_store(self.initial_capacity);
253    }
254    fn cache_size(&self) -> usize {
255        self.store.len()
256    }
257    fn cache_hits(&self) -> Option<u64> {
258        Some(self.hits)
259    }
260    fn cache_misses(&self) -> Option<u64> {
261        Some(self.misses)
262    }
263    fn cache_lifespan(&self) -> Option<u64> {
264        Some(self.seconds)
265    }
266
267    fn cache_set_lifespan(&mut self, seconds: u64) -> Option<u64> {
268        let old = self.seconds;
269        self.seconds = seconds;
270        Some(old)
271    }
272}
273
274impl<K: Hash + Eq + Clone, V: Clone> CloneCached<K, V> for TimedCache<K, V> {
275    fn cache_get_expired<Q>(&mut self, k: &Q) -> (Option<V>, bool)
276    where
277        K: std::borrow::Borrow<Q>,
278        Q: std::hash::Hash + Eq + ?Sized,
279    {
280        match self.status(k) {
281            Status::NotFound => {
282                self.misses += 1;
283                (None, false)
284            }
285            Status::Found => {
286                self.hits += 1;
287                (self.store.get(k).map(|stamped| &stamped.1).cloned(), false)
288            }
289            Status::Expired => {
290                self.misses += 1;
291                (self.store.remove(k).map(|stamped| stamped.1), true)
292            }
293        }
294    }
295}
296
297#[cfg(feature = "async")]
298#[async_trait]
299impl<K, V> CachedAsync<K, V> for TimedCache<K, V>
300where
301    K: Hash + Eq + Clone + Send,
302{
303    async fn get_or_set_with<F, Fut>(&mut self, k: K, f: F) -> &mut V
304    where
305        V: Send,
306        F: FnOnce() -> Fut + Send,
307        Fut: Future<Output = V> + Send,
308    {
309        match self.store.entry(k) {
310            Entry::Occupied(mut occupied) => {
311                if occupied.get().0.elapsed().as_secs() < self.seconds {
312                    if self.refresh {
313                        occupied.get_mut().0 = Instant::now();
314                    }
315                    self.hits += 1;
316                } else {
317                    self.misses += 1;
318                    occupied.insert((Instant::now(), f().await));
319                }
320                &mut occupied.into_mut().1
321            }
322            Entry::Vacant(vacant) => {
323                self.misses += 1;
324                &mut vacant.insert((Instant::now(), f().await)).1
325            }
326        }
327    }
328
329    async fn try_get_or_set_with<F, Fut, E>(&mut self, k: K, f: F) -> Result<&mut V, E>
330    where
331        V: Send,
332        F: FnOnce() -> Fut + Send,
333        Fut: Future<Output = Result<V, E>> + Send,
334    {
335        let v = match self.store.entry(k) {
336            Entry::Occupied(mut occupied) => {
337                if occupied.get().0.elapsed().as_secs() < self.seconds {
338                    if self.refresh {
339                        occupied.get_mut().0 = Instant::now();
340                    }
341                    self.hits += 1;
342                } else {
343                    self.misses += 1;
344                    occupied.insert((Instant::now(), f().await?));
345                }
346                &mut occupied.into_mut().1
347            }
348            Entry::Vacant(vacant) => {
349                self.misses += 1;
350                &mut vacant.insert((Instant::now(), f().await?)).1
351            }
352        };
353
354        Ok(v)
355    }
356}
357
358#[cfg(test)]
359/// Cache store tests
360mod tests {
361    use std::{thread::sleep, time::Duration};
362
363    use super::*;
364
365    #[test]
366    fn timed_cache() {
367        let mut c = TimedCache::with_lifespan(2);
368        assert!(c.cache_get(&1).is_none());
369        let misses = c.cache_misses().unwrap();
370        assert_eq!(1, misses);
371
372        assert_eq!(c.cache_set(1, 100), None);
373        assert!(c.cache_get(&1).is_some());
374        let hits = c.cache_hits().unwrap();
375        let misses = c.cache_misses().unwrap();
376        assert_eq!(1, hits);
377        assert_eq!(1, misses);
378
379        sleep(Duration::new(2, 0));
380        assert!(c.cache_get(&1).is_none());
381        let misses = c.cache_misses().unwrap();
382        assert_eq!(2, misses);
383
384        let old = c.cache_set_lifespan(1).unwrap();
385        assert_eq!(2, old);
386        assert_eq!(c.cache_set(1, 100), None);
387        assert!(c.cache_get(&1).is_some());
388        let hits = c.cache_hits().unwrap();
389        let misses = c.cache_misses().unwrap();
390        assert_eq!(2, hits);
391        assert_eq!(2, misses);
392
393        sleep(Duration::new(1, 0));
394        assert!(c.cache_get(&1).is_none());
395        let misses = c.cache_misses().unwrap();
396        assert_eq!(3, misses);
397    }
398
399    #[test]
400    fn timed_cache_refresh() {
401        let mut c = TimedCache::with_lifespan_and_refresh(2, true);
402        assert!(c.refresh());
403        assert_eq!(c.cache_get(&1), None);
404        let misses = c.cache_misses().unwrap();
405        assert_eq!(1, misses);
406
407        assert_eq!(c.cache_set(1, 100), None);
408        assert_eq!(c.cache_get(&1), Some(&100));
409        let hits = c.cache_hits().unwrap();
410        let misses = c.cache_misses().unwrap();
411        assert_eq!(1, hits);
412        assert_eq!(1, misses);
413
414        assert_eq!(c.cache_set(2, 200), None);
415        assert_eq!(c.cache_get(&2), Some(&200));
416        sleep(Duration::new(1, 0));
417        assert_eq!(c.cache_get(&1), Some(&100));
418        sleep(Duration::new(1, 0));
419        assert_eq!(c.cache_get(&1), Some(&100));
420        assert_eq!(c.cache_get(&2), None);
421    }
422
423    #[test]
424    fn clear() {
425        let mut c = TimedCache::with_lifespan(3600);
426
427        assert_eq!(c.cache_set(1, 100), None);
428        assert_eq!(c.cache_set(2, 200), None);
429        assert_eq!(c.cache_set(3, 300), None);
430        c.cache_clear();
431
432        assert_eq!(0, c.cache_size());
433    }
434
435    #[test]
436    fn reset() {
437        let mut c = TimedCache::with_lifespan(100);
438        assert_eq!(c.cache_set(1, 100), None);
439        assert_eq!(c.cache_set(2, 200), None);
440        assert_eq!(c.cache_set(3, 300), None);
441        assert!(3 <= c.store.capacity());
442
443        c.cache_reset();
444
445        assert_eq!(0, c.store.capacity());
446
447        let init_capacity = 1;
448        let mut c = TimedCache::with_lifespan_and_capacity(100, init_capacity);
449        assert_eq!(c.cache_set(1, 100), None);
450        assert_eq!(c.cache_set(2, 200), None);
451        assert_eq!(c.cache_set(3, 300), None);
452        assert!(3 <= c.store.capacity());
453
454        c.cache_reset();
455
456        assert!(init_capacity <= c.store.capacity());
457    }
458
459    #[test]
460    fn remove() {
461        let mut c = TimedCache::with_lifespan(3600);
462
463        assert_eq!(c.cache_set(1, 100), None);
464        assert_eq!(c.cache_set(2, 200), None);
465        assert_eq!(c.cache_set(3, 300), None);
466
467        assert_eq!(Some(100), c.cache_remove(&1));
468        assert_eq!(2, c.cache_size());
469    }
470
471    #[test]
472    fn remove_expired() {
473        let mut c = TimedCache::with_lifespan(1);
474
475        assert_eq!(c.cache_set(1, 100), None);
476        assert_eq!(c.cache_set(1, 200), Some(100));
477        assert_eq!(c.cache_size(), 1);
478
479        std::thread::sleep(std::time::Duration::from_secs(1));
480        assert_eq!(None, c.cache_remove(&1));
481        assert_eq!(0, c.cache_size());
482    }
483
484    #[test]
485    fn insert_expired() {
486        let mut c = TimedCache::with_lifespan(1);
487
488        assert_eq!(c.cache_set(1, 100), None);
489        assert_eq!(c.cache_set(1, 200), Some(100));
490        assert_eq!(c.cache_size(), 1);
491
492        std::thread::sleep(std::time::Duration::from_secs(1));
493        assert_eq!(1, c.cache_size());
494        assert_eq!(None, c.cache_set(1, 300));
495        assert_eq!(1, c.cache_size());
496    }
497
498    #[test]
499    fn get_expired() {
500        let mut c = TimedCache::with_lifespan(1);
501
502        assert_eq!(c.cache_set(1, 100), None);
503        assert_eq!(c.cache_set(1, 200), Some(100));
504        assert_eq!(c.cache_size(), 1);
505
506        std::thread::sleep(std::time::Duration::from_secs(1));
507        // still around until we try to get
508        assert_eq!(1, c.cache_size());
509        assert_eq!(None, c.cache_get(&1));
510        assert_eq!(0, c.cache_size());
511    }
512
513    #[test]
514    fn get_mut_expired() {
515        let mut c = TimedCache::with_lifespan(1);
516
517        assert_eq!(c.cache_set(1, 100), None);
518        assert_eq!(c.cache_set(1, 200), Some(100));
519        assert_eq!(c.cache_size(), 1);
520
521        std::thread::sleep(std::time::Duration::from_secs(1));
522        // still around until we try to get
523        assert_eq!(1, c.cache_size());
524        assert_eq!(None, c.cache_get_mut(&1));
525        assert_eq!(0, c.cache_size());
526    }
527
528    #[test]
529    fn flush_expired() {
530        let mut c = TimedCache::with_lifespan(1);
531
532        assert_eq!(c.cache_set(1, 100), None);
533        assert_eq!(c.cache_set(1, 200), Some(100));
534        assert_eq!(c.cache_size(), 1);
535
536        std::thread::sleep(std::time::Duration::from_secs(1));
537        // still around until we flush
538        assert_eq!(1, c.cache_size());
539        c.flush();
540        assert_eq!(0, c.cache_size());
541    }
542
543    #[test]
544    fn get_or_set_with() {
545        let mut c = TimedCache::with_lifespan(2);
546
547        assert_eq!(c.cache_get_or_set_with(0, || 0), &0);
548        assert_eq!(c.cache_get_or_set_with(1, || 1), &1);
549        assert_eq!(c.cache_get_or_set_with(2, || 2), &2);
550        assert_eq!(c.cache_get_or_set_with(3, || 3), &3);
551        assert_eq!(c.cache_get_or_set_with(4, || 4), &4);
552        assert_eq!(c.cache_get_or_set_with(5, || 5), &5);
553
554        assert_eq!(c.cache_misses(), Some(6));
555
556        assert_eq!(c.cache_get_or_set_with(0, || 0), &0);
557
558        assert_eq!(c.cache_misses(), Some(6));
559
560        assert_eq!(c.cache_get_or_set_with(0, || 42), &0);
561
562        assert_eq!(c.cache_misses(), Some(6));
563
564        sleep(Duration::new(2, 0));
565
566        assert_eq!(c.cache_get_or_set_with(1, || 42), &42);
567
568        assert_eq!(c.cache_misses(), Some(7));
569
570        c.cache_reset();
571        fn _try_get(n: usize) -> Result<usize, String> {
572            if n < 10 {
573                Ok(n)
574            } else {
575                Err("dead".to_string())
576            }
577        }
578
579        let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(10));
580        assert!(res.is_err());
581
582        let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(1));
583        assert_eq!(res.unwrap(), &1);
584        let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(5));
585        assert_eq!(res.unwrap(), &1);
586        sleep(Duration::new(2, 0));
587        let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(5));
588        assert_eq!(res.unwrap(), &5);
589    }
590}