Only return Poll::Pending if the inner future returns it

The waker stored in the future will only wake up and re-poll the stream
if the call to future returned Poll::Pending. Because of this, if we
didn't receive enough data to reconstruct the ResourceDiff, we need to
loop and poll the inner future again rather than return Poll::Pending
ourselves.

crates/rdsys-backend/src/lib.rs

@@ -5,12 +5,11 @@
 
 use bytes::{self, Buf, Bytes};
 use core::pin::Pin;
-use core::task::{Poll, ready};
 use futures_util::{Stream, StreamExt};
 use reqwest::Client;
 use std::io::{self, BufRead};
+use std::task::{ready, Context, Poll};
 use tokio::sync::mpsc;
-use tokio_stream::wrappers::ReceiverStream;
 use tokio_util::sync::ReusableBoxFuture;
 
 pub mod proto;
@@ -43,7 +42,7 @@ impl From<io::Error> for Error {
 /// An iterable wrapper of ResourceDiff items for the streamed chunks of Bytes
 /// received from the connection to the rdsys backend
 pub struct ResourceStream {
-    inner: ReusableBoxFuture<'static, (Option<Bytes>, ReceiverStream<Bytes>)>,
+    inner: ReusableBoxFuture<'static, (Option<Bytes>, mpsc::Receiver<Bytes>)>,
     buf: Vec<u8>,
     partial: Option<bytes::buf::Reader<Bytes>>,
 }
@@ -51,23 +50,24 @@ pub struct ResourceStream {
 impl ResourceStream {
     pub fn new(rx: mpsc::Receiver<Bytes>) -> ResourceStream {
         ResourceStream {
-            inner: ReusableBoxFuture::new(make_future(ReceiverStream::new(rx))),
+            inner: ReusableBoxFuture::new(make_future(rx)),
             buf: vec![],
             partial: None,
         }
     }
 }
 
-async fn make_future(mut rx: ReceiverStream<Bytes>) -> (Option<Bytes>, ReceiverStream<Bytes>) {
-    let result = rx.next().await;
+async fn make_future(mut rx: mpsc::Receiver<Bytes>) -> (Option<Bytes>, mpsc::Receiver<Bytes>) {
+    let result = rx.recv().await;
     (result, rx)
 }
 
 impl Stream for ResourceStream {
     type Item = proto::ResourceDiff;
-    fn poll_next(mut self: Pin<&mut Self>, cx: &mut core::task::Context) -> Poll<Option<Self::Item>> {
-        let parse =
-            |buffer: &mut bytes::buf::Reader<Bytes>, buf: &mut Vec<u8>| -> Result<Option<Self::Item>, Error> {
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let parse = |buffer: &mut bytes::buf::Reader<Bytes>,
+                     buf: &mut Vec<u8>|
+         -> Result<Option<Self::Item>, Error> {
             match buffer.read_until(b'\r', buf) {
                 Ok(_) => match buf.pop() {
                     Some(b'\r') => match serde_json::from_slice(buf) {
@@ -98,9 +98,8 @@ impl Stream for ResourceStream {
             }
         }
         self.buf = buf;
-        println!("HERE");
+        loop {
             let (result, rx) = ready!(self.inner.poll(cx));
-        println!("HERE2");
             self.inner.set(make_future(rx));
             match result {
                 Some(chunk) => {
@@ -110,11 +109,12 @@ impl Stream for ResourceStream {
                             self.partial = Some(buffer);
                             return Poll::Ready(Some(diff));
                         }
-                    Ok(None) => Poll::Pending, //maybe loop here?
+                        Ok(None) => continue,
                         Err(_) => return Poll::Ready(None),
                     }
                 }
-            None => Poll::Ready(None),
+                None => return Poll::Ready(None),
+            }
         }
     }
 }