Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
tokio: add AsyncReadExt::read_to_string (#1326)
- Loading branch information
1 parent
a298472
commit a88308e
Showing
4 changed files
with
138 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,72 @@ | ||
use super::read_to_end::read_to_end_internal; | ||
use futures_core::ready; | ||
use std::future::Future; | ||
use std::pin::Pin; | ||
use std::task::{Context, Poll}; | ||
use std::{io, mem, str}; | ||
use tokio_io::AsyncRead; | ||
|
||
/// Future for the [`read_to_string`](super::AsyncReadExt::read_to_string) method. | ||
#[derive(Debug)] | ||
#[must_use = "futures do nothing unless you `.await` or poll them"] | ||
pub struct ReadToString<'a, R: ?Sized + Unpin> { | ||
reader: &'a mut R, | ||
buf: &'a mut String, | ||
bytes: Vec<u8>, | ||
start_len: usize, | ||
} | ||
|
||
impl<R: ?Sized + Unpin> Unpin for ReadToString<'_, R> {} | ||
|
||
pub(crate) fn read_to_string<'a, R>(reader: &'a mut R, buf: &'a mut String) -> ReadToString<'a, R> | ||
where | ||
R: AsyncRead + ?Sized + Unpin, | ||
{ | ||
let start_len = buf.len(); | ||
ReadToString { | ||
reader, | ||
bytes: unsafe { mem::replace(buf.as_mut_vec(), Vec::new()) }, | ||
buf, | ||
start_len, | ||
} | ||
} | ||
|
||
fn read_to_string_internal<R: AsyncRead + ?Sized>( | ||
reader: Pin<&mut R>, | ||
cx: &mut Context<'_>, | ||
buf: &mut String, | ||
bytes: &mut Vec<u8>, | ||
start_len: usize, | ||
) -> Poll<io::Result<usize>> { | ||
let ret = ready!(read_to_end_internal(reader, cx, bytes, start_len)); | ||
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()); | ||
// 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<A> Future for ReadToString<'_, A> | ||
where | ||
A: AsyncRead + ?Sized + Unpin, | ||
{ | ||
type Output = io::Result<usize>; | ||
|
||
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { | ||
let Self { | ||
reader, | ||
buf, | ||
bytes, | ||
start_len, | ||
} = &mut *self; | ||
read_to_string_internal(Pin::new(reader), cx, buf, bytes, *start_len) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,40 @@ | ||
#![deny(warnings, rust_2018_idioms)] | ||
#![feature(async_await)] | ||
|
||
use tokio::io::{AsyncRead, AsyncReadExt}; | ||
use tokio_test::assert_ok; | ||
|
||
use std::pin::Pin; | ||
use std::task::{Context, Poll}; | ||
use std::{cmp, io}; | ||
|
||
#[tokio::test] | ||
async fn read_to_string() { | ||
struct Rd { | ||
val: &'static [u8], | ||
} | ||
|
||
impl AsyncRead for Rd { | ||
fn poll_read( | ||
mut self: Pin<&mut Self>, | ||
_cx: &mut Context<'_>, | ||
buf: &mut [u8], | ||
) -> Poll<io::Result<usize>> { | ||
let me = &mut *self; | ||
let len = cmp::min(buf.len(), me.val.len()); | ||
|
||
buf[..len].copy_from_slice(&me.val[..len]); | ||
me.val = &me.val[len..]; | ||
Poll::Ready(Ok(len)) | ||
} | ||
} | ||
|
||
let mut buf = String::new(); | ||
let mut rd = Rd { | ||
val: b"hello world", | ||
}; | ||
|
||
let n = assert_ok!(rd.read_to_string(&mut buf).await); | ||
assert_eq!(n, 11); | ||
assert_eq!(buf[..], "hello world"[..]); | ||
} |