#![allow(missing_docs)]
use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::Arc;
use async_trait::async_trait;
use futures::stream::{self, Stream};
use tokio::sync::{Mutex, Notify, RwLock};
use crate::event::{EventVersion, StreamEvent};
use crate::handler::NStreamsHandler;
use crate::namespace::Namespace;
use crate::publisher::EventPublisher;
use crate::store::EventStore;
use crate::stream::{ReadStreamBackend, WriteQueueBackend};
use crate::worker::{QueueConsumer, QueueMessage, StreamWorker};
#[derive(Clone, Default)]
pub struct InMemoryStore {
inner: Arc<Mutex<HashMap<String, Vec<StreamEvent>>>>,
}
#[async_trait]
impl EventStore for InMemoryStore {
async fn ensure_namespace<N: Namespace>(&self, namespace: &N) -> crate::Result<()> {
let mut inner = self.inner.lock().await;
inner.entry(namespace.as_str().to_string()).or_default();
Ok(())
}
async fn persist_next<N: Namespace>(
&self,
namespace: &N,
payload: &[u8],
filter_value: Option<&str>,
) -> crate::Result<StreamEvent> {
let mut inner = self.inner.lock().await;
let events = inner.entry(namespace.as_str().to_string()).or_default();
let version = events.last().map(|event| event.version + 1).unwrap_or(1);
let event = StreamEvent {
namespace: namespace.as_str().to_string(),
version,
payload: payload.to_vec(),
filter_value: filter_value.map(str::to_string),
};
events.push(event.clone());
Ok(event)
}
async fn load_history<N: Namespace>(
&self,
namespace: &N,
limit: u64,
) -> crate::Result<Vec<StreamEvent>> {
let inner = self.inner.lock().await;
let Some(events) = inner.get(namespace.as_str()) else {
return Ok(Vec::new());
};
let mut history = events
.iter()
.rev()
.take(limit as usize)
.cloned()
.collect::<Vec<_>>();
history.sort_by_key(|event| event.version);
Ok(history)
}
async fn load_after_version<N: Namespace>(
&self,
namespace: &N,
after_version: EventVersion,
limit: u64,
) -> crate::Result<Vec<StreamEvent>> {
let inner = self.inner.lock().await;
let Some(events) = inner.get(namespace.as_str()) else {
return Ok(Vec::new());
};
Ok(events
.iter()
.filter(|event| event.version > after_version)
.take(limit as usize)
.cloned()
.collect())
}
async fn latest_version<N: Namespace>(
&self,
namespace: &N,
) -> crate::Result<Option<EventVersion>> {
let inner = self.inner.lock().await;
Ok(inner
.get(namespace.as_str())
.and_then(|events| events.last().map(|event| event.version)))
}
}
#[derive(Clone, Default)]
pub struct InMemoryWriteQueue {
messages: Arc<Mutex<VecDeque<QueueMessage>>>,
notify: Arc<Notify>,
}
impl InMemoryWriteQueue {
pub fn consumer(self) -> InMemoryQueueConsumer {
InMemoryQueueConsumer {
messages: self.messages,
notify: self.notify,
}
}
}
#[async_trait]
impl WriteQueueBackend for InMemoryWriteQueue {
async fn publish<N: Namespace>(
&self,
namespace: &N,
payload: &[u8],
filter_value: Option<&str>,
) -> crate::Result<()> {
self.messages.lock().await.push_back(QueueMessage {
namespace: namespace.as_str().to_string(),
payload: payload.to_vec(),
filter_value: filter_value.map(str::to_string),
});
self.notify.notify_one();
Ok(())
}
}
pub struct InMemoryQueueConsumer {
messages: Arc<Mutex<VecDeque<QueueMessage>>>,
notify: Arc<Notify>,
}
#[async_trait]
impl QueueConsumer for InMemoryQueueConsumer {
async fn next(&mut self) -> Option<std::result::Result<QueueMessage, crate::Error>> {
loop {
if let Some(message) = self.messages.lock().await.pop_front() {
return Some(Ok(message));
}
self.notify.notified().await;
}
}
}
#[derive(Clone, Default)]
pub struct InMemoryReadStream {
streams: Arc<RwLock<HashSet<String>>>,
populated: Arc<RwLock<HashMap<String, Vec<StreamEvent>>>>,
live_senders: Arc<Mutex<HashMap<String, tokio::sync::mpsc::UnboundedSender<StreamEvent>>>>,
}
impl InMemoryReadStream {
pub async fn push_live(&self, namespace: &str, event: StreamEvent) {
let senders = self.live_senders.lock().await;
if let Some(sender) = senders.get(namespace) {
let _ = sender.send(event);
}
}
}
#[async_trait]
impl ReadStreamBackend for InMemoryReadStream {
async fn stream_exists<N: Namespace>(&self, namespace: &N) -> crate::Result<bool> {
Ok(self
.streams
.read()
.await
.contains(namespace.as_str()))
}
async fn create_read_stream<N: Namespace>(
&self,
namespace: &N,
_max_event_bytes: u64,
) -> crate::Result<()> {
self.streams
.write()
.await
.insert(namespace.as_str().to_string());
self.live_senders
.lock()
.await
.entry(namespace.as_str().to_string())
.or_insert_with(|| {
let (sender, _) = tokio::sync::mpsc::unbounded_channel();
sender
});
Ok(())
}
async fn publish_to_stream(&self, event: &StreamEvent) -> crate::Result<()> {
self.push_live(&event.namespace, event.clone()).await;
Ok(())
}
async fn populate_stream(&self, events: &[StreamEvent]) -> crate::Result<()> {
if let Some(first) = events.first() {
self.populated
.write()
.await
.insert(first.namespace.clone(), events.to_vec());
}
Ok(())
}
async fn subscribe_live<N: Namespace>(
&self,
namespace: &N,
_filter: Option<&crate::filter::StreamFilter>,
) -> crate::Result<impl Stream<Item = crate::Result<StreamEvent>> + Send> {
let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
self.live_senders
.lock()
.await
.insert(namespace.as_str().to_string(), sender);
Ok(stream::unfold(receiver, |mut receiver| async {
receiver.recv().await.map(|event| (Ok(event), receiver))
}))
}
}
pub struct TestHarness {
pub store: InMemoryStore,
pub write_queue: InMemoryWriteQueue,
pub read_stream: InMemoryReadStream,
pub handler: NStreamsHandler<InMemoryStore, InMemoryReadStream>,
pub publisher: EventPublisher<InMemoryWriteQueue>,
}
impl TestHarness {
pub fn new() -> Self {
let store = InMemoryStore::default();
let write_queue = InMemoryWriteQueue::default();
let read_stream = InMemoryReadStream::default();
let handler = NStreamsHandler::new(store.clone(), read_stream.clone());
let publisher = EventPublisher::new(write_queue.clone());
Self {
store,
write_queue,
read_stream,
handler,
publisher,
}
}
pub fn worker(&self) -> StreamWorker<InMemoryStore, InMemoryWriteQueue, InMemoryReadStream> {
StreamWorker::new(
self.store.clone(),
self.write_queue.clone(),
self.read_stream.clone(),
self.handler.clone(),
)
}
}
pub fn payload(value: serde_json::Value) -> Vec<u8> {
serde_json::to_vec(&value).expect("test payload serializes")
}
pub async fn collect_delivered(
subscription: &mut crate::Subscription,
count: usize,
) -> Vec<crate::event::DeliveredEvent> {
let mut events = Vec::with_capacity(count);
for _ in 0..count {
let item = tokio::time::timeout(std::time::Duration::from_secs(1), subscription.next())
.await
.expect("timed out waiting for event")
.expect("subscription ended");
events.push(item.expect("delivered event"));
}
events
}