cortexkit_bus_inmemory/
backend.rs

1use std::collections::{BTreeMap, BTreeSet, HashMap, VecDeque};
2use std::sync::{Arc, Mutex};
3use std::time::{Duration, Instant};
4
5use cortexkit_bus_naming::AccountNames;
6use cortexkit_bus_trait::{
7    AsyncPermissionViolation, BusError, BusResult, DeliveryToken, DurableOwner, Message,
8    RegisterUpdate, RegisterValue, Revision,
9};
10use tokio::sync::mpsc;
11
12#[derive(Debug, Clone)]
13pub struct WorkQueueConfig {
14    pub subject: String,
15    pub max_bytes: usize,
16    pub max_deliveries: u32,
17    pub max_nak_delay: Duration,
18}
19
20impl WorkQueueConfig {
21    pub fn new(subject: impl Into<String>, max_bytes: usize) -> Self {
22        Self {
23            subject: subject.into(),
24            max_bytes,
25            max_deliveries: 5,
26            max_nak_delay: Duration::from_secs(30),
27        }
28    }
29}
30
31/// The stream durables' ack wait unless a test sets another: the plane's
32/// consumers use 30 seconds.
33pub const DEFAULT_STREAM_ACK_WAIT: Duration = Duration::from_secs(30);
34
35#[derive(Debug, Clone)]
36pub struct InMemoryConfig {
37    /// Durable names, `c_{agent_id}` or `m_{module_id}`.
38    pub(crate) durables: BTreeSet<String>,
39    pub(crate) denied_subjects: BTreeSet<String>,
40    pub(crate) work_queue: Option<WorkQueueConfig>,
41    pub(crate) stream_ack_wait: Duration,
42}
43
44impl Default for InMemoryConfig {
45    fn default() -> Self {
46        Self {
47            durables: BTreeSet::new(),
48            denied_subjects: BTreeSet::new(),
49            work_queue: None,
50            stream_ack_wait: DEFAULT_STREAM_ACK_WAIT,
51        }
52    }
53}
54
55impl InMemoryConfig {
56    /// How long a stream delivery may stay unsettled before another cursor on
57    /// the same durable receives it again; applies to every durable.
58    pub fn with_stream_ack_wait(mut self, ack_wait: Duration) -> Self {
59        self.stream_ack_wait = ack_wait;
60        self
61    }
62
63    /// Creates the owner's durable, as ck-bus would before any cursor binds.
64    ///
65    /// # Panics
66    ///
67    /// When the owner's id is not a valid agent or module id token: the
68    /// fixture would otherwise hold a durable no cursor can bind.
69    pub fn with_durable(mut self, owner: DurableOwner<'_>) -> Self {
70        let name = durable_name(owner).expect("fixture durable owner must be a valid id token");
71        self.durables.insert(name);
72        self
73    }
74
75    pub fn deny_subject(mut self, subject: impl Into<String>) -> Self {
76        self.denied_subjects.insert(subject.into());
77        self
78    }
79
80    pub fn with_work_queue(mut self, config: WorkQueueConfig) -> Self {
81        self.work_queue = Some(config);
82        self
83    }
84}
85
86/// The owner's durable name, `c_{agent_id}` or `m_{module_id}`, from
87/// `AccountNames::consumer_name` and `AccountNames::module_consumer_name`.
88/// Both refuse an id that is not one plain subject token, since a dot or a
89/// wildcard in it would address some other consumer.
90pub(crate) fn durable_name(owner: DurableOwner<'_>) -> BusResult<String> {
91    let (id, name) = match owner {
92        DurableOwner::Agent(agent_id) => (agent_id, AccountNames::consumer_name(agent_id)),
93        DurableOwner::Module(module_id) => {
94            (module_id, AccountNames::module_consumer_name(module_id))
95        }
96    };
97    name.map_err(|error| BusError::denied(id, error.to_string()))
98}
99
100#[derive(Debug, Clone, PartialEq, Eq)]
101pub enum BackendEvent {
102    Published {
103        subject: String,
104        id: String,
105    },
106    Acknowledged {
107        token: DeliveryToken,
108    },
109    NegativelyAcknowledged {
110        token: DeliveryToken,
111    },
112    Terminated {
113        token: DeliveryToken,
114    },
115    /// A stream delivery settled with `StreamCursor::term`.
116    StreamTerminated {
117        stream_seq: u64,
118        reason: Option<String>,
119    },
120}
121
122#[derive(Clone)]
123pub struct InMemoryBus {
124    pub(crate) inner: Arc<Mutex<State>>,
125}
126
127impl InMemoryBus {
128    pub fn new(config: InMemoryConfig) -> Self {
129        let durable_cursors = config
130            .durables
131            .iter()
132            .cloned()
133            .map(|key| (key, DurableState::default()))
134            .collect();
135        Self {
136            inner: Arc::new(Mutex::new(State {
137                config,
138                next_stream_seq: 1,
139                stream_messages: Vec::new(),
140                durable_cursors,
141                register_revision: 0,
142                register_values: BTreeMap::new(),
143                watchers: Vec::new(),
144                queue: VecDeque::new(),
145                next_delivery_token: 1,
146                events: Vec::new(),
147                permission_violations: VecDeque::new(),
148            })),
149        }
150    }
151
152    pub fn events(&self) -> Vec<BackendEvent> {
153        self.inner
154            .lock()
155            .expect("in-memory bus poisoned")
156            .events
157            .clone()
158    }
159
160    pub fn queued_bytes(&self) -> usize {
161        self.inner
162            .lock()
163            .expect("in-memory bus poisoned")
164            .queue
165            .iter()
166            .map(|entry| entry.message.wire_size())
167            .sum()
168    }
169
170    pub fn inject_permission_violation(&self, subject: impl Into<String>) {
171        let subject = subject.into();
172        self.inner
173            .lock()
174            .expect("in-memory bus poisoned")
175            .permission_violations
176            .push_back(AsyncPermissionViolation {
177                subject,
178                reason: "fixture-configured permission violation".into(),
179            });
180    }
181}
182
183pub(crate) struct State {
184    pub config: InMemoryConfig,
185    pub next_stream_seq: u64,
186    pub stream_messages: Vec<StoredMessage>,
187    pub durable_cursors: HashMap<String, DurableState>,
188    pub register_revision: u64,
189    pub register_values: BTreeMap<String, RegisterValue>,
190    pub watchers: Vec<Watcher>,
191    pub queue: VecDeque<QueueEntry>,
192    pub next_delivery_token: u64,
193    pub events: Vec<BackendEvent>,
194    pub permission_violations: VecDeque<AsyncPermissionViolation>,
195}
196
197#[derive(Clone)]
198pub(crate) struct StoredMessage {
199    pub stream_seq: u64,
200    pub message: Message,
201}
202
203/// One durable's position. Delivery is strictly in order: the message at
204/// `next_index` is the only one a cursor can receive until it is acked or
205/// terminated.
206#[derive(Default)]
207pub(crate) struct DurableState {
208    pub next_index: usize,
209    /// Set by a nak: no cursor receives the head message before this instant.
210    pub available_at: Option<Instant>,
211    /// How many times the head message has been delivered.
212    pub delivery_count: u64,
213    /// The cursor holding the head message unsettled, if any.
214    pub lease: Option<Lease>,
215    pub next_cursor_id: u64,
216}
217
218/// An unsettled delivery of the head message. Until `deadline` no other
219/// cursor receives the message; after it the next `next` on any cursor
220/// redelivers it, as JetStream does once the ack wait expires.
221pub(crate) struct Lease {
222    pub cursor_id: u64,
223    pub deadline: Instant,
224}
225
226pub(crate) struct Watcher {
227    pub prefix: String,
228    pub sender: mpsc::UnboundedSender<RegisterUpdate>,
229}
230
231pub(crate) struct QueueEntry {
232    pub token: DeliveryToken,
233    pub message: Message,
234    pub delivery_count: u32,
235    pub available_at: Instant,
236    pub in_flight: bool,
237}
238
239impl State {
240    pub fn next_revision(&mut self) -> Revision {
241        self.register_revision += 1;
242        Revision(self.register_revision)
243    }
244
245    pub fn notify_watchers(&mut self, update: RegisterUpdate) {
246        let key = match &update {
247            RegisterUpdate::Put(entry) => &entry.key,
248            RegisterUpdate::Delete(tombstone) => &tombstone.key,
249        };
250        self.watchers.retain(|watcher| {
251            !key.starts_with(&watcher.prefix) || watcher.sender.send(update.clone()).is_ok()
252        });
253    }
254}
cortexkit_bus_inmemory/backend.rs

cortexkit_bus_inmemory/
backend.rs
