tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-17 00:00:11 +02:00
Code Issues Packages Projects Releases Wiki Activity

chore: migrate from pin-project to pin-project-lite (#1778)

Browse Source
This commit is contained in:
Taiki Endo
2019-11-16 09:14:40 -08:00
committed by Carl Lerche
parent 19f1fc36bd
commit 320c84a433
11 changed files with 226 additions and 263 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio-util/Cargo.toml
+1 -1
View File
@@ -26,7 +26,7 @@ bytes = "0.4.7"
futures-core = "0.3.0"
futures-sink = "0.3.0"
log = "0.4"
pin-project = "0.4"
pin-project-lite = "0.1.1"
[dev-dependencies]
tokio = { version = "=0.2.0-alpha.6", path = "../tokio" }
tokio-util/src/codec/framed.rs
+16 -14
View File
@@ -8,27 +8,29 @@ use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite};
use bytes::BytesMut;
use futures_core::Stream;
use futures_sink::Sink;
use pin_project::pin_project;
use pin_project_lite::pin_project;
use std::fmt;
use std::io::{self, BufRead, Read, Write};
use std::pin::Pin;
use std::task::{Context, Poll};
/// A unified `Stream` and `Sink` interface to an underlying I/O object, using
/// the `Encoder` and `Decoder` traits to encode and decode frames.
///
/// You can create a `Framed` instance by using the `AsyncRead::framed` adapter.
#[pin_project]
pub struct Framed<T, U> {
#[pin]
inner: FramedRead2<FramedWrite2<Fuse<T, U>>>,
pin_project! {
/// A unified `Stream` and `Sink` interface to an underlying I/O object, using
/// the `Encoder` and `Decoder` traits to encode and decode frames.
///
/// You can create a `Framed` instance by using the `AsyncRead::framed` adapter.
pub struct Framed<T, U> {
#[pin]
inner: FramedRead2<FramedWrite2<Fuse<T, U>>>,
}
}
#[pin_project]
pub(crate) struct Fuse<T, U> {
#[pin]
pub(crate) io: T,
pub(crate) codec: U,
pin_project! {
pub(crate) struct Fuse<T, U> {
#[pin]
pub(crate) io: T,
pub(crate) codec: U,
}
}
/// Abstracts over `FramedRead2` being either `FramedRead2<FramedWrite2<Fuse<T, U>>>` or
tokio-util/src/codec/framed_read.rs
+15 -13
View File
@@ -7,25 +7,27 @@ use bytes::BytesMut;
use futures_core::Stream;
use futures_sink::Sink;
use log::trace;
use pin_project::pin_project;
use pin_project_lite::pin_project;
use std::fmt;
use std::pin::Pin;
use std::task::{Context, Poll};
/// A `Stream` of messages decoded from an `AsyncRead`.
#[pin_project]
pub struct FramedRead<T, D> {
#[pin]
inner: FramedRead2<Fuse<T, D>>,
pin_project! {
/// A `Stream` of messages decoded from an `AsyncRead`.
pub struct FramedRead<T, D> {
#[pin]
inner: FramedRead2<Fuse<T, D>>,
}
}
#[pin_project]
pub(crate) struct FramedRead2<T> {
#[pin]
inner: T,
eof: bool,
is_readable: bool,
buffer: BytesMut,
pin_project! {
pub(crate) struct FramedRead2<T> {
#[pin]
inner: T,
eof: bool,
is_readable: bool,
buffer: BytesMut,
}
}
const INITIAL_CAPACITY: usize = 8 * 1024;
tokio-util/src/codec/framed_write.rs
+13 -11
View File
@@ -8,24 +8,26 @@ use bytes::BytesMut;
use futures_core::{ready, Stream};
use futures_sink::Sink;
use log::trace;
use pin_project::pin_project;
use pin_project_lite::pin_project;
use std::fmt;
use std::io::{self, BufRead, Read};
use std::pin::Pin;
use std::task::{Context, Poll};
/// A `Sink` of frames encoded to an `AsyncWrite`.
#[pin_project]
pub struct FramedWrite<T, E> {
#[pin]
inner: FramedWrite2<Fuse<T, E>>,
pin_project! {
/// A `Sink` of frames encoded to an `AsyncWrite`.
pub struct FramedWrite<T, E> {
#[pin]
inner: FramedWrite2<Fuse<T, E>>,
}
}
#[pin_project]
pub(crate) struct FramedWrite2<T> {
#[pin]
inner: T,
buffer: BytesMut,
pin_project! {
pub(crate) struct FramedWrite2<T> {
#[pin]
inner: T,
buffer: BytesMut,
}
}
const INITIAL_CAPACITY: usize = 8 * 1024;
tokio/Cargo.toml
+2 -3
View File
@@ -41,7 +41,7 @@ blocking = ["rt-core"]
dns = ["blocking"]
fs = ["blocking"]
io-driver = ["mio", "lazy_static", "sync"] # TODO: get rid of sync
io-util = ["pin-project", "pin-project-lite", "memchr"]
io-util = ["memchr"]
macros = ["tokio-macros"]
net = ["dns", "tcp", "udp", "uds"]
process = [
@@ -88,6 +88,7 @@ tokio-macros = { version = "=0.2.0-alpha.6", optional = true, path = "../tokio-m
bytes = "0.4"
iovec = "0.1"
pin-project-lite = "0.1.1"
# Everything else is optional...
fnv = { version = "1.0.6", optional = true }
@@ -96,8 +97,6 @@ lazy_static = { version = "1.0.2", optional = true }
memchr = { version = "2.2", optional = true }
mio = { version = "0.6.14", optional = true }
num_cpus = { version = "1.8.0", optional = true }
pin-project = { version = "0.4", optional = true }
pin-project-lite = { version = "0.1", optional = true }
# Backs `DelayQueue`
slab = { version = "0.4.1", optional = true }
tokio/src/future/try_join.rs
+39 -72
View File
@@ -1,5 +1,6 @@
use crate::future::{maybe_done, MaybeDone};
use pin_project_lite::pin_project;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
@@ -21,15 +22,20 @@ where
}
}
pub(crate) struct TryJoin3<F1, F2, F3>
where
F1: Future,
F2: Future,
F3: Future,
{
future1: MaybeDone<F1>,
future2: MaybeDone<F2>,
future3: MaybeDone<F3>,
pin_project! {
pub(crate) struct TryJoin3<F1, F2, F3>
where
F1: Future,
F2: Future,
F3: Future,
{
#[pin]
future1: MaybeDone<F1>,
#[pin]
future2: MaybeDone<F2>,
#[pin]
future3: MaybeDone<F3>,
}
}
impl<T1, F1, T2, F2, T3, F3, E> Future for TryJoin3<F1, F2, F3>
@@ -43,73 +49,34 @@ where
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut all_done = true;
// Safety: the fn takes `Pin`, we don't move any data out of `self`.
unsafe {
let me = self.get_unchecked_mut();
let mut me = self.project();
if Pin::new_unchecked(&mut me.future1).poll(cx).is_pending() {
all_done = false;
} else if Pin::new_unchecked(&mut me.future1)
.output_mut()
.unwrap()
.is_err()
{
return Poll::Ready(Err(Pin::new_unchecked(&mut me.future1)
.take_output()
.unwrap()
.err()
.unwrap()));
}
if me.future1.as_mut().poll(cx).is_pending() {
all_done = false;
} else if me.future1.as_mut().output_mut().unwrap().is_err() {
return Poll::Ready(Err(me.future1.take_output().unwrap().err().unwrap()));
}
if Pin::new_unchecked(&mut me.future2).poll(cx).is_pending() {
all_done = false;
} else if Pin::new_unchecked(&mut me.future2)
.output_mut()
.unwrap()
.is_err()
{
return Poll::Ready(Err(Pin::new_unchecked(&mut me.future2)
.take_output()
.unwrap()
.err()
.unwrap()));
}
if me.future2.as_mut().poll(cx).is_pending() {
all_done = false;
} else if me.future2.as_mut().output_mut().unwrap().is_err() {
return Poll::Ready(Err(me.future2.take_output().unwrap().err().unwrap()));
}
if Pin::new_unchecked(&mut me.future3).poll(cx).is_pending() {
all_done = false;
} else if Pin::new_unchecked(&mut me.future3)
.output_mut()
.unwrap()
.is_err()
{
return Poll::Ready(Err(Pin::new_unchecked(&mut me.future3)
.take_output()
.unwrap()
.err()
.unwrap()));
}
if me.future3.as_mut().poll(cx).is_pending() {
all_done = false;
} else if me.future3.as_mut().output_mut().unwrap().is_err() {
return Poll::Ready(Err(me.future3.take_output().unwrap().err().unwrap()));
}
if all_done {
Poll::Ready(Ok((
Pin::new_unchecked(&mut me.future1)
.take_output()
.unwrap()
.ok()
.unwrap(),
Pin::new_unchecked(&mut me.future2)
.take_output()
.unwrap()
.ok()
.unwrap(),
Pin::new_unchecked(&mut me.future3)
.take_output()
.unwrap()
.ok()
.unwrap(),
)))
} else {
Poll::Pending
}
if all_done {
Poll::Ready(Ok((
me.future1.take_output().unwrap().ok().unwrap(),
me.future2.take_output().unwrap().ok().unwrap(),
me.future3.take_output().unwrap().ok().unwrap(),
)))
} else {
Poll::Pending
}
}
}
tokio/src/io/util/buf_reader.rs
+31 -37
View File
@@ -1,35 +1,36 @@
use crate::io::util::DEFAULT_BUF_SIZE;
use crate::io::{AsyncBufRead, AsyncRead, AsyncWrite};
use pin_project::{pin_project, project};
use pin_project_lite::pin_project;
use std::io::{self, Read};
use std::pin::Pin;
use std::task::{Context, Poll};
use std::{cmp, fmt};
/// The `BufReader` struct adds buffering to any reader.
///
/// It can be excessively inefficient to work directly with a [`AsyncRead`]
/// instance. A `BufReader` performs large, infrequent reads on the underlying
/// [`AsyncRead`] and maintains an in-memory buffer of the results.
///
/// `BufReader` can improve the speed of programs that make *small* and
/// *repeated* read calls to the same file or network socket. It does not
/// help when reading very large amounts at once, or reading just one or a few
/// times. It also provides no advantage when reading from a source that is
/// already in memory, like a `Vec<u8>`.
///
/// When the `BufReader` is dropped, the contents of its buffer will be
/// discarded. Creating multiple instances of a `BufReader` on the same
/// stream can cause data loss.
// TODO: Examples
#[pin_project]
pub struct BufReader<R> {
#[pin]
pub(super) inner: R,
pub(super) buf: Box<[u8]>,
pub(super) pos: usize,
pub(super) cap: usize,
pin_project! {
/// The `BufReader` struct adds buffering to any reader.
///
/// It can be excessively inefficient to work directly with a [`AsyncRead`]
/// instance. A `BufReader` performs large, infrequent reads on the underlying
/// [`AsyncRead`] and maintains an in-memory buffer of the results.
///
/// `BufReader` can improve the speed of programs that make *small* and
/// *repeated* read calls to the same file or network socket. It does not
/// help when reading very large amounts at once, or reading just one or a few
/// times. It also provides no advantage when reading from a source that is
/// already in memory, like a `Vec<u8>`.
///
/// When the `BufReader` is dropped, the contents of its buffer will be
/// discarded. Creating multiple instances of a `BufReader` on the same
/// stream can cause data loss.
// TODO: Examples
pub struct BufReader<R> {
#[pin]
pub(super) inner: R,
pub(super) buf: Box<[u8]>,
pub(super) pos: usize,
pub(super) cap: usize,
}
}
impl<R: AsyncRead> BufReader<R> {
@@ -125,26 +126,19 @@ impl<R: AsyncRead> AsyncRead for BufReader<R> {
}
impl<R: AsyncRead> AsyncBufRead for BufReader<R> {
#[project]
fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
#[project]
let BufReader {
inner,
buf,
cap,
pos,
} = self.project();
let me = self.project();
// If we've reached the end of our internal buffer then we need to fetch
// some more data from the underlying reader.
// Branch using `>=` instead of the more correct `==`
// to tell the compiler that the pos..cap slice is always valid.
if *pos >= *cap {
debug_assert!(*pos == *cap);
*cap = ready!(inner.poll_read(cx, buf))?;
*pos = 0;
if *me.pos >= *me.cap {
debug_assert!(*me.pos == *me.cap);
*me.cap = ready!(me.inner.poll_read(cx, me.buf))?;
*me.pos = 0;
}
Poll::Ready(Ok(&buf[*pos..*cap]))
Poll::Ready(Ok(&me.buf[*me.pos..*me.cap]))
}
fn consume(self: Pin<&mut Self>, amt: usize) {
tokio/src/io/util/buf_stream.rs
+41 -33
View File
@@ -1,64 +1,70 @@
use crate::io::util::{BufReader, BufWriter};
use crate::io::{AsyncBufRead, AsyncRead, AsyncWrite};
use pin_project::pin_project;
use pin_project_lite::pin_project;
use std::io::{self};
use std::{
pin::Pin,
task::{Context, Poll},
};
/// Wraps a type that is [`AsyncWrite`] and [`AsyncRead`], and buffers its input and output.
///
/// It can be excessively inefficient to work directly with something that implements [`AsyncWrite`]
/// and [`AsyncRead`]. For example, every `write`, however small, has to traverse the syscall
/// interface, and similarly, every read has to do the same. The [`BufWriter`] and [`BufReader`]
/// types aid with these problems respectively, but do so in only one direction. `BufStream` wraps
/// one in the other so that both directions are buffered. See their documentation for details.
#[pin_project]
#[derive(Debug)]
pub struct BufStream<RW>(#[pin] BufReader<BufWriter<RW>>);
pin_project! {
/// Wraps a type that is [`AsyncWrite`] and [`AsyncRead`], and buffers its input and output.
///
/// It can be excessively inefficient to work directly with something that implements [`AsyncWrite`]
/// and [`AsyncRead`]. For example, every `write`, however small, has to traverse the syscall
/// interface, and similarly, every read has to do the same. The [`BufWriter`] and [`BufReader`]
/// types aid with these problems respectively, but do so in only one direction. `BufStream` wraps
/// one in the other so that both directions are buffered. See their documentation for details.
#[derive(Debug)]
pub struct BufStream<RW> {
#[pin]
inner: BufReader<BufWriter<RW>>,
}
}
impl<RW: AsyncRead + AsyncWrite> BufStream<RW> {
/// Wrap a type in both [`BufWriter`] and [`BufReader`].
///
/// See the documentation for those types and [`BufStream`] for details.
pub fn new(stream: RW) -> BufStream<RW> {
BufStream(BufReader::new(BufWriter::new(stream)))
BufStream {
inner: BufReader::new(BufWriter::new(stream)),
}
}
/// Gets a reference to the underlying I/O object.
///
/// It is inadvisable to directly read from the underlying I/O object.
pub fn get_ref(&self) -> &RW {
self.0.get_ref().get_ref()
self.inner.get_ref().get_ref()
}
/// Gets a mutable reference to the underlying I/O object.
///
/// It is inadvisable to directly read from the underlying I/O object.
pub fn get_mut(&mut self) -> &mut RW {
self.0.get_mut().get_mut()
self.inner.get_mut().get_mut()
}
/// Gets a pinned mutable reference to the underlying I/O object.
///
/// It is inadvisable to directly read from the underlying I/O object.
pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut RW> {
self.project().0.get_pin_mut().get_pin_mut()
self.project().inner.get_pin_mut().get_pin_mut()
}
/// Consumes this `BufStream`, returning the underlying I/O object.
///
/// Note that any leftover data in the internal buffer is lost.
pub fn into_inner(self) -> RW {
self.0.into_inner().into_inner()
self.inner.into_inner().into_inner()
}
}
impl<RW> From<BufReader<BufWriter<RW>>> for BufStream<RW> {
fn from(b: BufReader<BufWriter<RW>>) -> Self {
BufStream(b)
BufStream { inner: b }
}
}
@@ -77,16 +83,18 @@ impl<RW> From<BufWriter<BufReader<RW>>> for BufStream<RW> {
written,
} = b;
BufStream(BufReader {
inner: BufWriter {
inner,
buf: wbuf,
written,
BufStream {
inner: BufReader {
inner: BufWriter {
inner,
buf: wbuf,
written,
},
buf: rbuf,
pos,
cap,
},
buf: rbuf,
pos,
cap,
})
}
}
}
@@ -96,15 +104,15 @@ impl<RW: AsyncRead + AsyncWrite> AsyncWrite for BufStream<RW> {
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
self.project().0.poll_write(cx, buf)
self.project().inner.poll_write(cx, buf)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.project().0.poll_flush(cx)
self.project().inner.poll_flush(cx)
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.project().0.poll_shutdown(cx)
self.project().inner.poll_shutdown(cx)
}
}
@@ -114,22 +122,22 @@ impl<RW: AsyncRead + AsyncWrite> AsyncRead for BufStream<RW> {
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
self.project().0.poll_read(cx, buf)
self.project().inner.poll_read(cx, buf)
}
// we can't skip unconditionally because of the large buffer case in read.
unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
self.0.prepare_uninitialized_buffer(buf)
self.inner.prepare_uninitialized_buffer(buf)
}
}
impl<RW: AsyncBufRead + AsyncRead + AsyncWrite> AsyncBufRead for BufStream<RW> {
fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
self.project().0.poll_fill_buf(cx)
self.project().inner.poll_fill_buf(cx)
}
fn consume(self: Pin<&mut Self>, amt: usize) {
self.project().0.consume(amt)
self.project().inner.consume(amt)
}
}
tokio/src/io/util/buf_writer.rs
+37 -42
View File
@@ -1,39 +1,40 @@
use crate::io::util::DEFAULT_BUF_SIZE;
use crate::io::{AsyncBufRead, AsyncRead, AsyncWrite};
use pin_project::{pin_project, project};
use pin_project_lite::pin_project;
use std::fmt;
use std::io::{self, Write};
use std::pin::Pin;
use std::task::{Context, Poll};
/// Wraps a writer and buffers its output.
///
/// It can be excessively inefficient to work directly with something that
/// implements [`AsyncWrite`]. A `BufWriter` keeps an in-memory buffer of data and
/// writes it to an underlying writer in large, infrequent batches.
///
/// `BufWriter` can improve the speed of programs that make *small* and
/// *repeated* write calls to the same file or network socket. It does not
/// help when writing very large amounts at once, or writing just one or a few
/// times. It also provides no advantage when writing to a destination that is
/// in memory, like a `Vec<u8>`.
///
/// When the `BufWriter` is dropped, the contents of its buffer will be
/// discarded. Creating multiple instances of a `BufWriter` on the same
/// stream can cause data loss. If you need to write out the contents of its
/// buffer, you must manually call flush before the writer is dropped.
///
/// [`AsyncWrite`]: AsyncWrite
/// [`flush`]: super::AsyncWriteExt::flush
///
// TODO: Examples
#[pin_project]
pub struct BufWriter<W> {
#[pin]
pub(super) inner: W,
pub(super) buf: Vec<u8>,
pub(super) written: usize,
pin_project! {
/// Wraps a writer and buffers its output.
///
/// It can be excessively inefficient to work directly with something that
/// implements [`AsyncWrite`]. A `BufWriter` keeps an in-memory buffer of data and
/// writes it to an underlying writer in large, infrequent batches.
///
/// `BufWriter` can improve the speed of programs that make *small* and
/// *repeated* write calls to the same file or network socket. It does not
/// help when writing very large amounts at once, or writing just one or a few
/// times. It also provides no advantage when writing to a destination that is
/// in memory, like a `Vec<u8>`.
///
/// When the `BufWriter` is dropped, the contents of its buffer will be
/// discarded. Creating multiple instances of a `BufWriter` on the same
/// stream can cause data loss. If you need to write out the contents of its
/// buffer, you must manually call flush before the writer is dropped.
///
/// [`AsyncWrite`]: AsyncWrite
/// [`flush`]: super::AsyncWriteExt::flush
///
// TODO: Examples
pub struct BufWriter<W> {
#[pin]
pub(super) inner: W,
pub(super) buf: Vec<u8>,
pub(super) written: usize,
}
}
impl<W: AsyncWrite> BufWriter<W> {
@@ -52,19 +53,13 @@ impl<W: AsyncWrite> BufWriter<W> {
}
}
#[project]
fn flush_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
#[project]
let BufWriter {
mut inner,
buf,
written,
} = self.project();
let mut me = self.project();
let len = buf.len();
let len = me.buf.len();
let mut ret = Ok(());
while *written < len {
match ready!(inner.as_mut().poll_write(cx, &buf[*written..])) {
while *me.written < len {
match ready!(me.inner.as_mut().poll_write(cx, &me.buf[*me.written..])) {
Ok(0) => {
ret = Err(io::Error::new(
io::ErrorKind::WriteZero,
@@ -72,17 +67,17 @@ impl<W: AsyncWrite> BufWriter<W> {
));
break;
}
Ok(n) => *written += n,
Ok(n) => *me.written += n,
Err(e) => {
ret = Err(e);
break;
}
}
}
if *written > 0 {
buf.drain(..*written);
if *me.written > 0 {
me.buf.drain(..*me.written);
}
*written = 0;
*me.written = 0;
Poll::Ready(ret)
}
tokio/src/io/util/chain.rs
+16 -21
View File
@@ -1,20 +1,21 @@
use crate::io::{AsyncBufRead, AsyncRead};
use pin_project::{pin_project, project};
use pin_project_lite::pin_project;
use std::fmt;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
/// Stream for the [`chain`](super::AsyncReadExt::chain) method.
#[pin_project]
#[must_use = "streams do nothing unless polled"]
pub struct Chain<T, U> {
#[pin]
first: T,
#[pin]
second: U,
done_first: bool,
pin_project! {
/// Stream for the [`chain`](super::AsyncReadExt::chain) method.
#[must_use = "streams do nothing unless polled"]
pub struct Chain<T, U> {
#[pin]
first: T,
#[pin]
second: U,
done_first: bool,
}
}
pub(super) fn chain<T, U>(first: T, second: U) -> Chain<T, U>
@@ -104,24 +105,18 @@ where
T: AsyncBufRead,
U: AsyncBufRead,
{
#[project]
fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
#[project]
let Chain {
first,
second,
done_first,
} = self.project();
let me = self.project();
if !*done_first {
match ready!(first.poll_fill_buf(cx)?) {
if !*me.done_first {
match ready!(me.first.poll_fill_buf(cx)?) {
buf if buf.is_empty() => {
*done_first = true;
*me.done_first = true;
}
buf => return Poll::Ready(Ok(buf)),
}
}
second.poll_fill_buf(cx)
me.second.poll_fill_buf(cx)
}
fn consume(self: Pin<&mut Self>, amt: usize) {
tokio/src/io/util/take.rs
+15 -16
View File
@@ -1,19 +1,20 @@
use crate::io::{AsyncBufRead, AsyncRead};
use pin_project::{pin_project, project};
use pin_project_lite::pin_project;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::{cmp, io};
/// Stream for the [`take`](super::AsyncReadExt::take) method.
#[pin_project]
#[derive(Debug)]
#[must_use = "streams do nothing unless you `.await` or poll them"]
pub struct Take<R> {
#[pin]
inner: R,
// Add '_' to avoid conflicts with `limit` method.
limit_: u64,
pin_project! {
/// Stream for the [`take`](super::AsyncReadExt::take) method.
#[derive(Debug)]
#[must_use = "streams do nothing unless you `.await` or poll them"]
pub struct Take<R> {
#[pin]
inner: R,
// Add '_' to avoid conflicts with `limit` method.
limit_: u64,
}
}
pub(super) fn take<R: AsyncRead>(inner: R, limit: u64) -> Take<R> {
@@ -95,18 +96,16 @@ impl<R: AsyncRead> AsyncRead for Take<R> {
}
impl<R: AsyncBufRead> AsyncBufRead for Take<R> {
#[project]
fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
#[project]
let Take { inner, limit_ } = self.project();
let me = self.project();
// Don't call into inner reader at all at EOF because it may still block
if *limit_ == 0 {
if *me.limit_ == 0 {
return Poll::Ready(Ok(&[]));
}
let buf = ready!(inner.poll_fill_buf(cx)?);
let cap = cmp::min(buf.len() as u64, *limit_) as usize;
let buf = ready!(me.inner.poll_fill_buf(cx)?);
let cap = cmp::min(buf.len() as u64, *me.limit_) as usize;
Poll::Ready(Ok(&buf[..cap]))
}