Read multiple chunks into a slice of Bytes

Author: stammw

quinn-proto/src/connection/mod.rs

@@ -868,7 +868,7 @@ where
     ///
     /// While stream data is typically processed by applications in-order, unordered reads improve
     /// performance when packet loss occurs and data cannot be retransmitted before the flow control
-    /// window is filled. When in-order delivery is required, the sibling `read()` or `read_chunk()`
+    /// window is filled. When in-order delivery is required, the sibling `read()` or `read_chunk[s]()`
     /// methods should be used.
     ///
     /// The return value if `Ok` contains the bytes and their offset in the stream.
@@ -887,6 +887,18 @@ where
         }))
     }
 
+    /// Read the next ordered chunks from the given recv stream
+    pub fn read_chunks(
+        &mut self,
+        id: StreamId,
+        bufs: &mut [Bytes],
+    ) -> Result<Option<usize>, ReadError> {
+        Ok(self.streams.read_chunks(id, bufs)?.map(|result| {
+            self.add_read_credits(id, result.max_stream_data, result.max_data);
+            result.n_bufs
+        }))
+    }
+
     /// Read from the given recv stream
     pub fn read(&mut self, id: StreamId, buf: &mut [u8]) -> Result<Option<usize>, ReadError> {
         Ok(self.streams.read(id, buf)?.map(|result| {

quinn-proto/src/connection/streams.rs

@@ -213,6 +213,23 @@ impl Streams {
         }))
     }
 
+    pub fn read_chunks(
+        &mut self,
+        id: StreamId,
+        bufs: &mut [Bytes],
+    ) -> Result<Option<ReadChunksResult>, ReadError> {
+        let res = self.try_read(id, |rs: &mut Recv| rs.read_chunks(bufs))?;
+        Ok(res.map(|(n_bufs, max_stream_data)| {
+            let len = bufs[..n_bufs].iter().fold(0, |t, b| t + b.len());
+            let max_data = self.add_read_credits(len as u64);
+            ReadChunksResult {
+                n_bufs,
+                max_stream_data,
+                max_data,
+            }
+        }))
+    }
+
     fn try_read<T, U>(
         &mut self,
         id: StreamId,
@@ -1019,6 +1036,15 @@ pub struct ReadChunkResult {
     pub max_data: ShouldTransmit,
 }
 
+/// Result of a `Streams::read_chunks` call in case the stream had not ended yet
+#[derive(Debug, Eq, PartialEq)]
+#[must_use = "A frame might need to be enqueued"]
+pub struct ReadChunksResult {
+    pub n_bufs: usize,
+    pub max_stream_data: ShouldTransmit,
+    pub max_data: ShouldTransmit,
+}
+
 /// Result of a `Streams::read` call in case the stream had not ended yet
 #[derive(Debug, Eq, PartialEq, Copy, Clone)]
 #[must_use = "A frame might need to be enqueued"]
@@ -1147,6 +1173,28 @@ impl Recv {
         }
     }
 
+    fn read_chunks(&mut self, chunks: &mut [Bytes]) -> Result<Option<usize>, ReadError> {
+        if chunks.is_empty() {
+            return Ok(Some(0));
+        }
+
+        let mut size = 0;
+        while let Some(bytes) = self.assembler.read_chunk()? {
+            chunks[size] = bytes;
+            size += 1;
+
+            if size >= chunks.len() {
+                return Ok(Some(size));
+            }
+        }
+
+        if size > 0 {
+            return Ok(Some(size));
+        }
+
+        self.read_blocked().map(|()| None)
+    }
+
     fn read_blocked(&mut self) -> Result<(), ReadError> {
         match self.state {
             RecvState::ResetRecvd { error_code, .. } => {

quinn-proto/src/tests/mod.rs

@@ -1632,3 +1632,80 @@ fn repeated_request_response() {
         assert_matches!(pair.client_conn_mut(client_ch).read_unordered(s), Ok(None));
     }
 }
+
+#[test]
+fn read_chunks() {
+    let _guard = subscribe();
+    let server = ServerConfig {
+        transport: Arc::new(TransportConfig {
+            stream_window_bidi: 1u32.into(),
+            ..TransportConfig::default()
+        }),
+        ..server_config()
+    };
+    let mut pair = Pair::new(Default::default(), server);
+    let (client_ch, server_ch) = pair.connect();
+    let mut empty = vec![];
+    let mut chunks = vec![Bytes::new(), Bytes::new()];
+    const ONE: &[u8] = b"ONE";
+    const TWO: &[u8] = b"TWO";
+    const THREE: &[u8] = b"THREE";
+    for _ in 0..3 {
+        let s = pair.client_conn_mut(client_ch).open(Dir::Bi).unwrap();
+
+        pair.client_conn_mut(client_ch).write(s, ONE).unwrap();
+        pair.drive();
+        pair.client_conn_mut(client_ch).write(s, TWO).unwrap();
+        pair.drive();
+        pair.client_conn_mut(client_ch).write(s, THREE).unwrap();
+
+        pair.drive();
+
+        assert_eq!(pair.server_conn_mut(server_ch).accept(Dir::Bi), Some(s));
+
+        // Read into an empty slice can't do much you, but doesn't crash
+        assert_eq!(
+            pair.server_conn_mut(server_ch).read_chunks(s, &mut empty),
+            Ok(Some(0))
+        );
+
+        // Read until `chunks` is filled
+        assert_eq!(
+            pair.server_conn_mut(server_ch).read_chunks(s, &mut chunks),
+            Ok(Some(2))
+        );
+        assert_eq!(&chunks, &[ONE, TWO]);
+
+        // Read the rest
+        assert_eq!(
+            pair.server_conn_mut(server_ch).read_chunks(s, &mut chunks),
+            Ok(Some(1))
+        );
+        assert_eq!(&chunks[..1], &[THREE]);
+
+        // We've read everything, stream is now blocked
+        assert_eq!(
+            pair.server_conn_mut(server_ch).read_chunks(s, &mut chunks),
+            Err(ReadError::Blocked)
+        );
+
+        // Read a new chunk after we've been blocked
+        pair.client_conn_mut(client_ch).write(s, ONE).unwrap();
+        pair.drive();
+        assert_eq!(
+            pair.server_conn_mut(server_ch).read_chunks(s, &mut chunks),
+            Ok(Some(1))
+        );
+        assert_eq!(&chunks[..1], &[ONE]);
+
+        // Stream finishes by yeilding `Ok(None)`
+        pair.client_conn_mut(client_ch).finish(s).unwrap();
+        pair.drive();
+        assert_matches!(
+            pair.server_conn_mut(server_ch).read_chunks(s, &mut chunks),
+            Ok(None)
+        );
+
+        pair.drive();
+    }
+}

quinn/src/streams.rs

@@ -389,6 +389,27 @@ where
         self.poll_read_generic(cx, |conn, stream| conn.inner.read_chunk(stream))
     }
 
+    /// Read the next segments of data
+    ///
+    /// Fills `bufs` with the segments of data beginning immediately after the
+    /// last data yielded by `read` or `read_chunk`, or `None` if the stream was
+    /// finished.
+    ///
+    /// Slightly more efficient than `read` due to not copying. Chunk boundaries
+    /// do not correspond to peer writes, and hence cannot be used as framing.
+    pub fn read_chunks<'a>(&'a mut self, bufs: &'a mut [Bytes]) -> ReadChunks<'a, S> {
+        ReadChunks { stream: self, bufs }
+    }
+
+    /// Foundation of [`read_chunks()`]: RecvStream::read_chunks
+    fn poll_read_chunks(
+        &mut self,
+        cx: &mut Context,
+        bufs: &mut [Bytes],
+    ) -> Poll<Result<Option<usize>, ReadError>> {
+        self.poll_read_generic(cx, |conn, stream| conn.inner.read_chunks(stream, bufs))
+    }
+
     /// Convenience method to read all remaining data into a buffer
     ///
     /// The returned future fails with [`ReadToEndError::TooLong`] if it's longer than `size_limit`
@@ -795,6 +816,28 @@ where
     }
 }
 
+/// Future produced by [`RecvStream::read_chunks()`].
+///
+/// [`RecvStream::read_chunks()`]: crate::generic::RecvStream::read_chunks
+pub struct ReadChunks<'a, S>
+where
+    S: proto::crypto::Session,
+{
+    stream: &'a mut RecvStream<S>,
+    bufs: &'a mut [Bytes],
+}
+
+impl<'a, S> Future for ReadChunks<'a, S>
+where
+    S: proto::crypto::Session,
+{
+    type Output = Result<Option<usize>, ReadError>;
+    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
+        let this = self.get_mut();
+        this.stream.poll_read_chunks(cx, this.bufs)
+    }
+}
+
 /// Future produced by [`SendStream::write()`].
 ///
 /// [`SendStream::write()`]: crate::generic::SendStream::write