refactor

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-ipc/src/messages/reader_async.rs

@@ -12,28 +12,12 @@ use futures::Stream;
 use pin_project_lite::pin_project;
 use vortex_array::session::ArrayRegistry;
 use vortex_error::VortexResult;
+use vortex_error::vortex_err;
 
 use crate::messages::DecoderMessage;
 use crate::messages::MessageDecoder;
 use crate::messages::PollRead;
 
-/// The state of an in-progress read operation.
-#[derive(Default)]
-enum ReadState {
-    /// Ready to consult the decoder for the next operation.
-    #[default]
-    AwaitingDecoder,
-    /// Filling the buffer with data from the underlying reader.
-    ///
-    /// Async readers may return fewer bytes than requested (partial reads), especially over network
-    /// connections. This state persists across multiple `poll_next` calls until the buffer is
-    /// completely filled, at which point we transition back to [`Self::AwaitingDecoder`].
-    Filling {
-        /// The number of bytes read into the buffer so far.
-        total_bytes_read: usize,
-    },
-}
-
 pin_project! {
     /// An IPC message reader backed by an `AsyncRead` stream.
     pub struct AsyncMessageReader<R> {
@@ -56,6 +40,65 @@ impl<R> AsyncMessageReader<R> {
     }
 }
 
+/// The state of an in-progress read operation.
+#[derive(Default)]
+enum ReadState {
+    /// Ready to consult the decoder for the next operation.
+    #[default]
+    AwaitingDecoder,
+    /// Filling the buffer with data from the underlying reader.
+    ///
+    /// Async readers may return fewer bytes than requested (partial reads), especially over network
+    /// connections. This state persists across multiple `poll_next` calls until the buffer is
+    /// completely filled, at which point we transition back to [`Self::AwaitingDecoder`].
+    Filling {
+        /// The number of bytes read into the buffer so far.
+        total_bytes_read: usize,
+    },
+}
+
+/// Result of polling the reader to fill the buffer.
+enum FillResult {
+    /// The buffer has been completely filled.
+    Filled,
+    /// Need more data (partial read occurred).
+    Pending,
+    /// Clean EOF at a message boundary.
+    Eof,
+}
+
+/// Polls the reader to fill the buffer, handling partial reads.
+fn poll_fill_buffer<R: AsyncRead>(
+    read: Pin<&mut R>,
+    buffer: &mut [u8],
+    total_bytes_read: &mut usize,
+    cx: &mut Context<'_>,
+) -> Poll<VortexResult<FillResult>> {
+    let unfilled = &mut buffer[*total_bytes_read..];
+
+    let bytes_read = ready!(read.poll_read(cx, unfilled))?; // Poll until ready.
+
+    // `0` bytes read indicates an EOF.
+    Poll::Ready(if bytes_read == 0 {
+        if *total_bytes_read > 0 {
+            Err(vortex_err!(
+                "unexpected EOF during partial read: read {total_bytes_read} of {} expected bytes",
+                buffer.len()
+            ))
+        } else {
+            Ok(FillResult::Eof)
+        }
+    } else {
+        *total_bytes_read += bytes_read;
+        if *total_bytes_read == buffer.len() {
+            Ok(FillResult::Filled)
+        } else {
+            debug_assert!(*total_bytes_read < buffer.len());
+            Ok(FillResult::Pending)
+        }
+    })
+}
+
 impl<R: AsyncRead> Stream for AsyncMessageReader<R> {
     type Item = VortexResult<DecoderMessage>;
 
@@ -65,8 +108,6 @@ impl<R: AsyncRead> Stream for AsyncMessageReader<R> {
             match this.state {
                 ReadState::AwaitingDecoder => match this.decoder.read_next(this.buffer)? {
                     PollRead::Some(msg) => return Poll::Ready(Some(Ok(msg))),
-                    // The decoder needs more data to be read into the buffer, so resize for more
-                    // space to read from `this.read`.
                     PollRead::NeedMore(new_len) => {
                         this.buffer.resize(new_len, 0x00);
                         *this.state = ReadState::Filling {
@@ -75,22 +116,17 @@ impl<R: AsyncRead> Stream for AsyncMessageReader<R> {
                     }
                 },
                 ReadState::Filling { total_bytes_read } => {
-                    let unfilled = &mut this.buffer.as_mut()[*total_bytes_read..];
-
-                    match ready!(this.read.as_mut().poll_read(cx, unfilled)) {
-                        Err(e) => return Poll::Ready(Some(Err(e.into()))),
-                        Ok(0) => return Poll::Ready(None),
-                        // Continue to read until we fill the buffer. We need to do this in a loop
-                        // to handle partial reads.
-                        Ok(bytes_read) => {
-                            *total_bytes_read += bytes_read;
-
-                            if *total_bytes_read == this.buffer.len() {
-                                *this.state = ReadState::AwaitingDecoder;
-                            } else {
-                                debug_assert!(*total_bytes_read < this.buffer.len());
-                            }
-                        }
+                    // Poll until ready.
+                    match ready!(poll_fill_buffer(
+                        this.read.as_mut(),
+                        this.buffer,
+                        total_bytes_read,
+                        cx
+                    )) {
+                        Err(e) => return Poll::Ready(Some(Err(e))),
+                        Ok(FillResult::Eof) => return Poll::Ready(None),
+                        Ok(FillResult::Filled) => *this.state = ReadState::AwaitingDecoder,
+                        Ok(FillResult::Pending) => {}
                     }
                 }
             }