xmtp_api_grpc/streams/
non_blocking_stream.rs

1//! Compatibility layer for JS-Fetch POST streams & gRPC Tonic Web
2//!
3//! a web ['fetch' request](https://developer.mozilla.org/en-US/docs/Web/API/Fetch_API/Using_Fetch)
4//! may complete successfully, but the fetch promise does not resolve until the first bytes of the
5//! body are received by the browser.[issue](https://github.com/devashishdxt/tonic-web-wasm-client/issues/22).
6//!
7//! This poses a behavior inconsistency between gRPC native - HTTP/2 and gRPC-web HTTP/1.1. On
8//! native, gRPC streams do not block on the first body received, while web streams do.
9//! This is particularly obvious in tests, where:
10//! 1. stream is started
11//! 2. data is sent (for instance, group messages)
12//! 3. inspect sent data
13//!
14//! on web, we never get past step 1.)
15//!
16//! This solution models web stream request as part of the stream.
17//! Once the initial promise request resolves, the stream continues polling the
18//! resulting response object.
19//!
20//! This problem is not unique to grpc-web, and must be solved for grpc-gateway streams as well
21//! [code example for  grpc-gateway](https://github.com/xmtp/libxmtp/blob/87338b819730ade4c292937e3243b16e3cdee248/xmtp_api_http/src/http_stream.rs#L165)
22//!
23//! In context of gRPC, this should not break anything that already works with native -- grpc requests, even
24//! unary requests, are all modeled as streams (a unary request is a stream with a single message),
25//! and none block on receipt of the body. Ideally, we could check the header status and ensure the
26//! initial response is 200 (OK), although the browser environment constraints do not allow for
27//! this.
28
29use futures::{Stream, TryFuture, TryStream, stream::FusedStream};
30use pin_project::pin_project;
31use std::{
32    future::Future,
33    pin::Pin,
34    task::{Context, Poll, ready},
35};
36use tonic::Status;
37
38use crate::streams::IntoInner;
39
40#[pin_project]
41/// The establish future for the http post stream
42struct StreamEstablish<F> {
43    #[pin]
44    inner: F,
45}
46
47impl<F> StreamEstablish<F> {
48    fn new(inner: F) -> Self {
49        Self { inner }
50    }
51}
52
53impl<F> Future for StreamEstablish<F>
54where
55    F: TryFuture<Error = Status>,
56{
57    type Output = Result<F::Ok, Status>;
58
59    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
60        use Poll::*;
61        let this = self.as_mut().project();
62        let response = ready!(this.inner.try_poll(cx));
63        let response = response.inspect_err(|e| {
64            tracing::error!("Error during grpc-web subscription establishment {e}");
65        })?;
66        Ready(Ok(response))
67    }
68}
69
70#[pin_project(project = ProjectStream)]
71/// The establish future for the http post stream
72enum StreamState<F, S> {
73    NotStarted {
74        #[pin]
75        future: StreamEstablish<F>,
76    },
77    Started {
78        #[pin]
79        stream: S,
80    },
81    Terminated,
82}
83
84#[pin_project]
85pub struct NonBlockingWebStream<F, S> {
86    #[pin]
87    state: StreamState<F, S>,
88}
89
90impl<F, S> NonBlockingWebStream<F, S>
91where
92    F: TryFuture<Error = Status>,
93{
94    pub fn new(request: F) -> Self {
95        Self {
96            state: StreamState::NotStarted {
97                future: StreamEstablish::new(request),
98            },
99        }
100    }
101
102    /// Internal API to construct a started variant
103    pub fn started(stream: S) -> Self {
104        Self {
105            state: StreamState::Started { stream },
106        }
107    }
108}
109
110impl<F, S> Stream for NonBlockingWebStream<F, S>
111where
112    S: TryStream<Error = Status>,
113    F: TryFuture<Error = Status>,
114    F::Ok: IntoInner<Out = S>,
115{
116    type Item = Result<S::Ok, Status>;
117
118    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
119        use ProjectStream::*;
120        let mut this = self.as_mut().project();
121        match this.state.as_mut().project() {
122            NotStarted { future } => {
123                match ready!(future.poll(cx)) {
124                    Ok(stream) => {
125                        this.state.set(StreamState::Started {
126                            stream: stream.into_inner(),
127                        });
128                    }
129                    Err(e) => {
130                        this.state.set(StreamState::Terminated);
131                        return Poll::Ready(Some(Err(e)));
132                    }
133                }
134                tracing::trace!("stream ready, polling for the first time...");
135                cx.waker().wake_by_ref();
136                Poll::Pending
137            }
138            Started { mut stream } => {
139                let next = stream.as_mut().try_poll_next(cx);
140                if let Poll::Ready(None) = next {
141                    this.state.set(StreamState::Terminated);
142                }
143                next
144            }
145            Terminated => Poll::Ready(None),
146        }
147    }
148}
149
150impl<F, S> FusedStream for NonBlockingWebStream<F, S>
151where
152    F: TryFuture<Error = Status>,
153    S: TryStream<Error = Status> + FusedStream,
154    F::Ok: IntoInner<Out = S>,
155{
156    fn is_terminated(&self) -> bool {
157        match &self.state {
158            StreamState::Started { stream } => stream.is_terminated(),
159            StreamState::Terminated => true,
160            _ => false,
161        }
162    }
163}
164
165impl<F, S> std::fmt::Debug for NonBlockingWebStream<F, S> {
166    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
167        match self.state {
168            StreamState::NotStarted { .. } => write!(f, "not started"),
169            StreamState::Started { .. } => write!(f, "started"),
170            StreamState::Terminated => write!(f, "terminated"),
171        }
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use futures::{Stream, stream};
178    use futures_test::future::FutureTestExt;
179    use prost::bytes::Bytes;
180    use tonic::{Response, Streaming};
181
182    use super::*;
183
184    struct TestStream;
185    impl Stream for TestStream {
186        type Item = Result<Response<Bytes>, Status>;
187
188        fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
189            unreachable!()
190        }
191    }
192
193    impl FusedStream for TestStream {
194        fn is_terminated(&self) -> bool {
195            unreachable!()
196        }
197    }
198
199    impl<T> From<Streaming<T>> for TestStream {
200        fn from(_: Streaming<T>) -> Self {
201            unreachable!()
202        }
203    }
204
205    struct MockFut;
206    impl Future for MockFut {
207        type Output = Result<TestStream, Status>;
208
209        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
210            unimplemented!()
211        }
212    }
213
214    impl IntoInner for MockFut {
215        type Out = TestStream;
216
217        fn into_inner(self) -> Self::Out {
218            todo!()
219        }
220    }
221
222    #[xmtp_common::test]
223    fn handles_err_on_establish() {
224        let stream: NonBlockingWebStream<_, TestStream> =
225            NonBlockingWebStream::new(futures::future::ready({
226                // we just need something that creates a reqwest error
227                // we also use now_or_never to guarantee this will trigger an error on the first poll
228                Err::<MockFut, _>(Status::internal("test error"))
229            }));
230        futures::pin_mut!(stream);
231
232        assert!(matches!(stream.state, StreamState::NotStarted { .. }));
233        let cx = futures::task::noop_waker();
234        let mut cx = std::task::Context::from_waker(&cx);
235        assert!(matches!(
236            stream.as_mut().poll_next(&mut cx),
237            Poll::Ready(Some(Err(_)))
238        ));
239
240        assert!(FusedStream::is_terminated(&stream));
241        assert!(matches!(
242            stream.as_mut().poll_next(&mut cx),
243            Poll::Ready(None)
244        ));
245    }
246
247    #[xmtp_common::test]
248    fn happy_path_future() {
249        let fut = futures::future::ready(Ok(()));
250        let fut = fut.pending_once();
251        let fut = StreamEstablish::new(fut);
252        futures::pin_mut!(fut);
253        let mut context = futures_test::task::noop_context();
254        assert_eq!(
255            Poll::Pending,
256            fut.as_mut().poll(&mut context).map(Result::unwrap)
257        );
258        assert_eq!(Poll::Ready(()), fut.poll(&mut context).map(Result::unwrap));
259    }
260
261    struct FakeFuture<T>(T);
262
263    impl<T> FakeFuture<T> {
264        fn inner(self: Pin<&mut Self>) -> Pin<&mut T> {
265            // This is okay because `field` is pinned when `self` is.
266            unsafe { self.map_unchecked_mut(|s| &mut s.0) }
267        }
268    }
269
270    impl<T> Future for FakeFuture<T>
271    where
272        T: TryFuture<Error = Status>,
273    {
274        type Output = Result<T::Ok, Status>;
275
276        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
277            self.inner().try_poll(cx)
278        }
279    }
280
281    struct FakeStream<T>(T);
282
283    impl<T> FakeStream<T> {
284        fn inner(self: Pin<&mut Self>) -> Pin<&mut T> {
285            // This is okay because `field` is pinned when `self` is.
286            unsafe { self.map_unchecked_mut(|s| &mut s.0) }
287        }
288    }
289
290    impl<T> Stream for FakeStream<T>
291    where
292        T: TryStream<Error = Status>,
293    {
294        type Item = Result<T::Ok, Status>;
295
296        fn poll_next(
297            self: Pin<&mut Self>,
298            cx: &mut Context<'_>,
299        ) -> Poll<Option<Result<T::Ok, Status>>> {
300            self.inner().try_poll_next(cx)
301        }
302    }
303
304    impl<T> IntoInner for FakeStream<T> {
305        type Out = FakeStream<T>;
306
307        fn into_inner(self) -> Self::Out {
308            self
309        }
310    }
311
312    impl<T: TryStream<Error = Status>> FusedStream for FakeStream<T> {
313        fn is_terminated(&self) -> bool {
314            unreachable!()
315        }
316    }
317
318    fn item<T>(i: T) -> Result<T, Status> {
319        Ok(i)
320    }
321
322    #[xmtp_common::test]
323    fn establish_changes_state_to_started() {
324        let s = FakeStream(stream::iter(vec![item(0usize), item(1), item(2)]));
325        let fut = futures::future::ready(Ok(s));
326        let fut = FakeFuture(fut);
327        let fut = fut.pending_once();
328        let s =
329            NonBlockingWebStream::<_, FakeStream<stream::Iter<std::vec::IntoIter<_>>>>::new(fut);
330
331        futures::pin_mut!(s);
332        let mut context = futures_test::task::noop_context();
333        assert_eq!(
334            Poll::Pending,
335            s.as_mut()
336                .poll_next(&mut context)
337                .map(Option::unwrap)
338                .map(Result::unwrap)
339        );
340        assert!(matches!(s.state, StreamState::NotStarted { .. }));
341        assert_eq!(
342            Poll::Pending,
343            s.as_mut()
344                .poll_next(&mut context)
345                .map(Option::unwrap)
346                .map(Result::unwrap)
347        );
348        assert!(matches!(s.state, StreamState::Started { .. }));
349        for i in 0..3 {
350            assert_eq!(
351                Poll::Ready(i),
352                s.as_mut()
353                    .poll_next(&mut context)
354                    .map(Option::unwrap)
355                    .map(Result::unwrap)
356            );
357        }
358        // stream ended after going through all items
359        assert_eq!(
360            Poll::Ready(None),
361            s.as_mut()
362                .poll_next(&mut context)
363                .map(|o| o.map(Result::unwrap))
364        );
365        assert_eq!(
366            Poll::Ready(None),
367            s.as_mut()
368                .poll_next(&mut context)
369                .map(|o| o.map(Result::unwrap))
370        );
371        // state should be terminated
372        assert!(matches!(s.state, StreamState::Terminated));
373    }
374}
xmtp_api_grpc/streams/non_blocking_stream.rs

xmtp_api_grpc/streams/
non_blocking_stream.rs
