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#[derive(Debug, Clone)]
10pub struct CacheEntry {
11 pub value: Value,
12 pub event_id: i64,
13}
14
15pub struct PrefixState {
19 pub prefix: String,
20 items: RwLock<HashMap<String, CacheEntry>>,
21 pub latest_event_id: AtomicI64,
22 pub ready: AtomicBool,
23 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 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 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 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}