src/backup/chunk_stream.rs: fix behaviour for large input buffers

src/backup/chunk_stream.rs

@@ -9,13 +9,13 @@ pub struct ChunkStream<S: Stream<Item=Vec<u8>, Error=Error>> {
     input: S,
     chunker: Chunker,
     buffer: Option<Vec<u8>>,
-    rest: Option<Vec<u8>>,
+    scan: Option<Vec<u8>>,
 }
 
 impl <S: Stream<Item=Vec<u8>, Error=Error>> ChunkStream<S> {
 
     pub fn new(input: S) -> Self {
-        Self { input, chunker: Chunker::new(4 * 1024 * 1024), buffer: None, rest: None }
+        Self { input, chunker: Chunker::new(4 * 1024 * 1024), buffer: None, scan: None}
     }
 }
 
@@ -26,32 +26,8 @@ impl <S: Stream<Item=Vec<u8>, Error=Error>> Stream for ChunkStream<S> {
 
     fn poll(&mut self) -> Poll<Option<Vec<u8>>, Error> {
         loop {
-            match self.input.poll() {
-                Err(err) => {
-                    return Err(err);
-                }
-                Ok(Async::NotReady) => {
-                    return Ok(Async::NotReady);
-                }
-                Ok(Async::Ready(None)) => {
-                    let mut data = self.buffer.take().or_else(|| Some(vec![])).unwrap();
-                    if let Some(rest) = self.rest.take() { data.extend(rest); }
-
-                    if data.len() > 0 {
-                        return Ok(Async::Ready(Some(data)));
-                    } else {
-                        return Ok(Async::Ready(None));
-                    }
-                }
-                Ok(Async::Ready(Some(data))) => {
-
-                    let data = if let Some(mut rest) = self.rest.take() {
-                        rest.extend(data);
-                        rest
-                    } else {
-                        data
-                    };
 
+            if let Some(data) = self.scan.take() {
                 let buffer = self.buffer.get_or_insert_with(|| Vec::with_capacity(1024*1024));
                 let boundary = self.chunker.scan(&data);
 
@@ -64,15 +40,34 @@ impl <S: Stream<Item=Vec<u8>, Error=Error>> Stream for ChunkStream<S> {
                 } else if boundary < data.len() {
                     let (left, right) = data.split_at(boundary);
                     buffer.extend(left);
-
+                    self.scan = Some(right.to_vec());
-                        let rest = self.rest.get_or_insert_with(|| Vec::with_capacity(right.len()));
-                        rest.extend(right);
-
                     return Ok(Async::Ready(self.buffer.take()));
                 } else {
                     panic!("got unexpected chunk boundary from chunker");
                 }
             }
+
+            match self.input.poll() {
+                Err(err) => {
+                    return Err(err);
+                }
+                Ok(Async::NotReady) => {
+                    return Ok(Async::NotReady);
+                }
+                Ok(Async::Ready(None)) => {
+                    let mut data = self.buffer.take().or_else(|| Some(vec![])).unwrap();
+                    if let Some(rest) = self.scan.take() { data.extend(rest); }
+
+                    if data.len() > 0 {
+                        return Ok(Async::Ready(Some(data)));
+                    } else {
+                        return Ok(Async::Ready(None));
+                    }
+                }
+                Ok(Async::Ready(Some(data))) => {
+                    let scan = self.scan.get_or_insert_with(|| Vec::with_capacity(1024*1024));
+                    scan.extend(data);
+                }
             }
         }
     }