cortexkit_bus_trait/
census.rs1use std::collections::BTreeSet;
2
3use crate::{BusError, BusResult, RegisterUpdate, WatchEvent};
4
5#[derive(Debug, Default, Clone)]
8pub struct CensusWatchGuard {
9 present: BTreeSet<String>,
10 observed_deleted: BTreeSet<String>,
11 has_snapshot: bool,
12 connection_dropped: bool,
13}
14
15impl CensusWatchGuard {
16 pub fn new() -> Self {
17 Self::default()
18 }
19
20 pub fn observe(&mut self, event: &WatchEvent) {
21 if self.connection_dropped && !matches!(event, WatchEvent::Snapshot { .. }) {
22 return;
23 }
24 match event {
25 WatchEvent::Snapshot { entries, .. } => {
26 self.connection_dropped = false;
27 self.present = entries.iter().map(|entry| entry.key.clone()).collect();
28 self.observed_deleted.clear();
29 self.has_snapshot = true;
30 }
31 WatchEvent::Update(RegisterUpdate::Put(entry)) => {
32 self.present.insert(entry.key.clone());
33 self.observed_deleted.remove(&entry.key);
34 }
35 WatchEvent::Update(RegisterUpdate::Delete(tombstone)) => {
36 self.present.remove(&tombstone.key);
37 self.observed_deleted.insert(tombstone.key.clone());
38 }
39 }
40 }
41
42 pub fn connection_dropped(&mut self) {
44 self.present.clear();
45 self.observed_deleted.clear();
46 self.has_snapshot = false;
47 self.connection_dropped = true;
48 }
49
50 pub fn has_snapshot(&self) -> bool {
51 self.has_snapshot
52 }
53
54 pub fn check_publish(&self, census_key: &str, subject: &str) -> BusResult<()> {
55 if self.connection_dropped {
56 return Ok(());
57 }
58
59 let confirmed_absent = self.observed_deleted.contains(census_key)
60 || (self.has_snapshot && !self.present.contains(census_key));
61 if confirmed_absent {
62 return Err(BusError::denied(
63 subject,
64 format!("confirmed absent census key {census_key}"),
65 ));
66 }
67 Ok(())
68 }
69}
70
71#[cfg(test)]
72mod tests {
73 use super::*;
74 use crate::{RegisterEntry, Revision, Tombstone};
75
76 #[test]
77 fn no_snapshot_and_dropped_connection_are_both_neutral() {
78 let mut guard = CensusWatchGuard::new();
79 assert_eq!(
80 guard.check_publish("module-a", "ck.box.room.r.post"),
81 Ok(())
82 );
83
84 guard.observe(&WatchEvent::Snapshot {
85 entries: Vec::new(),
86 revision: Revision(0),
87 });
88 assert!(matches!(
89 guard.check_publish("module-a", "ck.box.room.r.post"),
90 Err(BusError::Denied { .. })
91 ));
92
93 guard.connection_dropped();
94 guard.observe(&WatchEvent::Update(RegisterUpdate::Delete(Tombstone {
95 key: "module-a".into(),
96 revision: Revision(1),
97 })));
98 assert_eq!(
99 guard.check_publish("module-a", "ck.box.room.r.post"),
100 Ok(())
101 );
102 }
103
104 #[test]
105 fn observed_delete_confirms_absence_even_without_a_snapshot() {
106 let mut guard = CensusWatchGuard::new();
107 guard.observe(&WatchEvent::Update(RegisterUpdate::Delete(Tombstone {
108 key: "module-a".into(),
109 revision: Revision(2),
110 })));
111
112 assert!(matches!(
113 guard.check_publish("module-a", "ck.box.wake.agent.fire"),
114 Err(BusError::Denied { subject, .. }) if subject == "ck.box.wake.agent.fire"
115 ));
116
117 guard.observe(&WatchEvent::Update(RegisterUpdate::Put(RegisterEntry {
118 key: "module-a".into(),
119 value: Vec::new(),
120 revision: Revision(3),
121 })));
122 assert_eq!(
123 guard.check_publish("module-a", "ck.box.wake.agent.fire"),
124 Ok(())
125 );
126 }
127}