From 0cfa120ba8c93822d6a84f2e7392008d84622010 Mon Sep 17 00:00:00 2001 From: Taiki Endo Date: Mon, 15 Jul 2019 15:21:09 +0900 Subject: [PATCH] tokio: add AsyncBufReadExt::read_until --- tokio/Cargo.toml | 3 +- tokio/src/io/async_buf_read_ext.rs | 30 +++++++++++- tokio/src/io/mod.rs | 1 + tokio/src/io/read_until.rs | 74 ++++++++++++++++++++++++++++++ tokio/tests/io_read_until.rs | 56 ++++++++++++++++++++++ 5 files changed, 162 insertions(+), 2 deletions(-) create mode 100644 tokio/src/io/read_until.rs create mode 100644 tokio/tests/io_read_until.rs diff --git a/tokio/Cargo.toml b/tokio/Cargo.toml index aff47d04c..d45b8d004 100644 --- a/tokio/Cargo.toml +++ b/tokio/Cargo.toml @@ -40,7 +40,7 @@ default = [ codec = ["io", "tokio-codec"] fs = ["tokio-fs"] -io = ["bytes", "tokio-io"] +io = ["bytes", "tokio-io", "memchr"] reactor = ["io", "tokio-reactor"] rt-full = [ "num_cpus", @@ -80,6 +80,7 @@ tokio-tcp = { version = "0.2.0", optional = true, path = "../tokio-tcp" } tokio-udp = { version = "0.2.0", optional = true, path = "../tokio-udp" } tokio-timer = { version = "0.3.0", optional = true, path = "../tokio-timer" } tracing-core = { version = "0.1", optional = true } +memchr = { version = "2.2", optional = true } # Needed for async/await preview support #tokio-futures = { version = "0.2.0", optional = true, path = "../tokio-futures" } diff --git a/tokio/src/io/async_buf_read_ext.rs b/tokio/src/io/async_buf_read_ext.rs index 7f0db43f8..8cbfec455 100644 --- a/tokio/src/io/async_buf_read_ext.rs +++ b/tokio/src/io/async_buf_read_ext.rs @@ -1,6 +1,34 @@ +use crate::io::read_until::{read_until, ReadUntil}; + use tokio_io::AsyncBufRead; /// An extension trait which adds utility methods to `AsyncBufRead` types. -pub trait AsyncBufReadExt: AsyncBufRead {} +pub trait AsyncBufReadExt: AsyncBufRead { + /// Creates a future which will read all the bytes associated with this I/O + /// object into `buf` until the delimiter `byte` or EOF is reached. + /// This method is the async equivalent to [`BufRead::read_until`](std::io::BufRead::read_until). + /// + /// This function will read bytes from the underlying stream until the + /// delimiter 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. + /// + /// # Examples + /// + /// ``` + /// unimplemented!(); + /// ``` + fn read_until<'a>(&'a mut self, byte: u8, buf: &'a mut Vec) -> ReadUntil<'a, Self> + where + Self: Unpin, + { + read_until(self, byte, buf) + } +} impl AsyncBufReadExt for R {} diff --git a/tokio/src/io/mod.rs b/tokio/src/io/mod.rs index 10d7b4b28..b185701a1 100644 --- a/tokio/src/io/mod.rs +++ b/tokio/src/io/mod.rs @@ -43,6 +43,7 @@ mod copy; mod read; mod read_exact; mod read_to_end; +mod read_until; mod write; mod write_all; diff --git a/tokio/src/io/read_until.rs b/tokio/src/io/read_until.rs new file mode 100644 index 000000000..2092176ba --- /dev/null +++ b/tokio/src/io/read_until.rs @@ -0,0 +1,74 @@ +use std::future::Future; +use std::io; +use std::mem; +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio_io::AsyncBufRead; + +/// Future for the [`read_until`](crate::io::AsyncBufReadExt::read_until) method. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct ReadUntil<'a, R: ?Sized + Unpin> { + reader: &'a mut R, + byte: u8, + buf: &'a mut Vec, + read: usize, +} + +impl Unpin for ReadUntil<'_, R> {} + +pub(crate) fn read_until<'a, R>( + reader: &'a mut R, + byte: u8, + buf: &'a mut Vec, +) -> ReadUntil<'a, R> +where + R: AsyncBufRead + ?Sized + Unpin, +{ + ReadUntil { + reader, + byte, + buf, + read: 0, + } +} + +pub(super) fn read_until_internal( + mut reader: Pin<&mut R>, + cx: &mut Context<'_>, + byte: u8, + buf: &mut Vec, + read: &mut usize, +) -> Poll> { + loop { + let (done, used) = { + let available = 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 Future for ReadUntil<'_, R> { + type Output = io::Result; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let Self { + reader, + byte, + buf, + read, + } = &mut *self; + read_until_internal(Pin::new(reader), cx, *byte, buf, read) + } +} diff --git a/tokio/tests/io_read_until.rs b/tokio/tests/io_read_until.rs new file mode 100644 index 000000000..22945f55b --- /dev/null +++ b/tokio/tests/io_read_until.rs @@ -0,0 +1,56 @@ +#![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_until() { + 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 = vec![]; + let mut rd = Rd { + val: b"hello world", + }; + + let n = assert_ok!(rd.read_until(b' ', &mut buf).await); + assert_eq!(n, 6); + assert_eq!(buf, b"hello "); + buf.clear(); + let n = assert_ok!(rd.read_until(b' ', &mut buf).await); + assert_eq!(n, 5); + assert_eq!(buf, b"world"); + buf.clear(); + let n = assert_ok!(rd.read_until(b' ', &mut buf).await); + assert_eq!(n, 0); + assert_eq!(buf, []); +}