From ab040bb498c96623f413efdc9fc16c46f72f9c63 Mon Sep 17 00:00:00 2001 From: Taiki Endo Date: Mon, 15 Jul 2019 15:23:42 +0900 Subject: [PATCH] tokio: add AsyncBufReadExt::read_line --- tokio/src/io/async_buf_read_ext.rs | 37 ++++++++++++++++ tokio/src/io/mod.rs | 1 + tokio/src/io/read_line.rs | 70 ++++++++++++++++++++++++++++++ tokio/tests/io_read_line.rs | 60 +++++++++++++++++++++++++ 4 files changed, 168 insertions(+) create mode 100644 tokio/src/io/read_line.rs create mode 100644 tokio/tests/io_read_line.rs diff --git a/tokio/src/io/async_buf_read_ext.rs b/tokio/src/io/async_buf_read_ext.rs index 8cbfec455..e81780742 100644 --- a/tokio/src/io/async_buf_read_ext.rs +++ b/tokio/src/io/async_buf_read_ext.rs @@ -1,3 +1,4 @@ +use crate::io::read_line::{read_line, ReadLine}; use crate::io::read_until::{read_until, ReadUntil}; use tokio_io::AsyncBufRead; @@ -29,6 +30,42 @@ pub trait AsyncBufReadExt: AsyncBufRead { { read_until(self, byte, buf) } + + /// Creates a future which will read all the bytes associated with this I/O + /// object into `buf` until a newline (the 0xA byte) or EOF is reached, + /// This method is the async equivalent to [`BufRead::read_line`](std::io::BufRead::read_line). + /// + /// This function will read bytes from the underlying stream until the + /// newline delimiter (the 0xA byte) or EOF is found. Once found, all bytes + /// up to, and including, the delimiter (if found) will be appended to + /// `buf`. + /// + /// The returned future will resolve to the number of bytes read once the read + /// operation is completed. + /// + /// In the case of an error the buffer and the object will be discarded, with + /// the error yielded. + /// + /// # Errors + /// + /// This function has the same error semantics as [`read_until`] and will + /// also return an error if the read bytes are not valid UTF-8. If an I/O + /// error is encountered then `buf` may contain some bytes already read in + /// the event that all data read so far was valid UTF-8. + /// + /// [`read_until`]: AsyncBufReadExt::read_until + /// + /// # Examples + /// + /// ``` + /// unimplemented!(); + /// ``` + fn read_line<'a>(&'a mut self, buf: &'a mut String) -> ReadLine<'a, Self> + where + Self: Unpin, + { + read_line(self, buf) + } } impl AsyncBufReadExt for R {} diff --git a/tokio/src/io/mod.rs b/tokio/src/io/mod.rs index b185701a1..bd1440ee2 100644 --- a/tokio/src/io/mod.rs +++ b/tokio/src/io/mod.rs @@ -42,6 +42,7 @@ mod async_write_ext; mod copy; mod read; mod read_exact; +mod read_line; mod read_to_end; mod read_until; mod write; diff --git a/tokio/src/io/read_line.rs b/tokio/src/io/read_line.rs new file mode 100644 index 000000000..987830d96 --- /dev/null +++ b/tokio/src/io/read_line.rs @@ -0,0 +1,70 @@ +use super::read_until::read_until_internal; +use std::future::Future; +use std::io; +use std::mem; +use std::pin::Pin; +use std::str; +use std::task::{Context, Poll}; +use tokio_io::AsyncBufRead; + +/// Future for the [`read_line`](crate::io::AsyncBufReadExt::read_line) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct ReadLine<'a, R: ?Sized + Unpin> { + reader: &'a mut R, + buf: &'a mut String, + bytes: Vec, + read: usize, +} + +impl Unpin for ReadLine<'_, R> {} + +pub(crate) fn read_line<'a, R>(reader: &'a mut R, buf: &'a mut String) -> ReadLine<'a, R> +where + R: AsyncBufRead + ?Sized + Unpin, +{ + ReadLine { + reader, + bytes: unsafe { mem::replace(buf.as_mut_vec(), Vec::new()) }, + buf, + read: 0, + } +} + +pub(super) fn read_line_internal( + reader: Pin<&mut R>, + cx: &mut Context<'_>, + buf: &mut String, + bytes: &mut Vec, + read: &mut usize, +) -> Poll> { + let ret = ready!(read_until_internal(reader, cx, b'\n', bytes, read)); + if str::from_utf8(&bytes).is_err() { + Poll::Ready(ret.and_then(|_| { + Err(io::Error::new( + io::ErrorKind::InvalidData, + "stream did not contain valid UTF-8", + )) + })) + } else { + debug_assert!(buf.is_empty()); + debug_assert_eq!(*read, 0); + // Safety: `bytes` is a valid UTF-8 because `str::from_utf8` returned `Ok`. + mem::swap(unsafe { buf.as_mut_vec() }, bytes); + Poll::Ready(ret) + } +} + +impl Future for ReadLine<'_, R> { + type Output = io::Result; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let Self { + reader, + buf, + bytes, + read, + } = &mut *self; + read_line_internal(Pin::new(reader), cx, buf, bytes, read) + } +} diff --git a/tokio/tests/io_read_line.rs b/tokio/tests/io_read_line.rs new file mode 100644 index 000000000..99b3b59e9 --- /dev/null +++ b/tokio/tests/io_read_line.rs @@ -0,0 +1,60 @@ +#![deny(warnings, rust_2018_idioms)] +#![feature(async_await)] + +use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncRead}; +use tokio_test::assert_ok; + +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; + +#[tokio::test] +async fn read_line() { + struct Rd { + val: &'static [u8], + } + + impl AsyncRead for Rd { + fn poll_read( + self: Pin<&mut Self>, + _: &mut Context<'_>, + _: &mut [u8], + ) -> Poll> { + unimplemented!() + } + } + + impl AsyncBufRead for Rd { + fn poll_fill_buf<'a>( + self: Pin<&'a mut Self>, + _: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(self.val)) + } + + fn consume(mut self: Pin<&mut Self>, amt: usize) { + self.val = &self.val[amt..]; + } + } + + let mut buf = String::new(); + let mut rd = Rd { + val: b"hello\nworld\n\n", + }; + + let n = assert_ok!(rd.read_line(&mut buf).await); + assert_eq!(n, 6); + assert_eq!(buf, "hello\n"); + buf.clear(); + let n = assert_ok!(rd.read_line(&mut buf).await); + assert_eq!(n, 6); + assert_eq!(buf, "world\n"); + buf.clear(); + let n = assert_ok!(rd.read_line(&mut buf).await); + assert_eq!(n, 1); + assert_eq!(buf, "\n"); + buf.clear(); + let n = assert_ok!(rd.read_line(&mut buf).await); + assert_eq!(n, 0); + assert_eq!(buf, ""); +}