1use std::borrow::Borrow;
2use std::cmp::Ord;
3use std::collections::btree_map;
4use std::collections::hash_map;
5use std::collections::BTreeMap;
6use std::collections::HashMap;
7use std::fmt::Debug;
8use std::hash::{BuildHasher, Hash};
9use std::time::{Duration, Instant};
10
11pub trait CacheMapExt<K, V> {
12 fn get_with_timeout<Q>(&self, k: &Q) -> Option<&V>
13 where
14 K: Borrow<Q> + Ord,
15 Q: Hash + Eq + Ord + ?Sized;
16
17 fn get_with_timeout_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
18 where
19 K: Borrow<Q> + Ord,
20 Q: Hash + Eq + Ord + ?Sized;
21
22 fn insert_with_timeout(&mut self, k: K, v: V, timeout: Option<Duration>) -> Option<V>;
23
24 fn remove_expired_values(&mut self);
25}
26
27pub trait EntryExt<'a, K, V> {
28 fn or_insert_with_timeout(self, default: V, timeout: Option<Duration>) -> &'a mut V;
29 fn or_insert_with_timeout_f<F: FnOnce() -> V>(
30 self,
31 default: F,
32 timeout: Option<Duration>,
33 ) -> &'a mut V;
34 fn or_insert_with_timeout_key_f<F: FnOnce(&K) -> V>(
35 self,
36 default: F,
37 timeout: Option<Duration>,
38 ) -> &'a mut V;
39 fn and_modify_with_timeout<F>(self, f: F) -> Self
40 where
41 F: FnOnce(&mut V);
42}
43
44impl<K, V, S> CacheMapExt<K, V> for HashMap<K, TimedValue<V>, S>
45where
46 K: Eq + Hash,
47 S: BuildHasher,
48{
49 fn get_with_timeout<Q>(&self, k: &Q) -> Option<&V>
50 where
51 K: Borrow<Q>,
52 Q: Hash + Eq + ?Sized,
53 {
54 self.get(k).and_then(|tv| {
55 if tv.is_expired() {
56 None
57 } else {
58 Some(tv.value())
59 }
60 })
61 }
62
63 fn get_with_timeout_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
64 where
65 K: Borrow<Q>,
66 Q: Hash + Eq + ?Sized,
67 {
68 self.get_mut(k).and_then(|tv| {
69 if tv.is_expired() {
70 None
71 } else {
72 Some(tv.value_mut())
73 }
74 })
75 }
76
77 fn insert_with_timeout(&mut self, k: K, v: V, timeout: Option<Duration>) -> Option<V> {
78 self.insert(k, TimedValue::new(v, timeout))
79 .map(|tv| tv.into_value())
80 }
81
82 fn remove_expired_values(&mut self) {
83 self.retain(|_, tv| !tv.is_expired());
84 }
85}
86
87impl<'a, K, V> EntryExt<'a, K, V> for hash_map::Entry<'a, K, TimedValue<V>>
88where
89 K: Eq + Hash,
90{
91 fn or_insert_with_timeout(self, default: V, timeout: Option<Duration>) -> &'a mut V {
92 match self {
93 hash_map::Entry::Occupied(entry) => {
94 let v = entry.into_mut();
95 if v.is_expired() {
96 *v = TimedValue::new(default, timeout);
97 }
98 v.value_mut()
99 }
100 hash_map::Entry::Vacant(entry) => {
101 entry.insert(TimedValue::new(default, timeout)).value_mut()
102 }
103 }
104 }
105
106 fn or_insert_with_timeout_f<F: FnOnce() -> V>(
107 self,
108 default: F,
109 timeout: Option<Duration>,
110 ) -> &'a mut V {
111 match self {
112 hash_map::Entry::Occupied(entry) => {
113 let v = entry.into_mut();
114 if v.is_expired() {
115 *v = TimedValue::new(default(), timeout);
116 }
117 v.value_mut()
118 }
119 hash_map::Entry::Vacant(entry) => entry
120 .insert(TimedValue::new(default(), timeout))
121 .value_mut(),
122 }
123 }
124
125 fn or_insert_with_timeout_key_f<F: FnOnce(&K) -> V>(
126 self,
127 default: F,
128 timeout: Option<Duration>,
129 ) -> &'a mut V {
130 match self {
131 hash_map::Entry::Occupied(entry) => {
132 let value = if entry.get().is_expired() {
133 Some(default(entry.key()))
134 } else {
135 None
136 };
137 let v = entry.into_mut();
138 if let Some(value) = value {
139 *v = TimedValue::new(value, timeout);
140 }
141 v.value_mut()
142 }
143 hash_map::Entry::Vacant(entry) => {
144 let value = default(entry.key());
145 entry.insert(TimedValue::new(value, timeout)).value_mut()
146 }
147 }
148 }
149
150 fn and_modify_with_timeout<F>(self, f: F) -> Self
151 where
152 F: FnOnce(&mut V),
153 {
154 self.and_modify(|v| f(v.value_mut()))
155 }
156}
157
158impl<K: Ord, V> CacheMapExt<K, V> for BTreeMap<K, TimedValue<V>> {
159 fn get_with_timeout<Q>(&self, k: &Q) -> Option<&V>
160 where
161 K: Borrow<Q>,
162 Q: Ord + ?Sized,
163 {
164 self.get(k).and_then(|tv| {
165 if tv.is_expired() {
166 None
167 } else {
168 Some(tv.value())
169 }
170 })
171 }
172
173 fn get_with_timeout_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
174 where
175 K: Borrow<Q>,
176 Q: Ord + ?Sized,
177 {
178 self.get_mut(k).and_then(|tv| {
179 if tv.is_expired() {
180 None
181 } else {
182 Some(tv.value_mut())
183 }
184 })
185 }
186
187 fn insert_with_timeout(&mut self, k: K, v: V, timeout: Option<Duration>) -> Option<V> {
188 self.insert(k, TimedValue::new(v, timeout))
189 .map(|tv| tv.into_value())
190 }
191
192 fn remove_expired_values(&mut self) {
193 self.retain(|_, tv| !tv.is_expired());
194 }
195}
196
197impl<'a, K, V> EntryExt<'a, K, V> for btree_map::Entry<'a, K, TimedValue<V>>
198where
199 K: Ord,
200{
201 fn or_insert_with_timeout(self, default: V, timeout: Option<Duration>) -> &'a mut V {
202 match self {
203 btree_map::Entry::Occupied(entry) => {
204 let v = entry.into_mut();
205 if v.is_expired() {
206 *v = TimedValue::new(default, timeout);
207 }
208 v.value_mut()
209 }
210 btree_map::Entry::Vacant(entry) => {
211 entry.insert(TimedValue::new(default, timeout)).value_mut()
212 }
213 }
214 }
215
216 fn or_insert_with_timeout_f<F: FnOnce() -> V>(
217 self,
218 default: F,
219 timeout: Option<Duration>,
220 ) -> &'a mut V {
221 match self {
222 btree_map::Entry::Occupied(entry) => {
223 let v = entry.into_mut();
224 if v.is_expired() {
225 *v = TimedValue::new(default(), timeout);
226 }
227 v.value_mut()
228 }
229 btree_map::Entry::Vacant(entry) => entry
230 .insert(TimedValue::new(default(), timeout))
231 .value_mut(),
232 }
233 }
234
235 fn or_insert_with_timeout_key_f<F: FnOnce(&K) -> V>(
236 self,
237 default: F,
238 timeout: Option<Duration>,
239 ) -> &'a mut V {
240 match self {
241 btree_map::Entry::Occupied(entry) => {
242 let value = if entry.get().is_expired() {
243 Some(default(entry.key()))
244 } else {
245 None
246 };
247 let v = entry.into_mut();
248 if let Some(value) = value {
249 *v = TimedValue::new(value, timeout);
250 }
251 v.value_mut()
252 }
253 btree_map::Entry::Vacant(entry) => {
254 let value = default(entry.key());
255 entry.insert(TimedValue::new(value, timeout)).value_mut()
256 }
257 }
258 }
259
260 fn and_modify_with_timeout<F>(self, f: F) -> Self
261 where
262 F: FnOnce(&mut V),
263 {
264 self.and_modify(|v| f(v.value_mut()))
265 }
266}
267
268impl<K: Ord + Clone, V> CacheMapExt<K, V> for dequemap::DequeBTreeMap<K, TimedValue<V>> {
269 fn get_with_timeout<Q>(&self, k: &Q) -> Option<&V>
270 where
271 K: Borrow<Q>,
272 Q: Ord + ?Sized,
273 {
274 self.get(k).and_then(|tv| {
275 if tv.is_expired() {
276 None
277 } else {
278 Some(tv.value())
279 }
280 })
281 }
282
283 fn get_with_timeout_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
284 where
285 K: Borrow<Q>,
286 Q: Ord + ?Sized,
287 {
288 self.get_mut(k).and_then(|tv| {
289 if tv.is_expired() {
290 None
291 } else {
292 Some(tv.value_mut())
293 }
294 })
295 }
296
297 fn insert_with_timeout(&mut self, k: K, v: V, timeout: Option<Duration>) -> Option<V> {
298 self.insert(k, TimedValue::new(v, timeout))
299 .map(|tv| tv.into_value())
300 }
301
302 fn remove_expired_values(&mut self) {
303 self.retain(|_, tv| !tv.is_expired());
304 }
305}
306
307impl<'a, K, V> EntryExt<'a, K, V> for dequemap::btreemap::Entry<'a, K, TimedValue<V>>
308where
309 K: Ord + Clone,
310{
311 fn or_insert_with_timeout(self, default: V, timeout: Option<Duration>) -> &'a mut V {
312 match self {
313 dequemap::btreemap::Entry::Occupied(entry) => {
314 let v = entry.into_mut();
315 if v.is_expired() {
316 *v = TimedValue::new(default, timeout);
317 }
318 v.value_mut()
319 }
320 dequemap::btreemap::Entry::Vacant(entry) => {
321 entry.insert(TimedValue::new(default, timeout)).value_mut()
322 }
323 }
324 }
325
326 fn or_insert_with_timeout_f<F: FnOnce() -> V>(
327 self,
328 default: F,
329 timeout: Option<Duration>,
330 ) -> &'a mut V {
331 match self {
332 dequemap::btreemap::Entry::Occupied(entry) => {
333 let v = entry.into_mut();
334 if v.is_expired() {
335 *v = TimedValue::new(default(), timeout);
336 }
337 v.value_mut()
338 }
339 dequemap::btreemap::Entry::Vacant(entry) => entry
340 .insert(TimedValue::new(default(), timeout))
341 .value_mut(),
342 }
343 }
344
345 fn or_insert_with_timeout_key_f<F: FnOnce(&K) -> V>(
346 self,
347 default: F,
348 timeout: Option<Duration>,
349 ) -> &'a mut V {
350 match self {
351 dequemap::btreemap::Entry::Occupied(entry) => {
352 let value = if entry.get().is_expired() {
353 Some(default(entry.key()))
354 } else {
355 None
356 };
357 let v = entry.into_mut();
358 if let Some(value) = value {
359 *v = TimedValue::new(value, timeout);
360 }
361 v.value_mut()
362 }
363 dequemap::btreemap::Entry::Vacant(entry) => {
364 let value = default(entry.key());
365 entry.insert(TimedValue::new(value, timeout)).value_mut()
366 }
367 }
368 }
369
370 fn and_modify_with_timeout<F>(self, f: F) -> Self
371 where
372 F: FnOnce(&mut V),
373 {
374 self.and_modify(|v| f(v.value_mut()))
375 }
376}
377
378impl<K, V, S> CacheMapExt<K, V> for dequemap::DequeHashMap<K, TimedValue<V>, S>
380where
381 K: Hash + Eq + Ord + Clone,
382 S: BuildHasher,
383{
384 fn get_with_timeout<Q>(&self, k: &Q) -> Option<&V>
385 where
386 K: Borrow<Q>,
387 Q: Hash + Eq + ?Sized,
388 {
389 self.get(k).and_then(|tv| {
390 if tv.is_expired() {
391 None
392 } else {
393 Some(tv.value())
394 }
395 })
396 }
397
398 fn get_with_timeout_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
399 where
400 K: Borrow<Q>,
401 Q: Hash + Eq + ?Sized,
402 {
403 self.get_mut(k).and_then(|tv| {
404 if tv.is_expired() {
405 None
406 } else {
407 Some(tv.value_mut())
408 }
409 })
410 }
411
412 fn insert_with_timeout(&mut self, k: K, v: V, timeout: Option<Duration>) -> Option<V> {
413 self.insert(k, TimedValue::new(v, timeout))
414 .map(|tv| tv.into_value())
415 }
416
417 fn remove_expired_values(&mut self) {
418 self.retain(|_, tv| !tv.is_expired());
419 }
420}
421
422impl<'a, K, V, S> EntryExt<'a, K, V> for dequemap::hashmap::Entry<'a, K, TimedValue<V>, S>
423where
424 K: Eq + Hash + Clone,
425 S: BuildHasher,
426{
427 fn or_insert_with_timeout(self, default: V, timeout: Option<Duration>) -> &'a mut V {
428 match self {
429 dequemap::hashmap::Entry::Occupied(entry) => {
430 let v = entry.into_mut();
431 if v.is_expired() {
432 *v = TimedValue::new(default, timeout);
433 }
434 v.value_mut()
435 }
436 dequemap::hashmap::Entry::Vacant(entry) => {
437 entry.insert(TimedValue::new(default, timeout)).value_mut()
438 }
439 }
440 }
441
442 fn or_insert_with_timeout_f<F: FnOnce() -> V>(
443 self,
444 default: F,
445 timeout: Option<Duration>,
446 ) -> &'a mut V {
447 match self {
448 dequemap::hashmap::Entry::Occupied(entry) => {
449 let v = entry.into_mut();
450 if v.is_expired() {
451 *v = TimedValue::new(default(), timeout);
452 }
453 v.value_mut()
454 }
455 dequemap::hashmap::Entry::Vacant(entry) => entry
456 .insert(TimedValue::new(default(), timeout))
457 .value_mut(),
458 }
459 }
460
461 fn or_insert_with_timeout_key_f<F: FnOnce(&K) -> V>(
462 self,
463 default: F,
464 timeout: Option<Duration>,
465 ) -> &'a mut V {
466 match self {
467 dequemap::hashmap::Entry::Occupied(entry) => {
468 let value = if entry.get().is_expired() {
469 Some(default(entry.key()))
470 } else {
471 None
472 };
473 let v = entry.into_mut();
474 if let Some(value) = value {
475 *v = TimedValue::new(value, timeout);
476 }
477 v.value_mut()
478 }
479 dequemap::hashmap::Entry::Vacant(entry) => {
480 let value = default(entry.key());
481 entry.insert(TimedValue::new(value, timeout)).value_mut()
482 }
483 }
484 }
485
486 fn and_modify_with_timeout<F>(self, f: F) -> Self
487 where
488 F: FnOnce(&mut V),
489 {
490 self.and_modify(|v| f(v.value_mut()))
491 }
492}
493
494#[derive(Clone, Debug)]
495pub struct TimedValue<V>(V, Option<Instant>);
496
497impl<V> TimedValue<V> {
498 pub fn new(value: V, timeout_duration: Option<Duration>) -> Self {
499 TimedValue(value, timeout_duration.map(|t| Instant::now() + t))
500 }
501
502 pub fn value(&self) -> &V {
503 &self.0
504 }
505
506 pub fn value_mut(&mut self) -> &mut V {
507 &mut self.0
508 }
509
510 pub fn into_value(self) -> V {
511 self.0
512 }
513
514 pub fn is_expired(&self) -> bool {
515 self.1.map(|e| Instant::now() >= e).unwrap_or(false)
516 }
517}
518
519impl<V> PartialEq for TimedValue<V>
520where
521 V: PartialEq,
522{
523 fn eq(&self, other: &TimedValue<V>) -> bool {
524 self.value() == other.value()
525 }
526}
527
528#[test]
531fn test_cache_map_ext() {
532 use std::collections::hash_map::RandomState;
533
534 let mut m: HashMap<_, _, RandomState> = HashMap::default();
535 let old1 = m.insert("k1", TimedValue::new(1, None));
536 let old2 = m.insert("k2", TimedValue::new(2, Some(Duration::from_millis(50))));
537 let old3 = m.insert("k3", TimedValue::new(3, Some(Duration::from_millis(80))));
538 let old4 = m.insert("k4", TimedValue::new(4, Some(Duration::from_millis(130))));
539 let old44 = m.insert("k4", TimedValue::new(44, None));
540
541 let old6 = m.insert_with_timeout("k6", 6, Some(Duration::from_secs(150)));
542 let old7 = m.insert_with_timeout("k7", 7, None);
543
544 let old66 = m.insert_with_timeout("k6", 66, Some(Duration::from_secs(60)));
545
546 assert_eq!(old1, None);
547 assert_eq!(old2, None);
548 assert_eq!(old3, None);
549 assert_eq!(old4, None);
550 assert_eq!(old6, None);
551 assert_eq!(old7, None);
552
553 println!("old44: {:?}", old44);
554 assert_eq!(old44, Some(TimedValue::new(4, None)));
555
556 assert_eq!(old66, Some(6));
557 println!("old66: {:?}", old66);
558
559 let v6 = m.get("k6");
560 println!("v6: {:?}", v6);
561 assert_eq!(v6, Some(&TimedValue::new(66, None)));
562
563 m.get_with_timeout_mut("k6").map(|v| *v = 666);
564 let v6 = m.get_with_timeout("k6");
565 println!("v6: {:?}", v6);
566 assert_eq!(v6, Some(&666));
567
568 for i in 0..20 {
569 m.remove_expired_values();
570 println!(
571 "{} map len: {}, map k1: {:?}, k2: {:?}",
572 i,
573 m.len(),
574 m.get("k1"),
575 m.get_with_timeout("k2")
576 );
577 std::thread::sleep(std::time::Duration::from_millis(30));
579 }
580}
581
582#[test]
583fn test_btree_map_ext() {
584 let mut m: BTreeMap<_, _> = BTreeMap::default();
585
586 let old1 = m.insert_with_timeout("k1", 1, None);
587 let old2 = m.insert_with_timeout("k2", 2, Some(Duration::from_millis(50)));
588 let old3 = m.insert_with_timeout("k3", 3, Some(Duration::from_millis(80)));
589 let old4 = m.insert_with_timeout("k4", 4, Some(Duration::from_millis(130)));
590 let old44 = m.insert_with_timeout("k4", 44, None);
591
592 let old6 = m.insert_with_timeout("k6", 6, Some(Duration::from_secs(150)));
593 let old7 = m.insert_with_timeout("k7", 7, None);
594 let old66 = m.insert_with_timeout("k6", 66, Some(Duration::from_secs(60)));
595
596 assert_eq!(old1, None);
597 assert_eq!(old2, None);
598 assert_eq!(old3, None);
599 assert_eq!(old4, None);
600 assert_eq!(old6, None);
601 assert_eq!(old7, None);
602
603 println!("old44: {:?}", old44);
604 assert_eq!(old44, Some(4));
605 assert_eq!(m.get_with_timeout("k4"), Some(&44));
606
607 assert_eq!(old66, Some(6));
608 println!("old66: {:?}", old66);
609
610 let v6 = m.get("k6");
611 println!("v6: {:?}", v6);
612 assert_eq!(v6, Some(&TimedValue::new(66, None)));
613
614 m.get_with_timeout_mut("k6").map(|v| *v = 666);
615 let v6 = m.get_with_timeout("k6");
616 println!("v6: {:?}", v6);
617 assert_eq!(v6, Some(&666));
618
619 m.entry("kk1").or_insert_with_timeout_f(|| 10, None);
620 assert_eq!(m.get_with_timeout("kk1"), Some(&10));
621 m.entry("kk1").and_modify_with_timeout(|v| {
622 *v = 100;
623 });
624 assert_eq!(m.get_with_timeout("kk1"), Some(&100));
625 println!("kk1: {:?}", m.get_with_timeout("kk1"));
626}
627
628#[test]
629fn test_btree_map_ext_removes() {
630 let mut m: dequemap::DequeBTreeMap<_, _> = dequemap::DequeBTreeMap::default();
631 m.push_back(3, TimedValue::new((), Some(Duration::from_millis(800))));
632 std::thread::sleep(Duration::from_millis(100));
633 m.push_back(1, TimedValue::new((), Some(Duration::from_millis(800))));
634 std::thread::sleep(Duration::from_millis(100));
635 m.push_back(6, TimedValue::new((), Some(Duration::from_millis(800))));
636 std::thread::sleep(Duration::from_millis(100));
637 m.push_back(8, TimedValue::new((), Some(Duration::from_millis(800))));
638 std::thread::sleep(Duration::from_millis(100));
639 m.push_back(3, TimedValue::new((), Some(Duration::from_millis(800))));
640
641 assert_eq!(m.len(), 4);
642 for (key, item) in m.iter() {
643 println!("key: {:?}, is_expired: {}", key, item.is_expired());
644 }
645 println!("--------------------------------------------------------------");
646 std::thread::sleep(Duration::from_millis(600));
647
648 while let Some((key, item)) = m.front() {
649 println!("clean expired, key: {}", key);
650 if item.is_expired() {
651 m.pop_front();
652 } else {
653 break;
654 }
655 }
656 println!("m.len(): {}", m.len());
657 for (key, item) in m.iter() {
658 println!("key: {:?}, is_expired: {}", key, item.is_expired());
659 }
660 assert_eq!(m.len(), 2);
661}