Skip to content

Commit

Permalink
io: add Join utility
Browse files Browse the repository at this point in the history
  • Loading branch information
devsnek committed Dec 16, 2023
1 parent 4aa7bbf commit f784ba4
Show file tree
Hide file tree
Showing 3 changed files with 218 additions and 0 deletions.
133 changes: 133 additions & 0 deletions tokio/src/io/join.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,133 @@
//! Join two values implementing `AsyncRead` and `AsyncWrite` into a single one.

use crate::io::{AsyncRead, AsyncWrite, ReadBuf};

use std::fmt;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};

cfg_io_util! {
pin_project_lite::pin_project! {
/// Joins two values implementing `AsyncRead` and `AsyncWrite` into a
/// single handle.
pub struct Join<R, W> {
#[pin]
reader: R,
#[pin]
writer: W,
}
}

impl<R, W> Join<R, W>
where
R: AsyncRead + Unpin,
W: AsyncWrite + Unpin,
{
/// Join two values implementing `AsyncRead` and `AsyncWrite` into a
/// single handle.
pub fn new(reader: R, writer: W) -> Self {
Self { reader, writer }
}

/// Splits this `Join` back into its `AsyncRead` and `AsyncWrite`
/// components.
pub fn split(self) -> (R, W) {
(self.reader, self.writer)
}

/// Returns a reference to the inner reader.
pub fn reader(&self) -> &R {
&self.reader
}

/// Returns a reference to the inner writer.
pub fn writer(&self) -> &W {
&self.writer
}

/// Returns a mutable reference to the inner reader.
pub fn reader_mut(&mut self) -> &mut R {
&mut self.reader
}

/// Returns a mutable reference to the inner writer.
pub fn writer_mut(&mut self) -> &mut W {
&mut self.writer
}
}

impl<R, W> AsyncRead for Join<R, W>
where
R: AsyncRead + Unpin,
{
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<Result<(), io::Error>> {
self.project().reader.poll_read(cx, buf)
}
}

impl<R, W> AsyncWrite for Join<R, W>
where
W: AsyncWrite + Unpin,
{
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, io::Error>> {
self.project().writer.poll_write(cx, buf)
}

fn poll_flush(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<(), io::Error>> {
self.project().writer.poll_flush(cx)
}

fn poll_shutdown(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<(), io::Error>> {
self.project().writer.poll_shutdown(cx)
}

fn poll_write_vectored(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
bufs: &[io::IoSlice<'_>]
) -> Poll<Result<usize, io::Error>> {
self.project().writer.poll_write_vectored(cx, bufs)
}

fn is_write_vectored(&self) -> bool {
self.writer.is_write_vectored()
}
}

impl<R, W> Clone for Join<R, W>
where R: Clone, W: Clone
{
fn clone(&self) -> Self {
Join {
writer: self.writer.clone(),
reader: self.reader.clone(),
}
}
}

impl<R, W> fmt::Debug for Join<R, W>
where R: fmt::Debug, W: fmt::Debug
{
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Join")
.field("writer", &self.writer)
.field("reader", &self.reader)
.finish()
}
}
}
2 changes: 2 additions & 0 deletions tokio/src/io/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -265,6 +265,8 @@ cfg_io_std! {
cfg_io_util! {
mod split;
pub use split::{split, ReadHalf, WriteHalf};
mod join;
pub use join::Join;

pub(crate) mod seek;
pub(crate) mod util;
Expand Down
83 changes: 83 additions & 0 deletions tokio/tests/io_join.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
#![warn(rust_2018_idioms)]
#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support panic recovery

use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, Join, ReadBuf};

use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};

struct R;

impl AsyncRead for R {
fn poll_read(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
buf.put_slice(&[b'z']);
Poll::Ready(Ok(()))
}
}

struct W;

impl AsyncWrite for W {
fn poll_write(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
_buf: &[u8],
) -> Poll<Result<usize, io::Error>> {
Poll::Ready(Ok(1))
}

fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
Poll::Ready(Ok(()))
}

fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
Poll::Ready(Ok(()))
}

fn poll_write_vectored(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
_bufs: &[io::IoSlice<'_>],
) -> Poll<Result<usize, io::Error>> {
Poll::Ready(Ok(2))
}

fn is_write_vectored(&self) -> bool {
true
}
}

#[test]
fn is_send_and_sync() {
fn assert_bound<T: Send + Sync>() {}

assert_bound::<Join<W, R>>();
}

#[test]
fn method_delegation() {
let mut rw = Join::new(R, W);
let mut buf = [0; 1];

tokio_test::block_on(async move {
assert_eq!(1, rw.read(&mut buf).await.unwrap());
assert_eq!(b'z', buf[0]);

assert_eq!(1, rw.write(&[b'x']).await.unwrap());
assert_eq!(
2,
rw.write_vectored(&[io::IoSlice::new(&[b'x'])])
.await
.unwrap()
);
assert!(rw.is_write_vectored());

assert!(rw.flush().await.is_ok());
assert!(rw.shutdown().await.is_ok());
});
}

0 comments on commit f784ba4

Please sign in to comment.