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#[derive(Debug)]
20pub(super) enum Status {
21 NotFound,
22 Found,
23 Expired,
24}
25
26#[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 #[must_use]
45 pub fn with_lifespan(seconds: u64) -> TimedCache<K, V> {
46 Self::with_lifespan_and_refresh(seconds, false)
47 }
48
49 #[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 #[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 #[must_use]
79 pub fn refresh(&self) -> bool {
80 self.refresh
81 }
82
83 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 #[must_use]
94 pub fn get_store(&self) -> &HashMap<K, (Instant, V)> {
95 &self.store
96 }
97
98 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)]
359mod 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 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 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 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}