Skip to main content

ferra_agent/
cache.rs

1use std::collections::HashMap;
2use std::sync::atomic::{AtomicBool, AtomicI64, Ordering};
3use std::sync::RwLock;
4
5use serde_json::Value;
6use tokio::sync::broadcast;
7
8/// One key's value plus the event_id at which we last updated it.
9#[derive(Debug, Clone)]
10pub struct CacheEntry {
11    pub value: Value,
12    pub event_id: i64,
13}
14
15/// State for one watched prefix: the cache subtree, the watch cursor, and a
16/// notifier that wakes long-pollers whenever anything under this prefix
17/// changes.
18pub struct PrefixState {
19    pub prefix: String,
20    items: RwLock<HashMap<String, CacheEntry>>,
21    pub latest_event_id: AtomicI64,
22    pub ready: AtomicBool,
23    /// Sends a `()` whenever any item under this prefix changes (set,
24    /// delete, or snapshot replace). Long-pollers subscribe to this and
25    /// re-check after each notification.
26    pub notify: broadcast::Sender<()>,
27}
28
29impl PrefixState {
30    pub fn new(prefix: String) -> Self {
31        let (tx, _) = broadcast::channel(64);
32        Self {
33            prefix,
34            items: RwLock::new(HashMap::new()),
35            latest_event_id: AtomicI64::new(0),
36            ready: AtomicBool::new(false),
37            notify: tx,
38        }
39    }
40
41    pub fn get(&self, key: &str) -> Option<CacheEntry> {
42        self.items.read().expect("cache rwlock").get(key).cloned()
43    }
44
45    pub fn list(&self) -> Vec<(String, CacheEntry)> {
46        let g = self.items.read().expect("cache rwlock");
47        let mut v: Vec<_> = g.iter().map(|(k, e)| (k.clone(), e.clone())).collect();
48        v.sort_by(|a, b| a.0.cmp(&b.0));
49        v
50    }
51
52    pub fn upsert(&self, key: String, value: Value, event_id: i64) {
53        {
54            let mut g = self.items.write().expect("cache rwlock");
55            g.insert(key, CacheEntry { value, event_id });
56        }
57        self.advance(event_id);
58        let _ = self.notify.send(());
59    }
60
61    pub fn remove(&self, key: &str, event_id: i64) {
62        {
63            let mut g = self.items.write().expect("cache rwlock");
64            g.remove(key);
65        }
66        self.advance(event_id);
67        let _ = self.notify.send(());
68    }
69
70    pub fn replace_snapshot(&self, items: HashMap<String, CacheEntry>, latest_event_id: i64) {
71        {
72            let mut g = self.items.write().expect("cache rwlock");
73            *g = items;
74        }
75        // Snapshot resets the cursor, even if downward (the server told us
76        // older event_ids are no longer valid).
77        self.latest_event_id.store(latest_event_id, Ordering::Relaxed);
78        self.ready.store(true, Ordering::Relaxed);
79        let _ = self.notify.send(());
80    }
81
82    fn advance(&self, event_id: i64) {
83        // Monotonic: only increase.
84        loop {
85            let cur = self.latest_event_id.load(Ordering::Relaxed);
86            if event_id <= cur {
87                return;
88            }
89            if self
90                .latest_event_id
91                .compare_exchange(cur, event_id, Ordering::Relaxed, Ordering::Relaxed)
92                .is_ok()
93            {
94                return;
95            }
96        }
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use serde_json::json;
104
105    #[test]
106    fn upsert_and_get() {
107        let p = PrefixState::new("services/payment/".into());
108        p.upsert("services/payment/timeout_ms".into(), json!(3000), 5);
109
110        let e = p.get("services/payment/timeout_ms").unwrap();
111        assert_eq!(e.value, json!(3000));
112        assert_eq!(e.event_id, 5);
113        assert_eq!(p.latest_event_id.load(Ordering::Relaxed), 5);
114    }
115
116    #[test]
117    fn remove_drops_entry() {
118        let p = PrefixState::new("p/".into());
119        p.upsert("p/a".into(), json!(1), 1);
120        p.remove("p/a", 2);
121        assert!(p.get("p/a").is_none());
122        assert_eq!(p.latest_event_id.load(Ordering::Relaxed), 2);
123    }
124
125    #[test]
126    fn replace_snapshot_resets_cursor_and_marks_ready() {
127        let p = PrefixState::new("p/".into());
128        p.upsert("p/old".into(), json!(1), 7);
129        assert!(!p.ready.load(Ordering::Relaxed));
130
131        let mut snap = HashMap::new();
132        snap.insert(
133            "p/new".into(),
134            CacheEntry {
135                value: json!(2),
136                event_id: 9,
137            },
138        );
139        p.replace_snapshot(snap, 9);
140
141        assert!(p.get("p/old").is_none());
142        assert_eq!(p.get("p/new").unwrap().value, json!(2));
143        assert_eq!(p.latest_event_id.load(Ordering::Relaxed), 9);
144        assert!(p.ready.load(Ordering::Relaxed));
145    }
146
147    #[test]
148    fn cursor_only_advances_forward() {
149        let p = PrefixState::new("p/".into());
150        p.upsert("p/a".into(), json!(1), 5);
151        // An older event_id should not push the cursor backwards.
152        p.upsert("p/b".into(), json!(2), 3);
153        assert_eq!(p.latest_event_id.load(Ordering::Relaxed), 5);
154    }
155
156    #[test]
157    fn list_is_sorted() {
158        let p = PrefixState::new("p/".into());
159        p.upsert("p/c".into(), json!(3), 3);
160        p.upsert("p/a".into(), json!(1), 1);
161        p.upsert("p/b".into(), json!(2), 2);
162        let keys: Vec<_> = p.list().into_iter().map(|(k, _)| k).collect();
163        assert_eq!(keys, vec!["p/a", "p/b", "p/c"]);
164    }
165}