cortexkit-bus-inmemory 0.3.2

Deterministic in-memory backend for the CortexKit NATS-shaped bus traits
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

use async_trait::async_trait;
use cortexkit_bus_trait::{
    BusError, BusResult, Register, RegisterEntry, RegisterUpdate, RegisterValue, RegisterWatch,
    Revision, Tombstone, WatchEvent,
};
use tokio::sync::mpsc;

use crate::backend::{InMemoryBus, Watcher};

pub struct InMemoryRegisterWatch {
    snapshot: Option<WatchEvent>,
    updates: mpsc::UnboundedReceiver<RegisterUpdate>,
    has_snapshot: bool,
}

impl InMemoryRegisterWatch {
    /// Simulates losing the transport before a pending snapshot can arrive.
    pub fn disconnect(&mut self) {
        self.snapshot = None;
        self.updates.close();
        self.has_snapshot = false;
    }
}

#[async_trait]
impl Register for InMemoryBus {
    type Watch = InMemoryRegisterWatch;

    async fn put(&self, key: &str, value: Vec<u8>) -> BusResult<Revision> {
        let mut state = self.inner.lock().expect("in-memory bus poisoned");
        let revision = state.next_revision();
        state.register_values.insert(
            key.into(),
            RegisterValue {
                value: value.clone(),
                revision,
            },
        );
        state.notify_watchers(RegisterUpdate::Put(RegisterEntry {
            key: key.into(),
            value,
            revision,
        }));
        Ok(revision)
    }

    async fn get(&self, key: &str) -> BusResult<RegisterValue> {
        self.inner
            .lock()
            .expect("in-memory bus poisoned")
            .register_values
            .get(key)
            .cloned()
            .ok_or_else(|| BusError::absent(format!("register key {key}")))
    }

    async fn delete(&self, key: &str) -> BusResult<Tombstone> {
        let mut state = self.inner.lock().expect("in-memory bus poisoned");
        if state.register_values.remove(key).is_none() {
            return Err(BusError::absent(format!("register key {key}")));
        }
        let tombstone = Tombstone {
            key: key.into(),
            revision: state.next_revision(),
        };
        state.notify_watchers(RegisterUpdate::Delete(tombstone.clone()));
        Ok(tombstone)
    }

    async fn watch(&self, prefix: &str) -> BusResult<Self::Watch> {
        let (sender, updates) = mpsc::unbounded_channel();
        let mut state = self.inner.lock().expect("in-memory bus poisoned");
        let entries = state
            .register_values
            .iter()
            .filter(|(key, _)| key.starts_with(prefix))
            .map(|(key, value)| RegisterEntry {
                key: key.clone(),
                value: value.value.clone(),
                revision: value.revision,
            })
            .collect();
        let revision = Revision(state.register_revision);
        state.watchers.push(Watcher {
            prefix: prefix.into(),
            sender,
        });
        Ok(InMemoryRegisterWatch {
            snapshot: Some(WatchEvent::Snapshot { entries, revision }),
            updates,
            has_snapshot: false,
        })
    }
}

#[async_trait]
impl RegisterWatch for InMemoryRegisterWatch {
    async fn next(&mut self) -> BusResult<WatchEvent> {
        if let Some(snapshot) = self.snapshot.take() {
            self.has_snapshot = true;
            return Ok(snapshot);
        }
        self.updates
            .recv()
            .await
            .map(WatchEvent::Update)
            .ok_or_else(|| BusError::unavailable("register watch connection dropped"))
    }

    fn has_snapshot(&self) -> bool {
        self.has_snapshot
    }
}