cortexkit_bus_inmemory/
register.rs

1use async_trait::async_trait;
2use cortexkit_bus_trait::{
3    BusError, BusResult, Register, RegisterEntry, RegisterUpdate, RegisterValue, RegisterWatch,
4    Revision, Tombstone, WatchEvent,
5};
6use tokio::sync::mpsc;
7
8use crate::backend::{InMemoryBus, Watcher};
9
10pub struct InMemoryRegisterWatch {
11    snapshot: Option<WatchEvent>,
12    updates: mpsc::UnboundedReceiver<RegisterUpdate>,
13    has_snapshot: bool,
14}
15
16impl InMemoryRegisterWatch {
17    /// Simulates losing the transport before a pending snapshot can arrive.
18    pub fn disconnect(&mut self) {
19        self.snapshot = None;
20        self.updates.close();
21        self.has_snapshot = false;
22    }
23}
24
25#[async_trait]
26impl Register for InMemoryBus {
27    type Watch = InMemoryRegisterWatch;
28
29    async fn put(&self, key: &str, value: Vec<u8>) -> BusResult<Revision> {
30        let mut state = self.inner.lock().expect("in-memory bus poisoned");
31        let revision = state.next_revision();
32        state.register_values.insert(
33            key.into(),
34            RegisterValue {
35                value: value.clone(),
36                revision,
37            },
38        );
39        state.notify_watchers(RegisterUpdate::Put(RegisterEntry {
40            key: key.into(),
41            value,
42            revision,
43        }));
44        Ok(revision)
45    }
46
47    async fn get(&self, key: &str) -> BusResult<RegisterValue> {
48        self.inner
49            .lock()
50            .expect("in-memory bus poisoned")
51            .register_values
52            .get(key)
53            .cloned()
54            .ok_or_else(|| BusError::absent(format!("register key {key}")))
55    }
56
57    async fn delete(&self, key: &str) -> BusResult<Tombstone> {
58        let mut state = self.inner.lock().expect("in-memory bus poisoned");
59        if state.register_values.remove(key).is_none() {
60            return Err(BusError::absent(format!("register key {key}")));
61        }
62        let tombstone = Tombstone {
63            key: key.into(),
64            revision: state.next_revision(),
65        };
66        state.notify_watchers(RegisterUpdate::Delete(tombstone.clone()));
67        Ok(tombstone)
68    }
69
70    async fn watch(&self, prefix: &str) -> BusResult<Self::Watch> {
71        let (sender, updates) = mpsc::unbounded_channel();
72        let mut state = self.inner.lock().expect("in-memory bus poisoned");
73        let entries = state
74            .register_values
75            .iter()
76            .filter(|(key, _)| key.starts_with(prefix))
77            .map(|(key, value)| RegisterEntry {
78                key: key.clone(),
79                value: value.value.clone(),
80                revision: value.revision,
81            })
82            .collect();
83        let revision = Revision(state.register_revision);
84        state.watchers.push(Watcher {
85            prefix: prefix.into(),
86            sender,
87        });
88        Ok(InMemoryRegisterWatch {
89            snapshot: Some(WatchEvent::Snapshot { entries, revision }),
90            updates,
91            has_snapshot: false,
92        })
93    }
94}
95
96#[async_trait]
97impl RegisterWatch for InMemoryRegisterWatch {
98    async fn next(&mut self) -> BusResult<WatchEvent> {
99        if let Some(snapshot) = self.snapshot.take() {
100            self.has_snapshot = true;
101            return Ok(snapshot);
102        }
103        self.updates
104            .recv()
105            .await
106            .map(WatchEvent::Update)
107            .ok_or_else(|| BusError::unavailable("register watch connection dropped"))
108    }
109
110    fn has_snapshot(&self) -> bool {
111        self.has_snapshot
112    }
113}
cortexkit_bus_inmemory/register.rs

cortexkit_bus_inmemory/
register.rs
