diff options
author | Joel Galenson <jgalenson@google.com> | 2020-10-23 09:39:31 -0700 |
---|---|---|
committer | Joel Galenson <jgalenson@google.com> | 2020-10-23 09:52:09 -0700 |
commit | d5495b03381a3ebe0805db353d198b285b535b5c (patch) | |
tree | 778b8524d15fca8b73db0253ee0e1919d0848bb6 /src/io/util/read_exact.rs | |
parent | ba45c5bedf31df8562364c61d3dfb5262f10642e (diff) | |
download | tokio-d5495b03381a3ebe0805db353d198b285b535b5c.tar.gz |
Update to tokio-0.3.1 and add new features
Test: Build
Change-Id: I5b5b9b386a21982a019653d0cf0bd3afc505cfac
Diffstat (limited to 'src/io/util/read_exact.rs')
-rw-r--r-- | src/io/util/read_exact.rs | 46 |
1 files changed, 19 insertions, 27 deletions
diff --git a/src/io/util/read_exact.rs b/src/io/util/read_exact.rs index 86b8412..1e8150e 100644 --- a/src/io/util/read_exact.rs +++ b/src/io/util/read_exact.rs @@ -1,7 +1,9 @@ -use crate::io::AsyncRead; +use crate::io::{AsyncRead, ReadBuf}; +use pin_project_lite::pin_project; use std::future::Future; use std::io; +use std::marker::PhantomPinned; use std::marker::Unpin; use std::pin::Pin; use std::task::{Context, Poll}; @@ -17,12 +19,12 @@ where { ReadExact { reader, - buf, - pos: 0, + buf: ReadBuf::new(buf), + _pin: PhantomPinned, } } -cfg_io_util! { +pin_project! { /// Creates a future which will read exactly enough bytes to fill `buf`, /// returning an error if EOF is hit sooner. /// @@ -31,8 +33,10 @@ cfg_io_util! { #[must_use = "futures do nothing unless you `.await` or poll them"] pub struct ReadExact<'a, A: ?Sized> { reader: &'a mut A, - buf: &'a mut [u8], - pos: usize, + buf: ReadBuf<'a>, + // Make this future `!Unpin` for compatibility with async trait methods. + #[pin] + _pin: PhantomPinned, } } @@ -46,32 +50,20 @@ where { type Output = io::Result<usize>; - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> { + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> { + let mut me = self.project(); + loop { // if our buffer is empty, then we need to read some data to continue. - if self.pos < self.buf.len() { - let me = &mut *self; - let n = ready!(Pin::new(&mut *me.reader).poll_read(cx, &mut me.buf[me.pos..]))?; - me.pos += n; - if n == 0 { + let rem = me.buf.remaining(); + if rem != 0 { + ready!(Pin::new(&mut *me.reader).poll_read(cx, &mut me.buf))?; + if me.buf.remaining() == rem { return Err(eof()).into(); } - } - - if self.pos >= self.buf.len() { - return Poll::Ready(Ok(self.pos)); + } else { + return Poll::Ready(Ok(me.buf.capacity())); } } } } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn assert_unpin() { - use std::marker::PhantomPinned; - crate::is_unpin::<ReadExact<'_, PhantomPinned>>(); - } -} |