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> {
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#[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#[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 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 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}