crdb_cache/
object_cache.rs

1use crdb_core::{DynSized, ObjectId, SystemTimeExt};
2use std::{
3    collections::HashMap,
4    sync::{
5        atomic::{AtomicI64, Ordering},
6        Arc,
7    },
8};
9use web_time::SystemTime;
10
11// TODO(test-med): test this
12
13pub struct ObjectCache {
14    watermark: usize,
15    // AtomicI64 here is the timestamp (in ms since unix epoch) of the last access to the object
16    objects: HashMap<ObjectId, (AtomicI64, Arc<dyn DynSized>)>,
17    approx_exclusive_size: usize,
18}
19
20impl ObjectCache {
21    pub fn new(watermark: usize) -> ObjectCache {
22        ObjectCache {
23            watermark,
24            objects: HashMap::new(),
25            approx_exclusive_size: 0,
26        }
27    }
28
29    fn add_approx_size(&mut self, size: usize) {
30        self.approx_exclusive_size = self.approx_exclusive_size.saturating_add(size);
31    }
32
33    fn rm_approx_size(&mut self, size: usize) {
34        self.approx_exclusive_size = self.approx_exclusive_size.saturating_sub(size);
35    }
36
37    fn recompute_exclusive_size(&mut self) {
38        self.approx_exclusive_size = 0;
39        for (_, v) in self.objects.values() {
40            if Arc::strong_count(v) == 1 {
41                self.approx_exclusive_size += v.deep_size_of();
42            }
43        }
44    }
45
46    pub fn set(&mut self, object_id: ObjectId, value: Arc<dyn DynSized>) {
47        let now = AtomicI64::new(SystemTime::now().ms_since_posix().unwrap());
48        self.add_approx_size(value.deep_size_of());
49        if let Some(previous) = self.objects.insert(object_id, (now, value)) {
50            self.rm_approx_size(previous.1.deep_size_of());
51        }
52        if self.approx_exclusive_size > self.watermark {
53            self.recompute_exclusive_size();
54            if self.approx_exclusive_size > self.watermark {
55                self.apply_watermark();
56            }
57        }
58    }
59
60    pub fn get(&self, id: &ObjectId) -> Option<Arc<dyn DynSized>> {
61        self.objects.get(id).map(|(access, v)| {
62            access.store(
63                SystemTime::now().ms_since_posix().unwrap(),
64                Ordering::Relaxed,
65            );
66            v.clone()
67        })
68    }
69
70    pub fn remove(&mut self, object_id: &ObjectId) {
71        if let Some(previous) = self.objects.remove(object_id) {
72            self.rm_approx_size(previous.1.deep_size_of());
73        }
74    }
75
76    fn apply_watermark(&mut self) {
77        let mut all_entries = self
78            .objects
79            .iter()
80            .map(|(id, (t, v))| (t.load(Ordering::Relaxed), *id, v.clone()))
81            .collect::<Vec<_>>();
82        all_entries.sort_unstable_by_key(|(t, _, _)| *t);
83        for (_, id, v) in all_entries {
84            if Arc::strong_count(&v) == 2 {
85                // One is `all_entries`, one is the HashMap Arc
86                self.objects.remove(&id);
87                self.rm_approx_size(v.deep_size_of());
88                if self.approx_exclusive_size <= self.watermark / 2 {
89                    break;
90                }
91            }
92        }
93        if self.approx_exclusive_size > self.watermark {
94            self.watermark = self.approx_exclusive_size;
95        }
96    }
97}