cecile_supercool_tracker/trackers/
epoch_db.rs

1use crate::track::TrackStatus;
2use anyhow::Result;
3use std::collections::HashMap;
4use std::sync::RwLock;
5
6pub trait EpochDb {
7    fn epoch_db(&self) -> &Option<RwLock<HashMap<u64, usize>>>;
8    fn max_idle_epochs(&self) -> usize;
9
10    fn skip_epochs_for_scene(&self, scene_id: u64, n: usize) {
11        if let Some(epoch_store) = self.epoch_db() {
12            let mut epoch_store = epoch_store.write().unwrap();
13            if let Some(epoch) = epoch_store.get_mut(&scene_id) {
14                *epoch += n;
15            } else {
16                epoch_store.insert(scene_id, n);
17            }
18        }
19    }
20
21    fn current_epoch_with_scene(&self, scene_id: u64) -> Option<usize> {
22        if let Some(epoch_store) = self.epoch_db() {
23            let mut epoch_store = epoch_store.write().unwrap();
24            let epoch = epoch_store.get_mut(&scene_id);
25            if let Some(epoch) = epoch {
26                Some(*epoch)
27            } else {
28                Some(0)
29            }
30        } else {
31            None
32        }
33    }
34
35    fn next_epoch(&self, scene_id: u64) -> Option<usize> {
36        if let Some(epoch_store) = self.epoch_db() {
37            let mut epoch_store = epoch_store.write().unwrap();
38            let epoch = epoch_store.get_mut(&scene_id);
39            if let Some(epoch) = epoch {
40                *epoch += 1;
41                Some(*epoch)
42            } else {
43                epoch_store.insert(scene_id, 1);
44                Some(1)
45            }
46        } else {
47            None
48        }
49    }
50
51    fn baked(&self, scene_id: u64, last_updated: usize) -> Result<TrackStatus> {
52        if let Some(current_epoch) = &self.epoch_db() {
53            let current_epoch = current_epoch.read().unwrap();
54            if last_updated + self.max_idle_epochs() < *current_epoch.get(&scene_id).unwrap_or(&0) {
55                Ok(TrackStatus::Wasted)
56            } else {
57                Ok(TrackStatus::Pending)
58            }
59        } else {
60            // If epoch expiration is not set the tracks are always ready.
61            // If set, then only when certain amount of epochs pass they are Wasted.
62            //
63            Ok(TrackStatus::Ready)
64        }
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use crate::track::TrackStatus;
71    use crate::trackers::epoch_db::EpochDb;
72    use std::collections::HashMap;
73    use std::sync::RwLock;
74
75    #[test]
76    fn test_epoch_db() {
77        #[derive(Debug, Default)]
78        pub struct DbOptions {
79            /// The map that stores current epochs for the scene_id
80            epoch_db: Option<RwLock<HashMap<u64, usize>>>,
81            /// The maximum number of epochs without update while the track is alive
82            max_idle_epochs: usize,
83        }
84
85        impl EpochDb for DbOptions {
86            fn epoch_db(&self) -> &Option<RwLock<HashMap<u64, usize>>> {
87                &self.epoch_db
88            }
89
90            fn max_idle_epochs(&self) -> usize {
91                self.max_idle_epochs
92            }
93        }
94
95        let db = DbOptions {
96            epoch_db: Some(RwLock::new(HashMap::default())),
97            max_idle_epochs: 2,
98        };
99
100        assert_eq!(db.next_epoch(0), Some(1));
101        assert_eq!(db.next_epoch(0), Some(2));
102
103        assert_eq!(db.next_epoch(1), Some(1));
104        assert_eq!(db.next_epoch(1), Some(2));
105
106        assert_eq!(db.current_epoch_with_scene(0), Some(2));
107        assert_eq!(db.current_epoch_with_scene(1), Some(2));
108        assert_eq!(db.current_epoch_with_scene(2), Some(0));
109
110        db.skip_epochs_for_scene(0, 10);
111        db.skip_epochs_for_scene(1, 2);
112
113        assert!(matches!(db.baked(0, 2), Ok(TrackStatus::Wasted)));
114        assert!(matches!(db.baked(1, 2), Ok(TrackStatus::Pending)));
115
116        db.skip_epochs_for_scene(1, 1);
117        assert!(matches!(db.baked(1, 2), Ok(TrackStatus::Wasted)));
118    }
119}