disk/qcow/
vec_cache.rs

1// Copyright 2018 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::hash_map::IterMut;
6use std::collections::HashMap;
7use std::io;
8use std::ops::Index;
9use std::ops::IndexMut;
10use std::slice::SliceIndex;
11
12/// Trait that allows for checking if an implementor is dirty. Useful for types that are cached so
13/// it can be checked if they need to be committed to disk.
14pub trait Cacheable {
15    /// Used to check if the item needs to be written out or if it can be discarded.
16    fn dirty(&self) -> bool;
17}
18
19#[derive(Debug)]
20/// Represents a vector that implements the `Cacheable` trait so it can be held in a cache.
21pub struct VecCache<T: 'static + Copy + Default> {
22    vec: Box<[T]>,
23    dirty: bool,
24}
25
26impl<T: 'static + Copy + Default> VecCache<T> {
27    /// Creates a `VecCache` that can hold `count` elements.
28    pub fn new(count: usize) -> VecCache<T> {
29        VecCache {
30            vec: vec![Default::default(); count].into_boxed_slice(),
31            dirty: true,
32        }
33    }
34
35    /// Creates a `VecCache` from the passed in `vec`.
36    pub fn from_vec(vec: Vec<T>) -> VecCache<T> {
37        VecCache {
38            vec: vec.into_boxed_slice(),
39            dirty: false,
40        }
41    }
42
43    pub fn get<I>(&self, index: I) -> Option<&<I as SliceIndex<[T]>>::Output>
44    where
45        I: SliceIndex<[T]>,
46    {
47        self.vec.get(index)
48    }
49
50    /// Gets a reference to the underlying vector.
51    pub fn get_values(&self) -> &[T] {
52        &self.vec
53    }
54
55    /// Mark this cache element as clean.
56    pub fn mark_clean(&mut self) {
57        self.dirty = false;
58    }
59
60    /// Returns the number of elements in the vector.
61    pub fn len(&self) -> usize {
62        self.vec.len()
63    }
64}
65
66impl<T: 'static + Copy + Default> Cacheable for VecCache<T> {
67    fn dirty(&self) -> bool {
68        self.dirty
69    }
70}
71
72impl<T: 'static + Copy + Default> Index<usize> for VecCache<T> {
73    type Output = T;
74
75    fn index(&self, index: usize) -> &T {
76        self.vec.index(index)
77    }
78}
79
80impl<T: 'static + Copy + Default> IndexMut<usize> for VecCache<T> {
81    fn index_mut(&mut self, index: usize) -> &mut T {
82        self.dirty = true;
83        self.vec.index_mut(index)
84    }
85}
86
87#[derive(Debug)]
88pub struct CacheMap<T: Cacheable> {
89    capacity: usize,
90    map: HashMap<usize, T>,
91}
92
93impl<T: Cacheable> CacheMap<T> {
94    pub fn new(capacity: usize) -> Self {
95        CacheMap {
96            capacity,
97            map: HashMap::with_capacity(capacity),
98        }
99    }
100
101    pub fn contains_key(&self, key: &usize) -> bool {
102        self.map.contains_key(key)
103    }
104
105    pub fn get(&self, index: &usize) -> Option<&T> {
106        self.map.get(index)
107    }
108
109    pub fn get_mut(&mut self, index: &usize) -> Option<&mut T> {
110        self.map.get_mut(index)
111    }
112
113    pub fn iter_mut(&mut self) -> IterMut<usize, T> {
114        self.map.iter_mut()
115    }
116
117    // Check if the refblock cache is full and we need to evict.
118    pub fn insert<F>(&mut self, index: usize, block: T, write_callback: F) -> io::Result<()>
119    where
120        F: FnOnce(usize, T) -> io::Result<()>,
121    {
122        if self.map.len() == self.capacity {
123            // TODO(dgreid) - smarter eviction strategy.
124            let to_evict = *self.map.iter().next().unwrap().0;
125            if let Some(evicted) = self.map.remove(&to_evict) {
126                if evicted.dirty() {
127                    write_callback(to_evict, evicted)?;
128                }
129            }
130        }
131        self.map.insert(index, block);
132        Ok(())
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    #[derive(Copy, Clone, Debug, Eq, PartialEq)]
141    struct NumCache(pub u64);
142    impl Cacheable for NumCache {
143        fn dirty(&self) -> bool {
144            true
145        }
146    }
147
148    #[test]
149    fn evicts_when_full() {
150        let mut cache = CacheMap::<NumCache>::new(3);
151        let mut evicted = None;
152        cache
153            .insert(0, NumCache(5), |index, _| {
154                evicted = Some(index);
155                Ok(())
156            })
157            .unwrap();
158        assert_eq!(evicted, None);
159        cache
160            .insert(1, NumCache(6), |index, _| {
161                evicted = Some(index);
162                Ok(())
163            })
164            .unwrap();
165        assert_eq!(evicted, None);
166        cache
167            .insert(2, NumCache(7), |index, _| {
168                evicted = Some(index);
169                Ok(())
170            })
171            .unwrap();
172        assert_eq!(evicted, None);
173        cache
174            .insert(3, NumCache(8), |index, _| {
175                evicted = Some(index);
176                Ok(())
177            })
178            .unwrap();
179        assert!(evicted.is_some());
180
181        // Check that three of the four items inserted are still there and that the most recently
182        // inserted is one of them.
183        let num_items = (0..=3).filter(|k| cache.contains_key(k)).count();
184        assert_eq!(num_items, 3);
185        assert!(cache.contains_key(&3));
186        assert_eq!(cache.get(&3), Some(&NumCache(8)));
187    }
188}