Change ResourceStream from an Iterator into a Stream

crates/rdsys-backend/Cargo.toml

@@ -10,6 +10,7 @@ serde_json = "1"
 futures-util = { version = "0.3"}
 serde = { version = "1", features = ["derive"]}
 bytes = "1"
-tokio = "1"
+tokio = { version = "1", features = ["macros"]}
-
 reqwest = { version = "0.11", features = ["stream"]}
+tokio-stream = "0.1.12"
+futures = "0.3.27"

crates/rdsys-backend/src/lib.rs

@@ -4,10 +4,13 @@
 //! https://gitlab.torproject.org/tpo/anti-censorship/rdsys/-/blob/main/doc/backend-api.md
 
 use bytes::{self, Buf, Bytes};
-use futures_util::StreamExt;
+use core::pin::Pin;
+use core::task::Poll;
+use futures_util::{Stream, StreamExt};
 use reqwest::Client;
 use std::io::{self, BufRead};
-use std::sync::mpsc;
+use tokio::sync::mpsc;
+use tokio_stream::wrappers::ReceiverStream;
 
 pub mod proto;
 
@@ -39,27 +42,37 @@ 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 {
-    rx: mpsc::Receiver<Bytes>,
+    rx: ReceiverStream<Bytes>,
     buf: Vec<u8>,
     partial: Option<bytes::buf::Reader<Bytes>>,
 }
 
