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