pikadick/util/
timed_cache.rs

1use dashmap::DashMap;
2use parking_lot::Mutex;
3use rand::seq::IteratorRandom;
4use std::{
5    borrow::Borrow,
6    hash::Hash,
7    sync::Arc,
8    time::{
9        Duration,
10        Instant,
11    },
12};
13
14/// 10 minutes
15const DEFAULT_EXPIRE_TIME: Duration = Duration::from_secs(10 * 60);
16
17/// A cache with entries that "expire" after a per-cache time limit
18pub struct TimedCache<K, V>(Arc<TimedCacheInner<K, V>>);
19
20struct TimedCacheInner<K, V> {
21    cache: DashMap<K, Arc<TimedCacheEntry<V>>>,
22    last_trim: Mutex<Instant>,
23
24    trim_time: Duration,
25    expiry_time: Duration,
26}
27
28impl<K, V> TimedCache<K, V>
29where
30    K: Eq + Hash + 'static,
31    V: 'static,
32{
33    /// Create a cache with timed entries with a default expire time
34    pub fn new() -> Self {
35        TimedCache(Arc::new(TimedCacheInner {
36            cache: DashMap::new(),
37            last_trim: Mutex::new(Instant::now()),
38
39            trim_time: DEFAULT_EXPIRE_TIME,
40            expiry_time: DEFAULT_EXPIRE_TIME,
41        }))
42    }
43
44    /// Get a value if fresh, or None if it doesn't exist or is expired
45    pub fn get_if_fresh<Q>(&self, key: &Q) -> Option<Arc<TimedCacheEntry<V>>>
46    where
47        K: Borrow<Q>,
48        Q: Hash + Eq + ?Sized,
49    {
50        self.0.cache.get(key).and_then(|entry| {
51            if entry.is_fresh(self.0.expiry_time) {
52                Some(entry.value().clone())
53            } else {
54                None
55            }
56        })
57    }
58
59    /// Get a random fresh value
60    pub fn get_random_if_fresh(&self) -> Option<Arc<TimedCacheEntry<V>>> {
61        self.0
62            .cache
63            .iter()
64            .filter(|entry| entry.is_fresh(self.0.expiry_time))
65            .choose(&mut rand::thread_rng())
66            .map(|v| v.value().clone())
67    }
68
69    /// Insert a K/V
70    pub fn insert(&self, key: K, value: V) {
71        self.0.cache.insert(
72            key,
73            Arc::new(TimedCacheEntry {
74                data: value,
75                last_update: Instant::now(),
76            }),
77        );
78    }
79
80    /// Insert a K/V and return the data for the newly inserted value
81    pub fn insert_and_get(&self, key: K, value: V) -> Arc<TimedCacheEntry<V>> {
82        let data = Arc::new(TimedCacheEntry {
83            data: value,
84            last_update: Instant::now(),
85        });
86        self.0.cache.insert(key, data.clone());
87        data
88    }
89
90    /// Trims expired entries
91    pub fn trim(&self) -> bool {
92        let mut last_trim = self.0.last_trim.lock();
93        if Instant::now().duration_since(*last_trim) > self.0.trim_time {
94            *last_trim = Instant::now();
95            drop(last_trim);
96            self.force_trim();
97
98            true
99        } else {
100            false
101        }
102    }
103
104    /// Trims expired entries, ignoring last trim time.
105    pub fn force_trim(&self) {
106        let expiry_time = self.0.expiry_time;
107        self.0.cache.retain(|_, v| !v.is_fresh(expiry_time));
108    }
109
110    /// Gets the number of entries. Includes expired entries.
111    pub fn len(&self) -> usize {
112        self.0.cache.len()
113    }
114
115    /// Checks if cache is empty. Included expired entries.
116    pub fn is_empty(&self) -> bool {
117        self.0.cache.is_empty()
118    }
119}
120
121impl<K, V> Default for TimedCache<K, V>
122where
123    K: Eq + Hash + 'static,
124    V: 'static,
125{
126    fn default() -> Self {
127        Self::new()
128    }
129}
130
131impl<K, V> Clone for TimedCache<K, V>
132where
133    K: Eq + Hash + 'static,
134    V: 'static,
135{
136    fn clone(&self) -> Self {
137        TimedCache(self.0.clone())
138    }
139}
140
141impl<K, V> std::fmt::Debug for TimedCache<K, V>
142where
143    K: Eq + std::fmt::Debug + Hash,
144    V: std::fmt::Debug,
145{
146    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
147        f.debug_struct("TimedCache")
148            .field("cache", &self.0.cache)
149            .finish()
150    }
151}
152
153#[derive(Debug)]
154pub struct TimedCacheEntry<T> {
155    data: T,
156    last_update: Instant,
157}
158
159impl<T> TimedCacheEntry<T> {
160    /// time is expire time
161    pub fn is_fresh(&self, time: Duration) -> bool {
162        self.last_update.elapsed() < time
163    }
164
165    /// Get data ref
166    pub fn data(&self) -> &T {
167        &self.data
168    }
169}