-impl Iterator for ResourceStream {
+impl ResourceStream {
+    pub fn new(rx: mpsc::Receiver<Bytes>) -> ResourceStream {
+        ResourceStream {
+            rx: ReceiverStream::new(rx),
+            buf: vec![],
+            partial: None,
+        }
+    }
+}
+
+impl Stream for ResourceStream {
     type Item = proto::ResourceDiff;
-    fn next(&mut self) -> Option<Self::Item> {
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut core::task::Context) -> Poll<Option<Self::Item>> {
-        let mut parse =
+        let parse =
-            |buffer: &mut bytes::buf::Reader<Bytes>| -> Result<Option<Self::Item>, Error> {
+            |buffer: &mut bytes::buf::Reader<Bytes>, buf: &mut Vec<u8>| -> Result<Option<Self::Item>, Error> {
-                match buffer.read_until(b'\r', &mut self.buf) {
+                match buffer.read_until(b'\r', buf) {
-                    Ok(_) => match self.buf.pop() {
+                    Ok(_) => match buf.pop() {
-                        Some(b'\r') => match serde_json::from_slice(&self.buf) {
+                        Some(b'\r') => match serde_json::from_slice(buf) {
                             Ok(diff) => {
-                                self.buf.clear();
+                                buf.clear();
                                 Ok(Some(diff))
                             }
                             Err(e) => Err(Error::JSON(e)),
                         },
                         Some(n) => {
-                            self.buf.push(n);
+                            buf.push(n);
                             Ok(None)
                         }
                         None => Ok(None),
@@ -67,26 +80,33 @@ impl Iterator for ResourceStream {
                     Err(e) => Err(Error::Io(e)),
                 }
             };
-
+        // This clone is here to avoid having multiple mutable references to self
+        // it's not optimal performance-wise but given that these resource streams aren't large
+        // this feels like an acceptable trade-off to the complexity of interior mutability
+        let mut buf = self.buf.clone();
         if let Some(p) = &mut self.partial {
-            match parse(p) {
+            match parse(p, &mut buf) {
-                Ok(Some(diff)) => return Some(diff),
+                Ok(Some(diff)) => return Poll::Ready(Some(diff)),
                 Ok(None) => self.partial = None,
-                Err(_) => return None,
+                Err(_) => return Poll::Ready(None),
             }
         }
-        for chunk in &self.rx {
+        self.buf = buf;
-            let mut buffer = chunk.reader();
+        match Pin::new(&mut self.rx).poll_next(cx) {
-            match parse(&mut buffer) {
+            Poll::Ready(Some(chunk)) => {
-                Ok(Some(diff)) => {
+                let mut buffer = chunk.reader();
-                    self.partial = Some(buffer);
+                match parse(&mut buffer, &mut self.buf) {
-                    return Some(diff);
+                    Ok(Some(diff)) => {
+                        self.partial = Some(buffer);
+                        return Poll::Ready(Some(diff));
+                    }
+                    Ok(None) => Poll::Pending, //maybe loop here?
+                    Err(_) => return Poll::Ready(None),
                 }
-                Ok(None) => continue,
+            }
-                Err(_) => return None,
+            Poll::Ready(None) => Poll::Ready(None),
-            };
+            Poll::Pending => Poll::Pending,
         }
-        None
     }
 }
 
@@ -94,72 +114,76 @@ impl Iterator for ResourceStream {
 mod tests {
     use super::*;
 
-    #[test]
+    #[tokio::test]
-    fn parse_resource() {
+    async fn parse_resource() {
+        let mut cx = std::task::Context::from_waker(futures::task::noop_waker_ref());
         let chunk = Bytes::from_static(
             b"{\"new\": null,\"changed\": null,\"gone\": null,\"full_update\": true}\r",
         );
-        let (tx, rx) = mpsc::channel();
+        let (tx, rx) = mpsc::channel(100);
-        tx.send(chunk).unwrap();
+        tx.send(chunk).await.unwrap();
-        let mut diffs = ResourceStream {
+        let mut diffs = ResourceStream::new(rx);
-            rx: rx,
+        let res = Pin::new(&mut diffs).poll_next(&mut cx);
-            partial: None,
+        assert_ne!(res, Poll::Ready(None));
-            buf: vec![],
+        assert_ne!(res, Poll::Pending);
-        };
+        if let Poll::Ready(Some(diff)) = res {
-        let res = diffs.next();
-        assert_ne!(res, None);
-        if let Some(diff) = res {
             assert_eq!(diff.new, None);
             assert_eq!(diff.full_update, true);
         }
     }
 
-    #[test]
+    #[tokio::test]
-    fn parse_across_chunks() {
+    async fn parse_across_chunks() {
+        let mut cx = std::task::Context::from_waker(futures::task::noop_waker_ref());
         let chunk1 = Bytes::from_static(b"{\"new\": null,\"changed\": null,");
         let chunk2 = Bytes::from_static(b"\"gone\": null,\"full_update\": true}\r");
-        let (tx, rx) = mpsc::channel();
+        let (tx, rx) = mpsc::channel(100);
-        tx.send(chunk1).unwrap();
+        tx.send(chunk1).await.unwrap();
-        tx.send(chunk2).unwrap();
+        tx.send(chunk2).await.unwrap();
-        let mut diffs = ResourceStream {
+        let mut diffs = ResourceStream::new(rx);
-            rx: rx,
+        let mut res = Pin::new(&mut diffs).poll_next(&mut cx);
-            partial: None,
+        while let Poll::Pending = res {
-            buf: vec![],
+            res = Pin::new(&mut diffs).poll_next(&mut cx);
-        };
+        }
-        let res = diffs.next();
+        assert_ne!(res, Poll::Ready(None));
-        assert_ne!(res, None);
+        assert_ne!(res, Poll::Pending);
-        if let Some(diff) = res {
+        if let Poll::Ready(Some(diff)) = res {
             assert_eq!(diff.new, None);
             assert_eq!(diff.full_update, true);
         }
     }
 
-    #[test]
+    #[tokio::test]
-    fn parse_multi_diff_partial_chunks() {
+    async fn parse_multi_diff_partial_chunks() {
+        let mut cx = std::task::Context::from_waker(futures::task::noop_waker_ref());
         let chunk1 = Bytes::from_static(b"{\"new\": null,\"changed\": null,");
         let chunk2 =
             Bytes::from_static(b"\"gone\": null,\"full_update\": true}\r{\"new\": null,\"changed");
         let chunk3 = Bytes::from_static(b"\": null,\"gone\": null,\"full_update\": true}");
         let chunk4 = Bytes::from_static(b"\r");
-        let (tx, rx) = mpsc::channel();
+        let (tx, rx) = mpsc::channel(100);
-        tx.send(chunk1).unwrap();
+        tx.send(chunk1).await.unwrap();
-        tx.send(chunk2).unwrap();
+        tx.send(chunk2).await.unwrap();
-        tx.send(chunk3).unwrap();
+        tx.send(chunk3).await.unwrap();
-        tx.send(chunk4).unwrap();
+        tx.send(chunk4).await.unwrap();
-        let mut diffs = ResourceStream {
+        let mut diffs = ResourceStream::new(rx);
-            rx: rx,
+        let mut res = Pin::new(&mut diffs).poll_next(&mut cx);
-            partial: None,
+        while let Poll::Pending = res {
-            buf: vec![],
+            res = Pin::new(&mut diffs).poll_next(&mut cx);
-        };
+        }
-        let mut res = diffs.next();
+        assert_ne!(res, Poll::Ready(None));
-        assert_ne!(res, None);
+        assert_ne!(res, Poll::Pending);
-        if let Some(diff) = res {
+        if let Poll::Ready(Some(diff)) = res {
             assert_eq!(diff.new, None);
             assert_eq!(diff.full_update, true);
         }
-        res = diffs.next();
+        res = Pin::new(&mut diffs).poll_next(&mut cx);
-        assert_ne!(res, None);
+        while let Poll::Pending = res {
-        if let Some(diff) = res {
+            res = Pin::new(&mut diffs).poll_next(&mut cx);
+        }
+        assert_ne!(res, Poll::Ready(None));
+        assert_ne!(res, Poll::Pending);
+        if let Poll::Ready(Some(diff)) = res {
             assert_eq!(diff.new, None);
             assert_eq!(diff.full_update, true);
         }
@@ -178,9 +202,13 @@ mod tests {
 /// let name = String::from("https");
 /// let token = String::from("HttpsApiTokenPlaceholder");
 /// let types = vec![String::from("obfs2"), String::from("scramblesuit")];
-/// let rx = start_stream(endpoint, name, token, types).await.unwrap();
+/// let stream = start_stream(endpoint, name, token, types).await.unwrap();
-/// for diff in rx {
+/// loop {
-///     println!("Received diff: {:?}", diff);
+///     match Pin::new(&mut stream).poll_next(&mut cx) {
+///         Poll::Ready(Some(diff)) => println!("Received diff: {:?}", diff),
+///         Poll::Ready(None) => break,
+///         Poll::Pending => continue,
+///     }
 /// }
 /// ```
 ///
@@ -190,7 +218,7 @@ pub async fn start_stream(
     token: String,
     resource_types: Vec<String>,
 ) -> Result<ResourceStream, Error> {
-    let (tx, rx) = mpsc::channel();
+    let (tx, rx) = mpsc::channel(100);
 
     let req = proto::ResourceRequest {
         request_origin: name,
@@ -218,12 +246,8 @@ pub async fn start_stream(
                     return;
                 }
             };
-            tx.send(bytes).unwrap();
+            tx.send(bytes).await.unwrap();
         }
     });
-    Ok(ResourceStream {
+    Ok(ResourceStream::new(rx))
-        rx,
-        buf: vec![],
-        partial: None,
-    })
 }