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}
139
140impl InMemoryBus {
141    pub fn new(config: InMemoryConfig) -> Self {
142        Self::new_with_deduplication_clock(config, Instant::now)
143    }
144
145    /// Injects a monotonic clock for deduplication without waiting in tests.
146    /// Delivery leases and work-queue delays still use the real clock.
147    /// Each bus has one stream; clones share its ids, while separate buses do not.
148    pub fn new_with_deduplication_clock(
149        config: InMemoryConfig,
150        clock: impl Fn() -> Instant + Send + Sync + 'static,
151    ) -> Self {
152        let durable_cursors = config
153            .durables
154            .iter()
155            .cloned()
156            .map(|key| (key, DurableState::default()))
157            .collect();
158        Self {
159            inner: Arc::new(Mutex::new(State {
160                config,
161                next_stream_seq: 1,
162                stream_messages: Vec::new(),
163                seen_ids: HashMap::new(),
164                durable_cursors,
165                register_revision: 0,
166                register_values: BTreeMap::new(),
167                watchers: Vec::new(),
168                queue: VecDeque::new(),
169                next_delivery_token: 1,
170                events: Vec::new(),
171                permission_violations: VecDeque::new(),
172            })),
173            deduplication_clock: Arc::new(clock),
174        }
175    }
176
177    pub fn events(&self) -> Vec<BackendEvent> {
178        self.inner
179            .lock()
180            .expect("in-memory bus poisoned")
181            .events
182            .clone()
183    }
184
185    pub fn queued_bytes(&self) -> usize {
186        self.inner
187            .lock()
188            .expect("in-memory bus poisoned")
189            .queue
190            .iter()
191            .map(|entry| entry.message.wire_size())
192            .sum()
193    }
194
195    pub fn inject_permission_violation(&self, subject: impl Into<String>) {
196        let subject = subject.into();
197        self.inner
198            .lock()
199            .expect("in-memory bus poisoned")
200            .permission_violations
201            .push_back(AsyncPermissionViolation {
202                subject,
203                reason: "fixture-configured permission violation".into(),
204            });
205    }
206}
207
208pub(crate) struct State {
209    pub config: InMemoryConfig,
210    pub next_stream_seq: u64,
211    pub stream_messages: Vec<StoredMessage>,
212    pub seen_ids: HashMap<String, (u64, Instant)>,
213    pub durable_cursors: HashMap<String, DurableState>,
214    pub register_revision: u64,
215    pub register_values: BTreeMap<String, RegisterValue>,
216    pub watchers: Vec<Watcher>,
217    pub queue: VecDeque<QueueEntry>,
218    pub next_delivery_token: u64,
219    pub events: Vec<BackendEvent>,
220    pub permission_violations: VecDeque<AsyncPermissionViolation>,
221}
222
223#[derive(Clone)]
224pub(crate) struct StoredMessage {
225    pub stream_seq: u64,
226    pub message: Message,
227}
228
229/// One durable's position. Delivery is strictly in order: the message at
230/// `next_index` is the only one a cursor can receive until it is acked or
231/// terminated.
232#[derive(Default)]
233pub(crate) struct DurableState {
234    pub next_index: usize,
235    /// Set by a nak: no cursor receives the head message before this instant.
236    pub available_at: Option<Instant>,
237    /// How many times the head message has been delivered.
238    pub delivery_count: u64,
239    /// The cursor holding the head message unsettled, if any.
240    pub lease: Option<Lease>,
241    pub next_cursor_id: u64,
242}
243
244/// An unsettled delivery of the head message. Until `deadline` no other
245/// cursor receives the message; after it the next `next` on any cursor
246/// redelivers it, as JetStream does once the ack wait expires.
247pub(crate) struct Lease {
248    pub cursor_id: u64,
249    pub deadline: Instant,
250}
251
252pub(crate) struct Watcher {
253    pub prefix: String,
254    pub sender: mpsc::UnboundedSender<RegisterUpdate>,
255}
256
257pub(crate) struct QueueEntry {
258    pub token: DeliveryToken,
259    pub message: Message,
260    pub delivery_count: u32,
261    pub available_at: Instant,
262    pub in_flight: bool,
263}
264
265impl State {
266    pub fn next_revision(&mut self) -> Revision {
267        self.register_revision += 1;
268        Revision(self.register_revision)
269    }
270
271    pub fn notify_watchers(&mut self, update: RegisterUpdate) {
272        let key = match &update {
273            RegisterUpdate::Put(entry) => &entry.key,
274            RegisterUpdate::Delete(tombstone) => &tombstone.key,
275        };
276        self.watchers.retain(|watcher| {
277            !key.starts_with(&watcher.prefix) || watcher.sender.send(update.clone()).is_ok()
278        });
279    }
280}
cortexkit_bus_inmemory/backend.rs

cortexkit_bus_inmemory/
backend.rs
