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::btree_map::Entry; 6 use std::collections::BTreeMap; 7 use std::collections::VecDeque; 8 use std::time::Duration; 9 use std::time::Instant; 10 11 pub struct ExpiringMap<K, V> { 12 limit: Duration, 13 map: BTreeMap<K, (V, Instant)>, 14 dq: VecDeque<(K, Instant)>, 15 } 16 17 impl<K, V> ExpiringMap<K, V> 18 where 19 K: Clone + Ord, 20 { new(limit: Duration) -> Self21 pub fn new(limit: Duration) -> Self { 22 Self { 23 limit, 24 map: Default::default(), 25 dq: Default::default(), 26 } 27 } 28 cleanup(&mut self, now: &Instant)29 fn cleanup(&mut self, now: &Instant) { 30 while let Some((k, prev)) = self.dq.front() { 31 if now.duration_since(*prev) < self.limit { 32 return; 33 } 34 self.map.remove(k); 35 self.dq.pop_front(); 36 } 37 } 38 39 #[cfg(test)] get(&mut self, key: &K) -> Option<&V>40 pub fn get(&mut self, key: &K) -> Option<&V> { 41 let now = Instant::now(); 42 self.cleanup(&now); 43 self.map.get(key).map(|(v, _)| v) 44 } 45 get_or_insert_with<F: FnOnce() -> std::result::Result<V, E>, E>( &mut self, key: &K, f: F, ) -> std::result::Result<&V, E>46 pub fn get_or_insert_with<F: FnOnce() -> std::result::Result<V, E>, E>( 47 &mut self, 48 key: &K, 49 f: F, 50 ) -> std::result::Result<&V, E> { 51 let now = Instant::now(); 52 self.cleanup(&now); 53 54 let (v, _) = match self.map.entry(key.clone()) { 55 Entry::Occupied(o) => o.into_mut(), 56 Entry::Vacant(v) => { 57 let value = f()?; 58 self.dq.push_back((key.clone(), now)); 59 v.insert((value, now)) 60 } 61 }; 62 Ok(v) 63 } 64 get_mut(&mut self, key: &K) -> Option<&mut V>65 pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { 66 let now = Instant::now(); 67 self.cleanup(&now); 68 self.map.get_mut(key).map(|(v, _)| v) 69 } 70 remove(&mut self, key: &K)71 pub fn remove(&mut self, key: &K) { 72 self.map.remove(key); 73 self.dq.retain(|(k, _)| k != key); 74 } 75 } 76