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