Skip to main content

rskit_messaging/bridge/
provider.rs

1//! Bridge adapters from messaging traits to [`rskit_provider`] traits.
2//!
3//! - [`ProducerSink`] wraps a [`MessageProducer`] as a [`Sink`].
4//! - [`ConsumerStream`] wraps a [`MessageConsumer`] as a [`Stream`].
5
6use std::sync::Arc;
7
8use async_trait::async_trait;
9use rskit_errors::AppResult;
10use rskit_provider::traits::{BoxStream, Provider, Sink, Stream};
11
12use crate::message::Message;
13use crate::traits::{MessageConsumer, MessageProducer};
14
15// ─── ProducerSink ────────────────────────────────────────────────────────────
16
17/// Wraps a [`MessageProducer`] as a [`Sink<Message<T>>`](Sink).
18///
19/// Each call to [`Sink::send`] publishes the message to the configured topic.
20pub struct ProducerSink<T: Send + Sync + 'static> {
21    name: &'static str,
22    producer: Arc<dyn MessageProducer<T>>,
23    topic: String,
24}
25
26#[async_trait]
27impl<T: Send + Sync + 'static> Provider for ProducerSink<T> {
28    fn name(&self) -> &'static str {
29        self.name
30    }
31}
32
33impl<T: Send + Sync + 'static> Sink<Message<T>> for ProducerSink<T> {
34    async fn send(&self, mut input: Message<T>) -> AppResult<()> {
35        if input.topic.is_empty() {
36            input.topic.clone_from(&self.topic);
37        }
38        self.producer.send(input).await
39    }
40}
41
42/// Create a [`ProducerSink`] that publishes messages via the given producer.
43///
44/// The `topic` is applied as a default when a message has no topic set.
45#[must_use]
46pub fn producer_as_sink<T: Send + Sync + 'static>(
47    name: &'static str,
48    producer: Arc<dyn MessageProducer<T>>,
49    topic: String,
50) -> ProducerSink<T> {
51    ProducerSink {
52        name,
53        producer,
54        topic,
55    }
56}
57
58// ─── ConsumerStream ──────────────────────────────────────────────────────────
59
60/// Wraps a [`MessageConsumer`] as a [`Stream<(), Message<T>>`](Stream).
61///
62/// Calling [`Stream::execute`] returns a backpressure-aware stream that receives
63/// one message at a time. Use [`consumer_as_bounded_stream`] when the stream
64/// should complete after a fixed number of messages.
65pub struct ConsumerStream<T: Send + Sync + 'static> {
66    name: &'static str,
67    consumer: Arc<dyn MessageConsumer<T>>,
68    max_messages: Option<usize>,
69}
70
71#[async_trait]
72impl<T: Send + Sync + 'static> Provider for ConsumerStream<T> {
73    fn name(&self) -> &'static str {
74        self.name
75    }
76}
77
78impl<T: Send + Sync + Clone + 'static> Stream<(), Message<T>> for ConsumerStream<T> {
79    async fn execute(&self, _input: ()) -> AppResult<BoxStream<Message<T>>> {
80        let consumer = Arc::clone(&self.consumer);
81        let max_messages = self.max_messages;
82        let stream = async_stream::try_stream! {
83            match max_messages {
84                Some(max_messages) => {
85                    for _ in 0..max_messages {
86                        let msg = consumer.recv().await?;
87                        yield msg;
88                    }
89                }
90                None => loop {
91                    let msg = consumer.recv().await?;
92                    yield msg;
93                },
94            }
95        };
96        let stream: BoxStream<Message<T>> = Box::pin(stream);
97        Ok(stream)
98    }
99}
100
101/// Create a [`ConsumerStream`] that yields messages from the given consumer.
102#[must_use]
103pub fn consumer_as_stream<T: Send + Sync + 'static>(
104    name: &'static str,
105    consumer: Arc<dyn MessageConsumer<T>>,
106) -> ConsumerStream<T> {
107    ConsumerStream {
108        name,
109        consumer,
110        max_messages: None,
111    }
112}
113
114/// Create a [`ConsumerStream`] that yields at most `max_messages` messages.
115#[must_use]
116pub fn consumer_as_bounded_stream<T: Send + Sync + 'static>(
117    name: &'static str,
118    consumer: Arc<dyn MessageConsumer<T>>,
119    max_messages: usize,
120) -> ConsumerStream<T> {
121    ConsumerStream {
122        name,
123        consumer,
124        max_messages: Some(max_messages),
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use std::sync::Arc;
131    use std::time::Duration;
132
133    use futures::StreamExt;
134
135    use crate::memory::InMemoryBroker;
136    use crate::traits::{MessageConsumer, MessageProducer};
137
138    use super::*;
139
140    #[tokio::test]
141    async fn producer_sink_sends_message() {
142        let broker = InMemoryBroker::<String>::new(16);
143        let producer = Arc::new(broker.producer());
144        let consumer = broker.consumer();
145        consumer.subscribe(&["sink-topic"]).await.unwrap();
146
147        let sink = producer_as_sink("test-sink", producer, "sink-topic".into());
148        let msg = Message::new("sink-topic", "hello".into());
149        Sink::send(&sink, msg).await.unwrap();
150
151        let received = consumer.recv().await.unwrap();
152        assert_eq!(received.topic, "sink-topic");
153        assert_eq!(received.payload, "hello");
154    }
155
156    #[tokio::test]
157    async fn producer_sink_applies_default_topic() {
158        let broker = InMemoryBroker::<String>::new(16);
159        let producer = Arc::new(broker.producer());
160        let consumer = broker.consumer();
161        consumer.subscribe(&["default-t"]).await.unwrap();
162
163        let sink = producer_as_sink("test-sink", producer, "default-t".into());
164        let mut msg = Message::new("", "data".into());
165        msg.topic = String::new();
166        Sink::send(&sink, msg).await.unwrap();
167
168        let received = consumer.recv().await.unwrap();
169        assert_eq!(received.topic, "default-t");
170    }
171
172    #[tokio::test]
173    async fn consumer_stream_yields_messages() {
174        let broker = InMemoryBroker::<String>::new(16);
175        let producer = broker.producer();
176        let consumer = Arc::new(broker.consumer());
177        consumer.subscribe(&["stream-t"]).await.unwrap();
178
179        let cs = consumer_as_stream("test-stream", consumer);
180
181        // Send messages before opening the stream
182        producer
183            .send(Message::new("stream-t", "a".into()))
184            .await
185            .unwrap();
186        producer
187            .send(Message::new("stream-t", "b".into()))
188            .await
189            .unwrap();
190
191        let stream = cs.execute(()).await.unwrap();
192        let items: Vec<_> = tokio::time::timeout(Duration::from_millis(200), async {
193            stream.take(2).collect::<Vec<_>>().await
194        })
195        .await
196        .unwrap();
197
198        assert_eq!(items.len(), 2);
199        assert_eq!(items[0].as_ref().unwrap().payload, "a");
200        assert_eq!(items[1].as_ref().unwrap().payload, "b");
201    }
202
203    #[tokio::test]
204    async fn bounded_consumer_stream_completes_after_limit() {
205        let broker = InMemoryBroker::<String>::new(16);
206        let producer = broker.producer();
207        let consumer = Arc::new(broker.consumer());
208        consumer.subscribe(&["stream-t"]).await.unwrap();
209
210        let cs = consumer_as_bounded_stream("test-stream", consumer, 2);
211
212        for value in ["a", "b", "c"] {
213            producer
214                .send(Message::new("stream-t", value.to_owned()))
215                .await
216                .unwrap();
217        }
218
219        let stream = cs.execute(()).await.unwrap();
220        let items = stream.collect::<Vec<_>>().await;
221
222        assert_eq!(items.len(), 2);
223        assert_eq!(items[0].as_ref().unwrap().payload, "a");
224        assert_eq!(items[1].as_ref().unwrap().payload, "b");
225    }
226
227    #[tokio::test]
228    async fn bounded_consumer_stream_with_zero_limit_completes_immediately() {
229        let broker = InMemoryBroker::<String>::new(16);
230        let consumer = Arc::new(broker.consumer());
231        consumer.subscribe(&["stream-t"]).await.unwrap();
232
233        let cs = consumer_as_bounded_stream("test-stream", consumer, 0);
234        let stream = cs.execute(()).await.unwrap();
235        let items = tokio::time::timeout(Duration::from_millis(50), stream.collect::<Vec<_>>())
236            .await
237            .unwrap();
238
239        assert!(items.is_empty());
240    }
241
242    #[tokio::test]
243    async fn producer_sink_has_provider_name() {
244        let broker = InMemoryBroker::<String>::new(4);
245        let producer = Arc::new(broker.producer());
246        let sink = producer_as_sink("my-sink", producer, "t".into());
247        assert_eq!(Provider::name(&sink), "my-sink");
248    }
249
250    #[tokio::test]
251    async fn consumer_stream_has_provider_name() {
252        let broker = InMemoryBroker::<String>::new(4);
253        let consumer = Arc::new(broker.consumer());
254        let cs = consumer_as_stream("my-cs", consumer);
255        assert_eq!(Provider::name(&cs), "my-cs");
256    }
257}