use std::sync::Arc;
use async_trait::async_trait;
use rskit_errors::AppResult;
use rskit_provider::traits::{BoxStream, Provider, Sink, Stream};
use crate::message::Message;
use crate::traits::{MessageConsumer, MessageProducer};
pub struct ProducerSink<T: Send + Sync + 'static> {
name: &'static str,
producer: Arc<dyn MessageProducer<T>>,
topic: String,
}
#[async_trait]
impl<T: Send + Sync + 'static> Provider for ProducerSink<T> {
fn name(&self) -> &'static str {
self.name
}
}
impl<T: Send + Sync + 'static> Sink<Message<T>> for ProducerSink<T> {
async fn send(&self, mut input: Message<T>) -> AppResult<()> {
if input.topic.is_empty() {
input.topic.clone_from(&self.topic);
}
self.producer.send(input).await
}
}
#[must_use]
pub fn producer_as_sink<T: Send + Sync + 'static>(
name: &'static str,
producer: Arc<dyn MessageProducer<T>>,
topic: String,
) -> ProducerSink<T> {
ProducerSink {
name,
producer,
topic,
}
}
pub struct ConsumerStream<T: Send + Sync + 'static> {
name: &'static str,
consumer: Arc<dyn MessageConsumer<T>>,
max_messages: Option<usize>,
}
#[async_trait]
impl<T: Send + Sync + 'static> Provider for ConsumerStream<T> {
fn name(&self) -> &'static str {
self.name
}
}
impl<T: Send + Sync + Clone + 'static> Stream<(), Message<T>> for ConsumerStream<T> {
async fn execute(&self, _input: ()) -> AppResult<BoxStream<Message<T>>> {
let consumer = Arc::clone(&self.consumer);
let max_messages = self.max_messages;
let stream = async_stream::try_stream! {
match max_messages {
Some(max_messages) => {
for _ in 0..max_messages {
let msg = consumer.recv().await?;
yield msg;
}
}
None => loop {
let msg = consumer.recv().await?;
yield msg;
},
}
};
let stream: BoxStream<Message<T>> = Box::pin(stream);
Ok(stream)
}
}
#[must_use]
pub fn consumer_as_stream<T: Send + Sync + 'static>(
name: &'static str,
consumer: Arc<dyn MessageConsumer<T>>,
) -> ConsumerStream<T> {
ConsumerStream {
name,
consumer,
max_messages: None,
}
}
#[must_use]
pub fn consumer_as_bounded_stream<T: Send + Sync + 'static>(
name: &'static str,
consumer: Arc<dyn MessageConsumer<T>>,
max_messages: usize,
) -> ConsumerStream<T> {
ConsumerStream {
name,
consumer,
max_messages: Some(max_messages),
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::time::Duration;
use futures::StreamExt;
use crate::memory::InMemoryBroker;
use crate::traits::{MessageConsumer, MessageProducer};
use super::*;
#[tokio::test]
async fn producer_sink_sends_message() {
let broker = InMemoryBroker::<String>::new(16);
let producer = Arc::new(broker.producer());
let consumer = broker.consumer();
consumer.subscribe(&["sink-topic"]).await.unwrap();
let sink = producer_as_sink("test-sink", producer, "sink-topic".into());
let msg = Message::new("sink-topic", "hello".into());
Sink::send(&sink, msg).await.unwrap();
let received = consumer.recv().await.unwrap();
assert_eq!(received.topic, "sink-topic");
assert_eq!(received.payload, "hello");
}
#[tokio::test]
async fn producer_sink_applies_default_topic() {
let broker = InMemoryBroker::<String>::new(16);
let producer = Arc::new(broker.producer());
let consumer = broker.consumer();
consumer.subscribe(&["default-t"]).await.unwrap();
let sink = producer_as_sink("test-sink", producer, "default-t".into());
let mut msg = Message::new("", "data".into());
msg.topic = String::new();
Sink::send(&sink, msg).await.unwrap();
let received = consumer.recv().await.unwrap();
assert_eq!(received.topic, "default-t");
}
#[tokio::test]
async fn consumer_stream_yields_messages() {
let broker = InMemoryBroker::<String>::new(16);
let producer = broker.producer();
let consumer = Arc::new(broker.consumer());
consumer.subscribe(&["stream-t"]).await.unwrap();
let cs = consumer_as_stream("test-stream", consumer);
producer
.send(Message::new("stream-t", "a".into()))
.await
.unwrap();
producer
.send(Message::new("stream-t", "b".into()))
.await
.unwrap();
let stream = cs.execute(()).await.unwrap();
let items: Vec<_> = tokio::time::timeout(Duration::from_millis(200), async {
stream.take(2).collect::<Vec<_>>().await
})
.await
.unwrap();
assert_eq!(items.len(), 2);
assert_eq!(items[0].as_ref().unwrap().payload, "a");
assert_eq!(items[1].as_ref().unwrap().payload, "b");
}
#[tokio::test]
async fn bounded_consumer_stream_completes_after_limit() {
let broker = InMemoryBroker::<String>::new(16);
let producer = broker.producer();
let consumer = Arc::new(broker.consumer());
consumer.subscribe(&["stream-t"]).await.unwrap();
let cs = consumer_as_bounded_stream("test-stream", consumer, 2);
for value in ["a", "b", "c"] {
producer
.send(Message::new("stream-t", value.to_owned()))
.await
.unwrap();
}
let stream = cs.execute(()).await.unwrap();
let items = stream.collect::<Vec<_>>().await;
assert_eq!(items.len(), 2);
assert_eq!(items[0].as_ref().unwrap().payload, "a");
assert_eq!(items[1].as_ref().unwrap().payload, "b");
}
#[tokio::test]
async fn bounded_consumer_stream_with_zero_limit_completes_immediately() {
let broker = InMemoryBroker::<String>::new(16);
let consumer = Arc::new(broker.consumer());
consumer.subscribe(&["stream-t"]).await.unwrap();
let cs = consumer_as_bounded_stream("test-stream", consumer, 0);
let stream = cs.execute(()).await.unwrap();
let items = tokio::time::timeout(Duration::from_millis(50), stream.collect::<Vec<_>>())
.await
.unwrap();
assert!(items.is_empty());
}
#[tokio::test]
async fn producer_sink_has_provider_name() {
let broker = InMemoryBroker::<String>::new(4);
let producer = Arc::new(broker.producer());
let sink = producer_as_sink("my-sink", producer, "t".into());
assert_eq!(Provider::name(&sink), "my-sink");
}
#[tokio::test]
async fn consumer_stream_has_provider_name() {
let broker = InMemoryBroker::<String>::new(4);
let consumer = Arc::new(broker.consumer());
let cs = consumer_as_stream("my-cs", consumer);
assert_eq!(Provider::name(&cs), "my-cs");
}
}