rustrade_integration/stream/
exchange.rs

1use crate::{Transformer, error::SocketError, protocol::StreamParser};
2use futures::Stream;
3use pin_project::pin_project;
4use std::{
5    collections::VecDeque,
6    marker::PhantomData,
7    pin::Pin,
8    task::{Context, Poll},
9};
10
11/// An [`ExchangeStream`] is a communication protocol agnostic [`Stream`]. It polls protocol
12/// messages from the inner [`Stream`], and transforms them into the desired output data structure.
13///
14/// Requires both the `stream` and `protocol` features.
15#[derive(Debug)]
16#[pin_project]
17pub struct ExchangeStream<Protocol, InnerStream, StreamTransformer>
18where
19    Protocol: StreamParser<StreamTransformer::Input>,
20    InnerStream: Stream,
21    StreamTransformer: Transformer,
22{
23    #[pin]
24    pub stream: InnerStream,
25    pub transformer: StreamTransformer,
26    pub buffer: VecDeque<Result<StreamTransformer::Output, StreamTransformer::Error>>,
27    pub protocol_marker: PhantomData<Protocol>,
28}
29
30impl<Protocol, InnerStream, StreamTransformer> Stream
31    for ExchangeStream<Protocol, InnerStream, StreamTransformer>
32where
33    Protocol: StreamParser<StreamTransformer::Input>,
34    InnerStream: Stream<Item = Result<Protocol::Message, Protocol::Error>> + Unpin,
35    StreamTransformer: Transformer,
36    StreamTransformer::Error: From<SocketError>,
37{
38    type Item = Result<StreamTransformer::Output, StreamTransformer::Error>;
39
40    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
41        loop {
42            // Flush Self::Item buffer if it is not currently empty
43            if let Some(output) = self.buffer.pop_front() {
44                return Poll::Ready(Some(output));
45            }
46
47            // Poll inner `Stream` for next the next input protocol message
48            let input = match self.as_mut().project().stream.poll_next(cx) {
49                Poll::Ready(Some(input)) => input,
50                Poll::Ready(None) => return Poll::Ready(None),
51                Poll::Pending => return Poll::Pending,
52            };
53
54            // Parse input protocol message into `ExchangeMessage`
55            let exchange_message = match Protocol::parse(input) {
56                // `StreamParser` successfully deserialised `ExchangeMessage`
57                Some(Ok(exchange_message)) => exchange_message,
58
59                // If `StreamParser` returns an Err pass it downstream
60                Some(Err(err)) => return Poll::Ready(Some(Err(err.into()))),
61
62                // If `StreamParser` returns None it's a safe-to-skip message
63                None => continue,
64            };
65
66            // Transform `ExchangeMessage` into `Transformer::OutputIter`
67            // ie/ IntoIterator<Item = Result<Output, SocketError>>
68            self.transformer
69                .transform(exchange_message)
70                .into_iter()
71                .for_each(
72                    |output_result: Result<StreamTransformer::Output, StreamTransformer::Error>| {
73                        self.buffer.push_back(output_result)
74                    },
75                );
76        }
77    }
78}
79
80impl<Protocol, InnerStream, StreamTransformer>
81    ExchangeStream<Protocol, InnerStream, StreamTransformer>
82where
83    Protocol: StreamParser<StreamTransformer::Input>,
84    InnerStream: Stream,
85    StreamTransformer: Transformer,
86{
87    pub fn new(
88        stream: InnerStream,
89        transformer: StreamTransformer,
90        buffer: VecDeque<Result<StreamTransformer::Output, StreamTransformer::Error>>,
91    ) -> Self {
92        Self {
93            stream,
94            transformer,
95            buffer,
96            protocol_marker: PhantomData,
97        }
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use futures::StreamExt;
105
106    /// Parses `0` as a message to skip, a negative number as a parse error, and anything else as
107    /// itself.
108    struct SkipZeroParser;
109
110    impl StreamParser<i32> for SkipZeroParser {
111        type Stream = futures::stream::Empty<i32>;
112        type Message = i32;
113        type Error = SocketError;
114
115        fn parse(input: Result<i32, SocketError>) -> Option<Result<i32, SocketError>> {
116            match input {
117                Ok(0) => None,
118                Ok(n) if n < 0 => Some(Err(SocketError::Sink)),
119                other => Some(other),
120            }
121        }
122    }
123
124    /// Transforms each message into two outputs, so the buffer has to hold the second.
125    struct TwoOutputs;
126
127    impl Transformer for TwoOutputs {
128        type Error = SocketError;
129        type Input = i32;
130        type Output = i32;
131        type OutputIter = Vec<Result<i32, SocketError>>;
132
133        fn transform(&mut self, input: i32) -> Self::OutputIter {
134            vec![Ok(input), Ok(input * 10)]
135        }
136    }
137
138    #[tokio::test]
139    async fn exchange_stream_parses_skips_and_buffers_until_the_inner_stream_ends() {
140        let inner = futures::stream::iter([Ok(1), Ok(0), Ok(-1), Ok(2)]);
141        let stream =
142            ExchangeStream::<SkipZeroParser, _, _>::new(inner, TwoOutputs, VecDeque::new());
143
144        let outputs = stream
145            .map(|result| result.map_err(|error| error.to_string()))
146            .collect::<Vec<_>>()
147            .await;
148
149        assert_eq!(
150            outputs,
151            vec![
152                Ok(1),
153                Ok(10),
154                Err(SocketError::Sink.to_string()),
155                Ok(2),
156                Ok(20),
157            ]
158        );
159    }
160}
rustrade_integration/stream/exchange.rs

rustrade_integration/stream/
exchange.rs
