xmtp_proto/traits/
stream.rs

1//! Default XMTP Stream
2
3use std::{
4    marker::PhantomData,
5    pin::Pin,
6    task::{Context, Poll, ready},
7};
8
9use crate::{
10    ApiEndpoint,
11    api::{ApiClientError, BytesStream},
12};
13use futures::{Stream, TryStream};
14use pin_project::pin_project;
15
16mod fake_empty;
17pub use fake_empty::*;
18
19#[pin_project]
20/// A stream which maps the tonic error to ApiClientError, and attaches endpoint metadata
21pub struct XmtpStream<T> {
22    #[pin]
23    inner: BytesStream,
24    endpoint: ApiEndpoint,
25    _marker: PhantomData<T>,
26}
27
28impl<T> XmtpStream<T> {
29    pub fn new(inner: BytesStream, endpoint: ApiEndpoint) -> Self {
30        Self {
31            inner,
32            endpoint,
33            _marker: PhantomData,
34        }
35    }
36}
37
38impl<T> Stream for XmtpStream<T>
39where
40    T: prost::Message + Default,
41{
42    type Item = Result<T, ApiClientError>;
43
44    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
45        let this = self.as_mut().project();
46        if let Some(item) = ready!(this.inner.try_poll_next(cx)) {
47            let res = item
48                .map_err(|e| e.endpoint(self.endpoint.clone()))
49                .and_then(|i| T::decode(i).map_err(ApiClientError::DecodeError));
50            Poll::Ready(Some(res))
51        } else {
52            Poll::Ready(None)
53        }
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use crate::api::NetworkError;
60
61    use super::*;
62    use bytes::Bytes;
63    use futures::{StreamExt, pin_mut, stream};
64    use prost::Message;
65    use xmtp_common::RetryableError;
66
67    #[derive(prost::Message)]
68    struct TestMessage {
69        #[prost(string, tag = "1")]
70        content: String,
71    }
72
73    #[xmtp_common::test]
74    async fn test_poll_next_successful_decode() {
75        let test_message = TestMessage {
76            content: "test content".to_string(),
77        };
78        let encoded_bytes = test_message.encode_to_vec();
79
80        let inner_stream = BytesStream::new(stream::once(async move {
81            Ok::<Bytes, ApiClientError>(Bytes::from(encoded_bytes))
82        }));
83        let xmtp_stream =
84            XmtpStream::<TestMessage>::new(inner_stream, ApiEndpoint::SubscribeGroupMessages);
85        pin_mut!(xmtp_stream);
86
87        let result = xmtp_stream.next().await.unwrap();
88        assert!(result.is_ok());
89        let decoded_message = result.unwrap();
90        assert_eq!(decoded_message.content, "test content");
91        // stream ends
92        let n = xmtp_stream.next().await;
93        assert!(n.is_none());
94    }
95
96    #[derive(thiserror::Error, Debug)]
97    enum TestError {
98        #[error("test error")]
99        DummyError,
100    }
101
102    impl RetryableError for TestError {
103        fn is_retryable(&self) -> bool {
104            false
105        }
106    }
107
108    #[xmtp_common::test]
109    async fn test_poll_next_error_mapping() {
110        let inner_stream = BytesStream::new(stream::once(async {
111            Err::<Bytes, ApiClientError>(ApiClientError::Client {
112                source: NetworkError::new(TestError::DummyError),
113            })
114        }));
115        let xmtp_stream =
116            XmtpStream::<TestMessage>::new(inner_stream, ApiEndpoint::SubscribeGroupMessages);
117        pin_mut!(xmtp_stream);
118
119        let result = xmtp_stream.next().await.unwrap();
120        assert!(result.is_err());
121
122        match result {
123            Err(ApiClientError::ClientWithEndpoint { endpoint, .. }) => {
124                assert_eq!(endpoint, "subscribe_group_messages");
125            }
126            _ => panic!("Expected ClientWithEndpoint error"),
127        }
128        // stream ends
129        let n = xmtp_stream.next().await;
130        assert!(n.is_none());
131    }
132}
xmtp_proto/traits/stream.rs

xmtp_proto/traits/
stream.rs
