mirror of
https://github.com/tokio-rs/tokio.git
synced 2026-08-23 00:00:10 +02:00
net: add UnixStream readiness and non-blocking ops (#3246)
This commit is contained in:
cssivision
@@ -1,5 +1,5 @@
|
||||
use crate::future::poll_fn;
|
||||
use crate::io::{AsyncRead, AsyncWrite, PollEvented, ReadBuf};
|
||||
use crate::io::{AsyncRead, AsyncWrite, Interest, PollEvented, ReadBuf, Ready};
|
||||
use crate::net::unix::split::{split, ReadHalf, WriteHalf};
|
||||
use crate::net::unix::split_owned::{split_owned, OwnedReadHalf, OwnedWriteHalf};
|
||||
use crate::net::unix::ucred::{self, UCred};
|
||||
@@ -7,7 +7,7 @@ use crate::net::unix::SocketAddr;
|
||||
|
||||
use std::convert::TryFrom;
|
||||
use std::fmt;
|
||||
use std::io;
|
||||
use std::io::{self, Read, Write};
|
||||
use std::net::Shutdown;
|
||||
use std::os::unix::io::{AsRawFd, RawFd};
|
||||
use std::os::unix::net;
|
||||
@@ -43,6 +43,314 @@ impl UnixStream {
|
||||
Ok(stream)
|
||||
}
|
||||
|
||||
/// Wait for any of the requested ready states.
|
||||
///
|
||||
/// This function is usually paired with `try_read()` or `try_write()`. It
|
||||
/// can be used to concurrently read / write to the same socket on a single
|
||||
/// task without splitting the socket.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// Concurrently read and write to the stream on the same task without
|
||||
/// splitting.
|
||||
///
|
||||
/// ```no_run
|
||||
/// use tokio::io::Interest;
|
||||
/// use tokio::net::UnixStream;
|
||||
/// use std::error::Error;
|
||||
/// use std::io;
|
||||
///
|
||||
/// #[tokio::main]
|
||||
/// async fn main() -> Result<(), Box<dyn Error>> {
|
||||
/// let dir = tempfile::tempdir().unwrap();
|
||||
/// let bind_path = dir.path().join("bind_path");
|
||||
/// let stream = UnixStream::connect(bind_path).await?;
|
||||
///
|
||||
/// loop {
|
||||
/// let ready = stream.ready(Interest::READABLE | Interest::WRITABLE).await?;
|
||||
///
|
||||
/// if ready.is_readable() {
|
||||
/// let mut data = vec![0; 1024];
|
||||
/// // Try to read data, this may still fail with `WouldBlock`
|
||||
/// // if the readiness event is a false positive.
|
||||
/// match stream.try_read(&mut data) {
|
||||
/// Ok(n) => {
|
||||
/// println!("read {} bytes", n);
|
||||
/// }
|
||||
/// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
|
||||
/// continue;
|
||||
/// }
|
||||
/// Err(e) => {
|
||||
/// return Err(e.into());
|
||||
/// }
|
||||
/// }
|
||||
///
|
||||
/// }
|
||||
///
|
||||
/// if ready.is_writable() {
|
||||
/// // Try to write data, this may still fail with `WouldBlock`
|
||||
/// // if the readiness event is a false positive.
|
||||
/// match stream.try_write(b"hello world") {
|
||||
/// Ok(n) => {
|
||||
/// println!("write {} bytes", n);
|
||||
/// }
|
||||
/// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
|
||||
/// continue;
|
||||
/// }
|
||||
/// Err(e) => {
|
||||
/// return Err(e.into());
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
pub async fn ready(&self, interest: Interest) -> io::Result<Ready> {
|
||||
let event = self.io.registration().readiness(interest).await?;
|
||||
Ok(event.ready)
|
||||
}
|
||||
|
||||
/// Wait for the socket to become readable.
|
||||
///
|
||||
/// This function is equivalent to `ready(Interest::READABLE)` and is usually
|
||||
/// paired with `try_read()`.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```no_run
|
||||
/// use tokio::net::UnixStream;
|
||||
/// use std::error::Error;
|
||||
/// use std::io;
|
||||
///
|
||||
/// #[tokio::main]
|
||||
/// async fn main() -> Result<(), Box<dyn Error>> {
|
||||
/// // Connect to a peer
|
||||
/// let dir = tempfile::tempdir().unwrap();
|
||||
/// let bind_path = dir.path().join("bind_path");
|
||||
/// let stream = UnixStream::connect(bind_path).await?;
|
||||
///
|
||||
/// let mut msg = vec![0; 1024];
|
||||
///
|
||||
/// loop {
|
||||
/// // Wait for the socket to be readable
|
||||
/// stream.readable().await?;
|
||||
///
|
||||
/// // Try to read data, this may still fail with `WouldBlock`
|
||||
/// // if the readiness event is a false positive.
|
||||
/// match stream.try_read(&mut msg) {
|
||||
/// Ok(n) => {
|
||||
/// msg.truncate(n);
|
||||
/// break;
|
||||
/// }
|
||||
/// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
|
||||
/// continue;
|
||||
/// }
|
||||
/// Err(e) => {
|
||||
/// return Err(e.into());
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
///
|
||||
/// println!("GOT = {:?}", msg);
|
||||
/// Ok(())
|
||||
/// }
|
||||
/// ```
|
||||
pub async fn readable(&self) -> io::Result<()> {
|
||||
self.ready(Interest::READABLE).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Polls for read readiness.
|
||||
///
|
||||
/// This function is intended for cases where creating and pinning a future
|
||||
/// via [`readable`] is not feasible. Where possible, using [`readable`] is
|
||||
/// preferred, as this supports polling from multiple tasks at once.
|
||||
///
|
||||
/// [`readable`]: method@Self::readable
|
||||
pub fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
self.io.registration().poll_read_ready(cx).map_ok(|_| ())
|
||||
}
|
||||
|
||||
/// Try to read data from the stream into the provided buffer, returning how
|
||||
/// many bytes were read.
|
||||
///
|
||||
/// Receives any pending data from the socket but does not wait for new data
|
||||
/// to arrive. On success, returns the number of bytes read. Because
|
||||
/// `try_read()` is non-blocking, the buffer does not have to be stored by
|
||||
/// the async task and can exist entirely on the stack.
|
||||
///
|
||||
/// Usually, [`readable()`] or [`ready()`] is used with this function.
|
||||
///
|
||||
/// [`readable()`]: UnixStream::readable()
|
||||
/// [`ready()`]: UnixStream::ready()
|
||||
///
|
||||
/// # Return
|
||||
///
|
||||
/// If data is successfully read, `Ok(n)` is returned, where `n` is the
|
||||
/// number of bytes read. `Ok(0)` indicates the stream's read half is closed
|
||||
/// and will no longer yield data. If the stream is not ready to read data
|
||||
/// `Err(io::ErrorKind::WouldBlock)` is returned.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```no_run
|
||||
/// use tokio::net::UnixStream;
|
||||
/// use std::error::Error;
|
||||
/// use std::io;
|
||||
///
|
||||
/// #[tokio::main]
|
||||
/// async fn main() -> Result<(), Box<dyn Error>> {
|
||||
/// // Connect to a peer
|
||||
/// let dir = tempfile::tempdir().unwrap();
|
||||
/// let bind_path = dir.path().join("bind_path");
|
||||
/// let stream = UnixStream::connect(bind_path).await?;
|
||||
///
|
||||
/// loop {
|
||||
/// // Wait for the socket to be readable
|
||||
/// stream.readable().await?;
|
||||
///
|
||||
/// // Creating the buffer **after** the `await` prevents it from
|
||||
/// // being stored in the async task.
|
||||
/// let mut buf = [0; 4096];
|
||||
///
|
||||
/// // Try to read data, this may still fail with `WouldBlock`
|
||||
/// // if the readiness event is a false positive.
|
||||
/// match stream.try_read(&mut buf) {
|
||||
/// Ok(0) => break,
|
||||
/// Ok(n) => {
|
||||
/// println!("read {} bytes", n);
|
||||
/// }
|
||||
/// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
|
||||
/// continue;
|
||||
/// }
|
||||
/// Err(e) => {
|
||||
/// return Err(e.into());
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
///
|
||||
/// Ok(())
|
||||
/// }
|
||||
/// ```
|
||||
pub fn try_read(&self, buf: &mut [u8]) -> io::Result<usize> {
|
||||
self.io
|
||||
.registration()
|
||||
.try_io(Interest::READABLE, || (&*self.io).read(buf))
|
||||
}
|
||||
|
||||
/// Wait for the socket to become writable.
|
||||
///
|
||||
/// This function is equivalent to `ready(Interest::WRITABLE)` and is usually
|
||||
/// paired with `try_write()`.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```no_run
|
||||
/// use tokio::net::UnixStream;
|
||||
/// use std::error::Error;
|
||||
/// use std::io;
|
||||
///
|
||||
/// #[tokio::main]
|
||||
/// async fn main() -> Result<(), Box<dyn Error>> {
|
||||
/// // Connect to a peer
|
||||
/// let dir = tempfile::tempdir().unwrap();
|
||||
/// let bind_path = dir.path().join("bind_path");
|
||||
/// let stream = UnixStream::connect(bind_path).await?;
|
||||
///
|
||||
/// loop {
|
||||
/// // Wait for the socket to be writable
|
||||
/// stream.writable().await?;
|
||||
///
|
||||
/// // Try to write data, this may still fail with `WouldBlock`
|
||||
/// // if the readiness event is a false positive.
|
||||
/// match stream.try_write(b"hello world") {
|
||||
/// Ok(n) => {
|
||||
/// break;
|
||||
/// }
|
||||
/// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
|
||||
/// continue;
|
||||
/// }
|
||||
/// Err(e) => {
|
||||
/// return Err(e.into());
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
///
|
||||
/// Ok(())
|
||||
/// }
|
||||
/// ```
|
||||
pub async fn writable(&self) -> io::Result<()> {
|
||||
self.ready(Interest::WRITABLE).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Polls for write readiness.
|
||||
///
|
||||
/// This function is intended for cases where creating and pinning a future
|
||||
/// via [`writable`] is not feasible. Where possible, using [`writable`] is
|
||||
/// preferred, as this supports polling from multiple tasks at once.
|
||||
///
|
||||
/// [`writable`]: method@Self::writable
|
||||
pub fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
self.io.registration().poll_write_ready(cx).map_ok(|_| ())
|
||||
}
|
||||
|
||||
/// Try to write a buffer to the stream, returning how many bytes were
|
||||
/// written.
|
||||
///
|
||||
/// The function will attempt to write the entire contents of `buf`, but
|
||||
/// only part of the buffer may be written.
|
||||
///
|
||||
/// This function is usually paired with `writable()`.
|
||||
///
|
||||
/// # Return
|
||||
///
|
||||
/// If data is successfully written, `Ok(n)` is returned, where `n` is the
|
||||
/// number of bytes written. If the stream is not ready to write data,
|
||||
/// `Err(io::ErrorKind::WouldBlock)` is returned.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```no_run
|
||||
/// use tokio::net::UnixStream;
|
||||
/// use std::error::Error;
|
||||
/// use std::io;
|
||||
///
|
||||
/// #[tokio::main]
|
||||
/// async fn main() -> Result<(), Box<dyn Error>> {
|
||||
/// // Connect to a peer
|
||||
/// let dir = tempfile::tempdir().unwrap();
|
||||
/// let bind_path = dir.path().join("bind_path");
|
||||
/// let stream = UnixStream::connect(bind_path).await?;
|
||||
///
|
||||
/// loop {
|
||||
/// // Wait for the socket to be writable
|
||||
/// stream.writable().await?;
|
||||
///
|
||||
/// // Try to write data, this may still fail with `WouldBlock`
|
||||
/// // if the readiness event is a false positive.
|
||||
/// match stream.try_write(b"hello world") {
|
||||
/// Ok(n) => {
|
||||
/// break;
|
||||
/// }
|
||||
/// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
|
||||
/// continue;
|
||||
/// }
|
||||
/// Err(e) => {
|
||||
/// return Err(e.into());
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
///
|
||||
/// Ok(())
|
||||
/// }
|
||||
/// ```
|
||||
pub fn try_write(&self, buf: &[u8]) -> io::Result<usize> {
|
||||
self.io
|
||||
.registration()
|
||||
.try_io(Interest::WRITABLE, || (&*self.io).write(buf))
|
||||
}
|
||||
|
||||
/// Creates new `UnixStream` from a `std::os::unix::net::UnixStream`.
|
||||
///
|
||||
/// This function is intended to be used to wrap a UnixStream from the
|
||||
|
||||
+199
-3
@@ -2,10 +2,14 @@
|
||||
#![warn(rust_2018_idioms)]
|
||||
#![cfg(unix)]
|
||||
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::{UnixListener, UnixStream};
|
||||
use std::io;
|
||||
use std::task::Poll;
|
||||
|
||||
use futures::future::try_join;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt, Interest};
|
||||
use tokio::net::{UnixListener, UnixStream};
|
||||
use tokio_test::{assert_ok, assert_pending, assert_ready_ok, task};
|
||||
|
||||
use futures::future::{poll_fn, try_join};
|
||||
|
||||
#[tokio::test]
|
||||
async fn accept_read_write() -> std::io::Result<()> {
|
||||
@@ -56,3 +60,195 @@ async fn shutdown() -> std::io::Result<()> {
|
||||
assert_eq!(n, 0);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn try_read_write() -> std::io::Result<()> {
|
||||
let msg = b"hello world";
|
||||
|
||||
let dir = tempfile::tempdir()?;
|
||||
let bind_path = dir.path().join("bind.sock");
|
||||
|
||||
// Create listener
|
||||
let listener = UnixListener::bind(&bind_path)?;
|
||||
|
||||
// Create socket pair
|
||||
let client = UnixStream::connect(&bind_path).await?;
|
||||
|
||||
let (server, _) = listener.accept().await?;
|
||||
let mut written = msg.to_vec();
|
||||
|
||||
// Track the server receiving data
|
||||
let mut readable = task::spawn(server.readable());
|
||||
assert_pending!(readable.poll());
|
||||
|
||||
// Write data.
|
||||
client.writable().await?;
|
||||
assert_eq!(msg.len(), client.try_write(msg)?);
|
||||
|
||||
// The task should be notified
|
||||
while !readable.is_woken() {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
|
||||
// Fill the write buffer
|
||||
loop {
|
||||
// Still ready
|
||||
let mut writable = task::spawn(client.writable());
|
||||
assert_ready_ok!(writable.poll());
|
||||
|
||||
match client.try_write(msg) {
|
||||
Ok(n) => written.extend(&msg[..n]),
|
||||
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
|
||||
break;
|
||||
}
|
||||
Err(e) => panic!("error = {:?}", e),
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
// Write buffer full
|
||||
let mut writable = task::spawn(client.writable());
|
||||
assert_pending!(writable.poll());
|
||||
|
||||
// Drain the socket from the server end
|
||||
let mut read = vec![0; written.len()];
|
||||
let mut i = 0;
|
||||
|
||||
while i < read.len() {
|
||||
server.readable().await?;
|
||||
|
||||
match server.try_read(&mut read[i..]) {
|
||||
Ok(n) => i += n,
|
||||
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
|
||||
Err(e) => panic!("error = {:?}", e),
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(read, written);
|
||||
}
|
||||
|
||||
// Now, we listen for shutdown
|
||||
drop(client);
|
||||
|
||||
loop {
|
||||
let ready = server.ready(Interest::READABLE).await?;
|
||||
|
||||
if ready.is_read_closed() {
|
||||
break;
|
||||
} else {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_pair() -> (UnixStream, UnixStream) {
|
||||
let dir = assert_ok!(tempfile::tempdir());
|
||||
let bind_path = dir.path().join("bind.sock");
|
||||
|
||||
let listener = assert_ok!(UnixListener::bind(&bind_path));
|
||||
|
||||
let accept = listener.accept();
|
||||
let connect = UnixStream::connect(&bind_path);
|
||||
let ((server, _), client) = assert_ok!(try_join(accept, connect).await);
|
||||
|
||||
(client, server)
|
||||
}
|
||||
|
||||
macro_rules! assert_readable_by_polling {
|
||||
($stream:expr) => {
|
||||
assert_ok!(poll_fn(|cx| $stream.poll_read_ready(cx)).await);
|
||||
};
|
||||
}
|
||||
|
||||
macro_rules! assert_not_readable_by_polling {
|
||||
($stream:expr) => {
|
||||
poll_fn(|cx| {
|
||||
assert_pending!($stream.poll_read_ready(cx));
|
||||
Poll::Ready(())
|
||||
})
|
||||
.await;
|
||||
};
|
||||
}
|
||||
|
||||
macro_rules! assert_writable_by_polling {
|
||||
($stream:expr) => {
|
||||
assert_ok!(poll_fn(|cx| $stream.poll_write_ready(cx)).await);
|
||||
};
|
||||
}
|
||||
|
||||
macro_rules! assert_not_writable_by_polling {
|
||||
($stream:expr) => {
|
||||
poll_fn(|cx| {
|
||||
assert_pending!($stream.poll_write_ready(cx));
|
||||
Poll::Ready(())
|
||||
})
|
||||
.await;
|
||||
};
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn poll_read_ready() {
|
||||
let (mut client, mut server) = create_pair().await;
|
||||
|
||||
// Initial state - not readable.
|
||||
assert_not_readable_by_polling!(server);
|
||||
|
||||
// There is data in the buffer - readable.
|
||||
assert_ok!(client.write_all(b"ping").await);
|
||||
assert_readable_by_polling!(server);
|
||||
|
||||
// Readable until calls to `poll_read` return `Poll::Pending`.
|
||||
let mut buf = [0u8; 4];
|
||||
assert_ok!(server.read_exact(&mut buf).await);
|
||||
assert_readable_by_polling!(server);
|
||||
read_until_pending(&mut server);
|
||||
assert_not_readable_by_polling!(server);
|
||||
|
||||
// Detect the client disconnect.
|
||||
drop(client);
|
||||
assert_readable_by_polling!(server);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn poll_write_ready() {
|
||||
let (mut client, server) = create_pair().await;
|
||||
|
||||
// Initial state - writable.
|
||||
assert_writable_by_polling!(client);
|
||||
|
||||
// No space to write - not writable.
|
||||
write_until_pending(&mut client);
|
||||
assert_not_writable_by_polling!(client);
|
||||
|
||||
// Detect the server disconnect.
|
||||
drop(server);
|
||||
assert_writable_by_polling!(client);
|
||||
}
|
||||
|
||||
fn read_until_pending(stream: &mut UnixStream) {
|
||||
let mut buf = vec![0u8; 1024 * 1024];
|
||||
loop {
|
||||
match stream.try_read(&mut buf) {
|
||||
Ok(_) => (),
|
||||
Err(err) => {
|
||||
assert_eq!(err.kind(), io::ErrorKind::WouldBlock);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn write_until_pending(stream: &mut UnixStream) {
|
||||
let buf = vec![0u8; 1024 * 1024];
|
||||
loop {
|
||||
match stream.try_write(&buf) {
|
||||
Ok(_) => (),
|
||||
Err(err) => {
|
||||
assert_eq!(err.kind(), io::ErrorKind::WouldBlock);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user