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

cortexkit_bus_inmemory/
backend.rs
