use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, Mutex};
use super::source::{AsyncMessageSource, ReceivedMessage};
use super::{run_source, Bus, BusConsumer, MessageRouter, RunOptions, TransportError};
use super::{Message, MessageKind};
type Queues = Arc<Mutex<HashMap<String, VecDeque<Message>>>>;
type Topics = Arc<Mutex<HashMap<String, Vec<Message>>>>;
fn lock_poisoned(what: &str) -> TransportError {
TransportError::permanent(format!("in-memory bus {what} lock poisoned"))
}
#[derive(Clone, Default)]
pub struct InMemoryBus {
queues: Queues,
topics: Topics,
}
impl InMemoryBus {
pub fn new() -> Self {
Self::default()
}
fn enqueue(&self, message: Message) -> Result<(), TransportError> {
self.queues
.lock()
.map_err(|_| lock_poisoned("queue"))?
.entry(message.name().to_string())
.or_default()
.push_back(message);
Ok(())
}
fn append(&self, message: Message) -> Result<(), TransportError> {
self.topics
.lock()
.map_err(|_| lock_poisoned("topic"))?
.entry(message.name().to_string())
.or_default()
.push(message);
Ok(())
}
}
impl Bus for InMemoryBus {
async fn send(&self, name: &str, payload: Vec<u8>) -> Result<(), TransportError> {
self.enqueue(Message::new(name, MessageKind::Command, payload))
}
async fn publish(&self, name: &str, payload: Vec<u8>) -> Result<(), TransportError> {
self.append(Message::new(name, MessageKind::Event, payload))
}
async fn send_message(&self, message: Message) -> Result<(), TransportError> {
self.enqueue(message)
}
async fn publish_message(&self, message: Message) -> Result<(), TransportError> {
self.append(message)
}
}
impl BusConsumer for InMemoryBus {
async fn listen<R: MessageRouter>(
&self,
router: Arc<R>,
options: RunOptions,
) -> Result<(), TransportError> {
let names = router.subscription_plan().commands;
let source = QueueSource {
queues: self.queues.clone(),
names,
};
run_source(router, source, options).await
}
async fn subscribe<R: MessageRouter>(
&self,
router: Arc<R>,
options: RunOptions,
) -> Result<(), TransportError> {
let names = router.subscription_plan().events;
let source = TopicSource {
topics: self.topics.clone(),
names,
cursors: HashMap::new(),
};
run_source(router, source, options).await
}
}
struct QueueSource {
queues: Queues,
names: Vec<String>,
}
impl AsyncMessageSource for QueueSource {
type Received = InMemoryReceived;
async fn recv(&mut self) -> Result<Option<Self::Received>, TransportError> {
let mut queues = self.queues.lock().map_err(|_| lock_poisoned("queue"))?;
for name in &self.names {
if let Some(message) = queues.get_mut(name).and_then(VecDeque::pop_front) {
return Ok(Some(InMemoryReceived { message }));
}
}
Ok(None)
}
}
struct TopicSource {
topics: Topics,
names: Vec<String>,
cursors: HashMap<String, usize>,
}
impl AsyncMessageSource for TopicSource {
type Received = InMemoryReceived;
async fn recv(&mut self) -> Result<Option<Self::Received>, TransportError> {
let topics = self.topics.lock().map_err(|_| lock_poisoned("topic"))?;
for name in &self.names {
let Some(log) = topics.get(name) else {
continue;
};
let cursor = self.cursors.entry(name.clone()).or_insert(0);
if *cursor < log.len() {
let message = log[*cursor].clone();
*cursor += 1;
return Ok(Some(InMemoryReceived { message }));
}
}
Ok(None)
}
}
pub struct InMemoryReceived {
message: Message,
}
impl ReceivedMessage for InMemoryReceived {
fn message(&self) -> &Message {
&self.message
}
async fn ack(self) -> Result<(), TransportError> {
Ok(())
}
async fn nack(self, _reason: &str) -> Result<(), TransportError> {
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::bus::Handlers;
use std::future::Future;
fn block_on<F: Future>(future: F) -> F::Output {
use std::ptr;
use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
const VTABLE: RawWakerVTable = RawWakerVTable::new(
|_| RawWaker::new(ptr::null(), &VTABLE),
|_| {},
|_| {},
|_| {},
);
let waker = unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &VTABLE)) };
let mut cx = Context::from_waker(&waker);
let mut future = std::pin::pin!(future);
loop {
if let Poll::Ready(output) = future.as_mut().poll(&mut cx) {
return output;
}
}
}
fn recorder() -> Arc<Mutex<Vec<String>>> {
Arc::new(Mutex::new(Vec::new()))
}
fn command_service(rec: Arc<Mutex<Vec<String>>>) -> Arc<Handlers> {
Arc::new(Handlers::new().on_command("work", move |msg: &Message| {
let rec = rec.clone();
let name = msg.name().to_string();
async move {
rec.lock().unwrap().push(name);
Ok(())
}
}))
}
fn event_service(rec: Arc<Mutex<Vec<String>>>) -> Arc<Handlers> {
Arc::new(Handlers::new().on_event("evt", move |msg: &Message| {
let rec = rec.clone();
let id = msg.id().unwrap_or("?").to_string();
async move {
rec.lock().unwrap().push(id);
Ok(())
}
}))
}
#[test]
fn send_then_listen_dispatches_each_command() {
let bus = InMemoryBus::new();
for _ in 0..3 {
block_on(bus.send("work", b"{}".to_vec())).unwrap();
}
let rec = recorder();
block_on(bus.listen(command_service(rec.clone()), RunOptions::idempotent())).unwrap();
assert_eq!(
rec.lock().unwrap().len(),
3,
"the listener handles all 3 commands"
);
}
#[test]
fn listen_is_point_to_point_each_message_popped_once() {
let bus = InMemoryBus::new();
for i in 0..4 {
block_on(bus.send_message(
Message::new("work", MessageKind::Command, b"{}".to_vec()).with_id(format!("m{i}")),
))
.unwrap();
}
let mut a = QueueSource {
queues: bus.queues.clone(),
names: vec!["work".to_string()],
};
let mut b = QueueSource {
queues: bus.queues.clone(),
names: vec!["work".to_string()],
};
let mut got = Vec::new();
for _ in 0..4 {
if let Some(r) = block_on(a.recv()).unwrap() {
got.push(r.message().id().unwrap().to_string());
}
if let Some(r) = block_on(b.recv()).unwrap() {
got.push(r.message().id().unwrap().to_string());
}
}
got.sort();
assert_eq!(
got,
vec!["m0", "m1", "m2", "m3"],
"each message delivered exactly once"
);
assert!(block_on(a.recv()).unwrap().is_none());
assert!(block_on(b.recv()).unwrap().is_none());
}
#[test]
fn publish_then_subscribe_fans_out_to_every_subscriber() {
let bus = InMemoryBus::new();
for i in 0..3 {
block_on(bus.publish_message(
Message::new("evt", MessageKind::Event, b"{}".to_vec()).with_id(format!("e{i}")),
))
.unwrap();
}
let a = recorder();
let b = recorder();
block_on(bus.subscribe(event_service(a.clone()), RunOptions::idempotent())).unwrap();
block_on(bus.subscribe(event_service(b.clone()), RunOptions::idempotent())).unwrap();
let mut a_ids = a.lock().unwrap().clone();
let mut b_ids = b.lock().unwrap().clone();
a_ids.sort();
b_ids.sort();
assert_eq!(a_ids, vec!["e0", "e1", "e2"]);
assert_eq!(b_ids, vec!["e0", "e1", "e2"]);
}
#[test]
fn unknown_command_is_acked_and_ignored() {
let bus = InMemoryBus::new();
block_on(bus.send("unrelated", b"{}".to_vec())).unwrap();
block_on(bus.send("work", b"{}".to_vec())).unwrap();
let rec = recorder();
block_on(bus.listen(command_service(rec.clone()), RunOptions::idempotent())).unwrap();
assert_eq!(rec.lock().unwrap().clone(), vec!["work"]);
}
#[test]
fn handler_error_does_not_panic_the_loop() {
let bus = InMemoryBus::new();
block_on(bus.send("work", b"{}".to_vec())).unwrap();
let handlers: Arc<Handlers> = Arc::new(
Handlers::new().on_command("work", |_: &Message| async move {
Err(TransportError::permanent("no"))
}),
);
block_on(bus.listen(handlers, RunOptions::idempotent())).unwrap();
}
}