1 // Copyright 2023 The ChromiumOS Authors 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 use std::collections::BTreeMap; 6 use std::collections::VecDeque; 7 use std::time::Duration; 8 use std::time::Instant; 9 10 pub struct ExpiringMap<K, V> { 11 limit: Duration, 12 map: BTreeMap<K, (V, Instant)>, 13 dq: VecDeque<(K, Instant)>, 14 } 15 16 impl<K, V> ExpiringMap<K, V> 17 where 18 K: Clone + Ord, 19 { new(limit: Duration) -> Self20 pub fn new(limit: Duration) -> Self { 21 Self { 22 limit, 23 map: Default::default(), 24 dq: Default::default(), 25 } 26 } 27 cleanup(&mut self, now: &Instant)28 fn cleanup(&mut self, now: &Instant) { 29 while let Some((k, prev)) = self.dq.front() { 30 if now.duration_since(*prev) < self.limit { 31 return; 32 } 33 self.map.remove(k); 34 self.dq.pop_front(); 35 } 36 } 37 38 #[cfg(test)] get(&mut self, key: &K) -> Option<&V>39 pub fn get(&mut self, key: &K) -> Option<&V> { 40 let now = Instant::now(); 41 self.cleanup(&now); 42 self.map.get(key).map(|(v, _)| v) 43 } 44 get_or_insert_with<F: FnOnce() -> std::result::Result<V, E>, E>( &mut self, key: &K, f: F, ) -> std::result::Result<&V, E>45 pub fn get_or_insert_with<F: FnOnce() -> std::result::Result<V, E>, E>( 46 &mut self, 47 key: &K, 48 f: F, 49 ) -> std::result::Result<&V, E> { 50 let now = Instant::now(); 51 self.cleanup(&now); 52 53 if self.map.get(key).is_some() { 54 Ok(self.map.get(key).map(|(v, _)| v).expect("must exist")) 55 } else { 56 self.dq.push_back((key.clone(), now)); 57 self.map.insert(key.clone(), (f()?, now)); 58 Ok(self.map.get(key).map(|(v, _)| v).expect("must exist")) 59 } 60 } 61 get_mut(&mut self, key: &K) -> Option<&mut V>62 pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { 63 let now = Instant::now(); 64 self.cleanup(&now); 65 self.map.get_mut(key).map(|(v, _)| v) 66 } 67 remove(&mut self, key: &K)68 pub fn remove(&mut self, key: &K) { 69 self.map.remove(key); 70 self.dq = self 71 .dq 72 .iter() 73 .filter(|(k, _)| k != key) 74 .map(|(k, time)| (k.clone(), *time)) 75 .collect(); 76 } 77 } 78