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, ErrorCode};
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                    let mut delivered = 0;
86                    while delivered < max_messages {
87                        match consumer.recv(std::time::Duration::from_secs(1)).await {
88                            Ok(msg) => {
89                                yield msg;
90                                delivered += 1;
91                            }
92                            Err(e) if e.code() == ErrorCode::Timeout => {}
93                            Err(e) => Err(e)?,
94                        }
95                    }
96                }
97                None => loop {
98                    match consumer.recv(std::time::Duration::from_secs(1)).await {
99                        Ok(msg) => yield msg,
100                        Err(e) if e.code() == ErrorCode::Timeout => {}
101                        Err(e) => Err(e)?,
102                    }
103                },
104            }
105        };
106        let stream: BoxStream<Message<T>> = Box::pin(stream);
107        Ok(stream)
108    }
109}
110
111/// Create a [`ConsumerStream`] that yields messages from the given consumer.
112#[must_use]
113pub fn consumer_as_stream<T: Send + Sync + 'static>(
114    name: &'static str,
115    consumer: Arc<dyn MessageConsumer<T>>,
116) -> ConsumerStream<T> {
117    ConsumerStream {
118        name,
119        consumer,
120        max_messages: None,
121    }
122}
123
124/// Create a [`ConsumerStream`] that yields at most `max_messages` messages.
125#[must_use]
126pub fn consumer_as_bounded_stream<T: Send + Sync + 'static>(
127    name: &'static str,
128    consumer: Arc<dyn MessageConsumer<T>>,
129    max_messages: usize,
130) -> ConsumerStream<T> {
131    ConsumerStream {
132        name,
133        consumer,
134        max_messages: Some(max_messages),
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use std::sync::Arc;
141    use std::time::Duration;
142
143    use futures::StreamExt;
144
145    use crate::memory::InMemoryBroker;
146    use crate::traits::{MessageConsumer, MessageProducer};
147
148    use super::*;
149
150    #[tokio::test]
151    async fn producer_sink_sends_message() {
152        let broker = InMemoryBroker::<String>::new(16);
153        let producer = Arc::new(broker.producer());
154        let consumer = broker.consumer();
155        consumer.subscribe(&["sink-topic"]).await.unwrap();
156
157        let sink = producer_as_sink("test-sink", producer, "sink-topic".into());
158        let msg = Message::new("sink-topic", "hello".into());
159        Sink::send(&sink, msg).await.unwrap();
160
161        let received = consumer
162            .recv(std::time::Duration::from_secs(1))
163            .await
164            .unwrap();
165        assert_eq!(received.topic, "sink-topic");
166        assert_eq!(received.payload, "hello");
167    }
168
169    #[tokio::test]
170    async fn producer_sink_applies_default_topic() {
171        let broker = InMemoryBroker::<String>::new(16);
172        let producer = Arc::new(broker.producer());
173        let consumer = broker.consumer();
174        consumer.subscribe(&["default-t"]).await.unwrap();
175
176        let sink = producer_as_sink("test-sink", producer, "default-t".into());
177        let mut msg = Message::new("", "data".into());
178        msg.topic = String::new();
179        Sink::send(&sink, msg).await.unwrap();
180
181        let received = consumer
182            .recv(std::time::Duration::from_secs(1))
183            .await
184            .unwrap();
185        assert_eq!(received.topic, "default-t");
186    }
187
188    #[tokio::test]
189    async fn consumer_stream_yields_messages() {
190        let broker = InMemoryBroker::<String>::new(16);
191        let producer = broker.producer();
192        let consumer = Arc::new(broker.consumer());
193        consumer.subscribe(&["stream-t"]).await.unwrap();
194
195        let cs = consumer_as_stream("test-stream", consumer);
196
197        // Send messages before opening the stream
198        producer
199            .send(Message::new("stream-t", "a".into()))
200            .await
201            .unwrap();
202        producer
203            .send(Message::new("stream-t", "b".into()))
204            .await
205            .unwrap();
206
207        let stream = cs.execute(()).await.unwrap();
208        let items: Vec<_> = tokio::time::timeout(Duration::from_millis(200), async {
209            stream.take(2).collect::<Vec<_>>().await
210        })
211        .await
212        .unwrap();
213
214        assert_eq!(items.len(), 2);
215        assert_eq!(items[0].as_ref().unwrap().payload, "a");
216        assert_eq!(items[1].as_ref().unwrap().payload, "b");
217    }
218
219    #[tokio::test]
220    async fn bounded_consumer_stream_completes_after_limit() {
221        let broker = InMemoryBroker::<String>::new(16);
222        let producer = broker.producer();
223        let consumer = Arc::new(broker.consumer());
224        consumer.subscribe(&["stream-t"]).await.unwrap();
225
226        let cs = consumer_as_bounded_stream("test-stream", consumer, 2);
227
228        for value in ["a", "b", "c"] {
229            producer
230                .send(Message::new("stream-t", value.to_owned()))
231                .await
232                .unwrap();
233        }
234
235        let stream = cs.execute(()).await.unwrap();
236        let items = stream.collect::<Vec<_>>().await;
237
238        assert_eq!(items.len(), 2);
239        assert_eq!(items[0].as_ref().unwrap().payload, "a");
240        assert_eq!(items[1].as_ref().unwrap().payload, "b");
241    }
242
243    #[tokio::test]
244    async fn bounded_consumer_stream_with_zero_limit_completes_immediately() {
245        let broker = InMemoryBroker::<String>::new(16);
246        let consumer = Arc::new(broker.consumer());
247        consumer.subscribe(&["stream-t"]).await.unwrap();
248
249        let cs = consumer_as_bounded_stream("test-stream", consumer, 0);
250        let stream = cs.execute(()).await.unwrap();
251        let items = tokio::time::timeout(Duration::from_millis(50), stream.collect::<Vec<_>>())
252            .await
253            .unwrap();
254
255        assert!(items.is_empty());
256    }
257
258    #[tokio::test(start_paused = true)]
259    async fn consumer_stream_survives_idle_recv_timeouts() {
260        let broker = InMemoryBroker::<String>::new(16);
261        let producer = broker.producer();
262        let consumer = Arc::new(broker.consumer());
263        consumer.subscribe(&["idle-t"]).await.unwrap();
264
265        let cs = consumer_as_stream("idle-stream", consumer);
266        let stream = cs.execute(()).await.unwrap();
267        let collector = tokio::spawn(async move { stream.take(1).collect::<Vec<_>>().await });
268
269        // Let several per-recv timeouts elapse with no traffic; the stream must
270        // keep polling instead of terminating on the first idle second.
271        tokio::time::sleep(Duration::from_secs(3)).await;
272        producer
273            .send(Message::new("idle-t", "late".into()))
274            .await
275            .unwrap();
276
277        let items = tokio::time::timeout(Duration::from_secs(5), collector)
278            .await
279            .expect("stream terminated instead of surviving idle recv timeouts")
280            .unwrap();
281        assert_eq!(items.len(), 1);
282        assert_eq!(items[0].as_ref().unwrap().payload, "late");
283    }
284
285    #[tokio::test]
286    async fn producer_sink_has_provider_name() {
287        let broker = InMemoryBroker::<String>::new(4);
288        let producer = Arc::new(broker.producer());
289        let sink = producer_as_sink("my-sink", producer, "t".into());
290        assert_eq!(Provider::name(&sink), "my-sink");
291    }
292
293    #[tokio::test]
294    async fn consumer_stream_has_provider_name() {
295        let broker = InMemoryBroker::<String>::new(4);
296        let consumer = Arc::new(broker.consumer());
297        let cs = consumer_as_stream("my-cs", consumer);
298        assert_eq!(Provider::name(&cs), "my-cs");
299    }
300}