Correct the behavior of read_until when Pending is returned

Author: taiki-e

futures-util/src/io/read_until.rs

@@ -2,6 +2,7 @@ use futures_core::future::Future;
 use futures_core::task::{Context, Poll};
 use futures_io::AsyncBufRead;
 use std::io;
+use std::mem;
 use std::pin::Pin;
 
 /// Future for the [`read_until`](super::AsyncBufReadExt::read_until) method.
@@ -10,38 +11,48 @@ pub struct ReadUntil<'a, R: ?Sized + Unpin> {
     reader: &'a mut R,
     byte: u8,
     buf: &'a mut Vec<u8>,
+    read: usize,
 }
 
 impl<R: ?Sized + Unpin> Unpin for ReadUntil<'_, R> {}
 
 impl<'a, R: AsyncBufRead + ?Sized + Unpin> ReadUntil<'a, R> {
     pub(super) fn new(reader: &'a mut R, byte: u8, buf: &'a mut Vec<u8>) -> Self {
-        Self { reader, byte, buf }
+        Self { reader, byte, buf, read: 0 }
+    }
+}
+
+fn read_until_internal<R: AsyncBufRead + ?Sized + Unpin>(
+    mut reader: Pin<&mut R>,
+    byte: u8,
+    buf: &mut Vec<u8>,
+    read: &mut usize,
+    cx: &mut Context<'_>,
+) -> Poll<io::Result<usize>> {
+    loop {
+        let (done, used) = {
+            let available = try_ready!(reader.as_mut().poll_fill_buf(cx));
+            if let Some(i) = memchr::memchr(byte, available) {
+                buf.extend_from_slice(&available[..=i]);
+                (true, i + 1)
+            } else {
+                buf.extend_from_slice(available);
+                (false, available.len())
+            }
+        };
+        reader.as_mut().consume(used);
+        *read += used;
+        if done || used == 0 {
+            return Poll::Ready(Ok(mem::replace(read, 0)));
+        }
     }
 }
 
 impl<R: AsyncBufRead + ?Sized + Unpin> Future for ReadUntil<'_, R> {
     type Output = io::Result<usize>;
 
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let this = &mut *self;
-        let mut read = 0;
-        loop {
-            let (done, used) = {
-                let available = try_ready!(Pin::new(&mut this.reader).poll_fill_buf(cx));
-                if let Some(i) = memchr::memchr(this.byte, available) {
-                    this.buf.extend_from_slice(&available[..=i]);
-                    (true, i + 1)
-                } else {
-                    this.buf.extend_from_slice(available);
-                    (false, available.len())
-                }
-            };
-            Pin::new(&mut this.reader).consume(used);
-            read += used;
-            if done || used == 0 {
-                return Poll::Ready(Ok(read));
-            }
-        }
+        let Self { reader, byte, buf, read } = &mut *self;
+        read_until_internal(Pin::new(reader), *byte, buf, read, cx)
     }
 }

futures/tests/io_read_until.rs

@@ -1,6 +1,11 @@
 use futures::executor::block_on;
-use futures::io::AsyncBufReadExt;
-use std::io::Cursor;
+use futures::future::Future;
+use futures::io::{AsyncRead, AsyncBufRead, AsyncBufReadExt};
+use futures::task::{Context, Poll};
+use futures_test::task::noop_context;
+use std::cmp;
+use std::io::{self, Cursor};
+use std::pin::Pin;
 
 #[test]
 fn read_until() {
@@ -20,3 +25,73 @@ fn read_until() {
     assert_eq!(block_on(buf.read_until(b'3', &mut v)).unwrap(), 0);
     assert_eq!(v, []);
 }
+
+fn run<F: Future + Unpin>(mut f: F) -> F::Output {
+    let mut cx = noop_context();
+    loop {
+        if let Poll::Ready(x) = Pin::new(&mut f).poll(&mut cx) {
+            return x;
+        }
+    }
+}
+
+struct MaybePending<'a> {
+    inner: &'a [u8],
+    ready: bool,
+}
+
+impl<'a> MaybePending<'a> {
+    fn new(inner: &'a [u8]) -> Self {
+        Self { inner, ready: false }
+    }
+}
+
+impl AsyncRead for MaybePending<'_> {
+    fn poll_read(self: Pin<&mut Self>, _: &mut Context<'_>, _: &mut [u8])
+        -> Poll<io::Result<usize>>
+    {
+        unimplemented!()
+    }
+}
+
+impl AsyncBufRead for MaybePending<'_> {
+    fn poll_fill_buf<'a>(mut self: Pin<&'a mut Self>, _: &mut Context<'_>)
+        -> Poll<io::Result<&'a [u8]>>
+    {
+        if self.ready {
+            self.ready = false;
+            if self.inner.is_empty() { return Poll::Ready(Ok(&[])) }
+            let len = cmp::min(2, self.inner.len());
+            Poll::Ready(Ok(&self.inner[0..len]))
+        } else {
+            self.ready = true;
+            Poll::Pending
+        }
+    }
+
+    fn consume(mut self: Pin<&mut Self>, amt: usize) {
+        self.inner = &self.inner[amt..];
+    }
+}
+
+#[test]
+fn maybe_pending() {
+    let mut buf = MaybePending::new(b"12");
+    let mut v = Vec::new();
+    assert_eq!(run(buf.read_until(b'3', &mut v)).unwrap(), 2);
+    assert_eq!(v, b"12");
+
+    let mut buf = MaybePending::new(b"12333");
+    let mut v = Vec::new();
+    assert_eq!(run(buf.read_until(b'3', &mut v)).unwrap(), 3);
+    assert_eq!(v, b"123");
+    v.clear();
+    assert_eq!(run(buf.read_until(b'3', &mut v)).unwrap(), 1);
+    assert_eq!(v, b"3");
+    v.clear();
+    assert_eq!(run(buf.read_until(b'3', &mut v)).unwrap(), 1);
+    assert_eq!(v, b"3");
+    v.clear();
+    assert_eq!(run(buf.read_until(b'3', &mut v)).unwrap(), 0);
+    assert_eq!(v, []);
+}