devices/virtio/fs/
expiring_map.rs

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
5use std::collections::btree_map::Entry;
6use std::collections::BTreeMap;
7use std::collections::VecDeque;
8use std::time::Duration;
9use std::time::Instant;
10
11pub struct ExpiringMap<K, V> {
12    limit: Duration,
13    map: BTreeMap<K, (V, Instant)>,
14    dq: VecDeque<(K, Instant)>,
15}
16
17impl<K, V> ExpiringMap<K, V>
18where
19    K: Clone + Ord,
20{
21    pub fn new(limit: Duration) -> Self {
22        Self {
23            limit,
24            map: Default::default(),
25            dq: Default::default(),
26        }
27    }
28
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)]
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
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
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
71    pub fn remove(&mut self, key: &K) {
72        self.map.remove(key);
73        self.dq.retain(|(k, _)| k != key);
74    }
75}