Skip to content

Commit

Permalink
io: make duplex stream cooperative (tokio-rs#4470)
Browse files Browse the repository at this point in the history
Add coop checks on pipe poll_read and poll_write.

Fixes: tokio-rs#4470
Refs: tokio-rs#4291, tokio-rs#4300
  • Loading branch information
GongLG committed Feb 7, 2022
1 parent fc4deaa commit 0b9acc6
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 0 deletions.
16 changes: 16 additions & 0 deletions tokio/src/io/util/mem.rs
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,7 @@ impl AsyncRead for Pipe {
cx: &mut task::Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
ready!(poll_proceed_and_make_progress(cx));
if self.buffer.has_remaining() {
let max = self.buffer.remaining().min(buf.remaining());
buf.put_slice(&self.buffer[..max]);
Expand Down Expand Up @@ -212,6 +213,7 @@ impl AsyncWrite for Pipe {
cx: &mut task::Context<'_>,
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
ready!(poll_proceed_and_make_progress(cx));
if self.is_closed {
return Poll::Ready(Err(std::io::ErrorKind::BrokenPipe.into()));
}
Expand Down Expand Up @@ -241,3 +243,17 @@ impl AsyncWrite for Pipe {
Poll::Ready(Ok(()))
}
}

cfg_coop! {
fn poll_proceed_and_make_progress(cx: &mut task::Context<'_>) -> Poll<()> {
let coop = ready!(crate::coop::poll_proceed(cx));
coop.made_progress();
Poll::Ready(())
}
}

cfg_not_coop! {
fn poll_proceed_and_make_progress(_: &mut Context<'_>) -> Poll<()> {
Poll::Ready(())
}
}
19 changes: 19 additions & 0 deletions tokio/tests/io_mem_stream.rs
Original file line number Diff line number Diff line change
Expand Up @@ -100,3 +100,22 @@ async fn max_write_size() {
// drop b only after task t1 finishes writing
drop(b);
}

#[tokio::test]
async fn duplex_is_cooperative() {
let (mut tx, mut rx) = tokio::io::duplex(1024 * 8);

tokio::select! {
biased;

_ = async {
loop {
let buf = [3u8; 4096];
let _ = tx.write_all(&buf).await;
let mut buf = [0u8; 4096];
let _ = rx.read(&mut buf).await;
}
} => {},
_ = tokio::task::yield_now() => {}
}
}

0 comments on commit 0b9acc6

Please sign in to comment.