summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorLucio Franco <luciofranco14@gmail.com>2020-12-11 23:40:24 -0500
committerGitHub <noreply@github.com>2020-12-11 20:40:24 -0800
commitb01b2dacf2e4136c0237977dac27a3688467d2ea (patch)
tree1e0bbb86739731038cc9fd69fe112cad54662d16
parent68717c7efaced76651915696495dcb04c890be25 (diff)
net: update `TcpStream::poll_peek` to use `ReadBuf` (#3259)
Closes #2987
-rw-r--r--tokio/src/net/tcp/split.rs12
-rw-r--r--tokio/src/net/tcp/split_owned.rs12
-rw-r--r--tokio/src/net/tcp/stream.rs21
3 files changed, 35 insertions, 10 deletions
diff --git a/tokio/src/net/tcp/split.rs b/tokio/src/net/tcp/split.rs
index 28c94eb4..296b469d 100644
--- a/tokio/src/net/tcp/split.rs
+++ b/tokio/src/net/tcp/split.rs
@@ -60,7 +60,7 @@ impl ReadHalf<'_> {
/// # Examples
///
/// ```no_run
- /// use tokio::io;
+ /// use tokio::io::{self, ReadBuf};
/// use tokio::net::TcpStream;
///
/// use futures::future::poll_fn;
@@ -70,6 +70,7 @@ impl ReadHalf<'_> {
/// let mut stream = TcpStream::connect("127.0.0.1:8000").await?;
/// let (mut read_half, _) = stream.split();
/// let mut buf = [0; 10];
+ /// let mut buf = ReadBuf::new(&mut buf);
///
/// poll_fn(|cx| {
/// read_half.poll_peek(cx, &mut buf)
@@ -80,7 +81,11 @@ impl ReadHalf<'_> {
/// ```
///
/// [`TcpStream::poll_peek`]: TcpStream::poll_peek
- pub fn poll_peek(&mut self, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
+ pub fn poll_peek(
+ &mut self,
+ cx: &mut Context<'_>,
+ buf: &mut ReadBuf<'_>,
+ ) -> Poll<io::Result<usize>> {
self.0.poll_peek(cx, buf)
}
@@ -124,7 +129,8 @@ impl ReadHalf<'_> {
/// [`read`]: fn@crate::io::AsyncReadExt::read
/// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt
pub async fn peek(&mut self, buf: &mut [u8]) -> io::Result<usize> {
- poll_fn(|cx| self.poll_peek(cx, buf)).await
+ let mut buf = ReadBuf::new(buf);
+ poll_fn(|cx| self.poll_peek(cx, &mut buf)).await
}
}
diff --git a/tokio/src/net/tcp/split_owned.rs b/tokio/src/net/tcp/split_owned.rs
index 8d77c8ca..725d7411 100644
--- a/tokio/src/net/tcp/split_owned.rs
+++ b/tokio/src/net/tcp/split_owned.rs
@@ -115,7 +115,7 @@ impl OwnedReadHalf {
/// # Examples
///
/// ```no_run
- /// use tokio::io;
+ /// use tokio::io::{self, ReadBuf};
/// use tokio::net::TcpStream;
///
/// use futures::future::poll_fn;
@@ -125,6 +125,7 @@ impl OwnedReadHalf {
/// let stream = TcpStream::connect("127.0.0.1:8000").await?;
/// let (mut read_half, _) = stream.into_split();
/// let mut buf = [0; 10];
+ /// let mut buf = ReadBuf::new(&mut buf);
///
/// poll_fn(|cx| {
/// read_half.poll_peek(cx, &mut buf)
@@ -135,7 +136,11 @@ impl OwnedReadHalf {
/// ```
///
/// [`TcpStream::poll_peek`]: TcpStream::poll_peek
- pub fn poll_peek(&mut self, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
+ pub fn poll_peek(
+ &mut self,
+ cx: &mut Context<'_>,
+ buf: &mut ReadBuf<'_>,
+ ) -> Poll<io::Result<usize>> {
self.inner.poll_peek(cx, buf)
}
@@ -179,7 +184,8 @@ impl OwnedReadHalf {
/// [`read`]: fn@crate::io::AsyncReadExt::read
/// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt
pub async fn peek(&mut self, buf: &mut [u8]) -> io::Result<usize> {
- poll_fn(|cx| self.poll_peek(cx, buf)).await
+ let mut buf = ReadBuf::new(buf);
+ poll_fn(|cx| self.poll_peek(cx, &mut buf)).await
}
}
diff --git a/tokio/src/net/tcp/stream.rs b/tokio/src/net/tcp/stream.rs
index 83e9f2a7..c4a0d12c 100644
--- a/tokio/src/net/tcp/stream.rs
+++ b/tokio/src/net/tcp/stream.rs
@@ -291,7 +291,7 @@ impl TcpStream {
/// # Examples
///
/// ```no_run
- /// use tokio::io;
+ /// use tokio::io::{self, ReadBuf};
/// use tokio::net::TcpStream;
///
/// use futures::future::poll_fn;
@@ -300,6 +300,7 @@ impl TcpStream {
/// async fn main() -> io::Result<()> {
/// let stream = TcpStream::connect("127.0.0.1:8000").await?;
/// let mut buf = [0; 10];
+ /// let mut buf = ReadBuf::new(&mut buf);
///
/// poll_fn(|cx| {
/// stream.poll_peek(cx, &mut buf)
@@ -308,12 +309,24 @@ impl TcpStream {
/// Ok(())
/// }
/// ```
- pub fn poll_peek(&self, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
+ pub fn poll_peek(
+ &self,
+ cx: &mut Context<'_>,
+ buf: &mut ReadBuf<'_>,
+ ) -> Poll<io::Result<usize>> {
loop {
let ev = ready!(self.io.registration().poll_read_ready(cx))?;
- match self.io.peek(buf) {
- Ok(ret) => return Poll::Ready(Ok(ret)),
+ let b = unsafe {
+ &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit<u8>] as *mut [u8])
+ };
+
+ match self.io.peek(b) {
+ Ok(ret) => {
+ unsafe { buf.assume_init(ret) };
+ buf.advance(ret);
+ return Poll::Ready(Ok(ret));
+ }
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.registration().clear_readiness(ev);
}