rskit_messaging/bridge/
provider.rs1use 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
15pub 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#[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
58pub 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#[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#[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 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}