io: fix the hang-forever issue of Mock when the write buf is empty

Signed-off-by: ADD-SP <[email protected]>

Author: ADD-SP

tokio-test/Cargo.toml

@@ -24,6 +24,7 @@ futures-core = "0.3.0"
 [dev-dependencies]
 tokio = { version = "1.2.0", path = "../tokio", features = ["full"] }
 futures-util = "0.3.0"
+futures-test = "0.3.5"
 
 [package.metadata.docs.rs]
 all-features = true

tokio-test/src/io.rs

@@ -23,6 +23,7 @@ use tokio::sync::mpsc;
 use tokio::time::{self, Duration, Instant, Sleep};
 use tokio_stream::wrappers::UnboundedReceiverStream;
 
+use core::panic;
 use futures_core::Stream;
 use std::collections::VecDeque;
 use std::fmt;
@@ -235,7 +236,9 @@ impl Inner {
                 let err = Arc::try_unwrap(err).expect("There are no other references.");
                 Err(err)
             }
-            Some(_) => {
+            Some(&mut Action::Write(_))
+            | Some(&mut Action::WriteError(_))
+            | Some(&mut Action::Wait(_)) => {
                 // Either waiting or expecting a write
                 Err(io::ErrorKind::WouldBlock.into())
             }
@@ -280,10 +283,8 @@ impl Inner {
                 Action::Wait(..) | Action::WriteError(..) => {
                     break;
                 }
-                _ => {}
+                Action::Read(_) | Action::ReadError(_) => (),
             }
-
-            // TODO: remove write
         }
 
         Ok(ret)
@@ -441,6 +442,7 @@ impl AsyncWrite for Mock {
                         panic!("unexpected WouldBlock {}", self.pmsg());
                     }
                 }
+                Ok(0) if buf.is_empty() => return Poll::Ready(Ok(0)),
                 Ok(0) => {
                     // TODO: Is this correct?
                     if !self.inner.actions.is_empty() {
@@ -494,7 +496,7 @@ impl Drop for Mock {
                 "There is still data left to write. {}",
                 self.pmsg()
             ),
-            _ => (),
+            Action::ReadError(_) | Action::WriteError(_) | Action::Wait(_) => (),
         });
     }
 }

tokio-test/tests/io.rs

@@ -1,8 +1,14 @@
 #![warn(rust_2018_idioms)]
 
+use futures_test::task::{noop_context, panic_waker};
+use futures_util::pin_mut;
+use std::future::Future;
 use std::io;
+use std::task::{Context, Poll};
+use tokio::io::AsyncWrite;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::time::{Duration, Instant};
+use tokio_test::assert_pending;
 use tokio_test::io::Builder;
 
 #[tokio::test]
@@ -170,3 +176,84 @@ async fn multiple_wait() {
         start.elapsed().as_millis()
     );
 }
+
+// No matter which usecase, it doesn't make sense for a read
+// hang forever. However, currently, if there is no sequenced read
+// action, it will hang forever.
+//
+// Since we want be aware of the fixing of this bug,
+// no matter intentionally or unintentionally,
+// we add this test to catch the behavior change.
+//
+// It looks like fixing it is not hard, but not sure the downstream
+// impact, which might be a breaking change due to the
+// `Mock::inner::read_wait` field, so we keep it as is for now.
+//
+// TODO: fix this bug
+#[test]
+fn should_hang_forever_on_read_but_no_sequenced_read_action() {
+    let mut mock = Builder::new()
+        .write_error(io::Error::new(io::ErrorKind::Other, "cruel"))
+        .build();
+
+    let mut buf = [0; 1];
+    let read_exact_fut = mock.read(&mut buf);
+    pin_mut!(read_exact_fut);
+    assert_pending!(read_exact_fut.poll(&mut Context::from_waker(&panic_waker())));
+}
+
+// The `Mock` is expected to always panic if there is an unconsumed error action,
+// rather than silently ignoring it. However,
+// currently it only panics on unconsumed read/write actions,
+// not on error actions. Fixing this requires a breaking change.
+//
+// This test verifies that it does not panic yet,
+// to prevent accidentally introducing the breaking change prematurely.
+//
+// TODO: fix this bug in the next major release
+#[test]
+fn do_not_panic_unconsumed_error() {
+    let _mock = Builder::new()
+        .read_error(io::Error::new(io::ErrorKind::Other, "cruel"))
+        .build();
+}
+
+// The `Mock` must never panic, even if cloned multiple times.
+// However, at present, cloning the builder under certain
+// conditions causes a panic.
+//
+// Fixing this would require making `Mock` non-`Clone`,
+// which is a breaking change.
+//
+// Since we want be aware of the fixing of this bug,
+// no matter intentionally or unintentionally,
+// we add this test to catch the behavior change.
+//
+// TODO: fix this bug in the next major release
+#[tokio::test]
+#[should_panic = "There are no other references.: Custom { kind: Other, error: \"cruel\" }"]
+async fn panic_if_clone_the_build_with_error_action() {
+    let mut builder = Builder::new();
+    builder.write_error(io::Error::new(io::ErrorKind::Other, "cruel"));
+    let mut builder2 = builder.clone();
+
+    let mut mock = builder.build();
+    let _mock2 = builder2.build();
+
+    // this write_all will panic due to unwrapping the error from `Arc`
+    mock.write_all(b"hello").await.unwrap();
+    unreachable!();
+}
+
+#[tokio::test]
+async fn should_not_hang_forever_on_zero_length_write() {
+    let mock = Builder::new().write(b"write").build();
+    pin_mut!(mock);
+    match mock.as_mut().poll_write(&mut noop_context(), &[0u8; 0]) {
+        // drain the remaining write action to avoid panic at drop of the `mock`
+        Poll::Ready(Ok(0)) => mock.write_all(b"write").await.unwrap(),
+        Poll::Ready(Ok(n)) => panic!("expected to write 0 bytes, wrote {n} bytes instead"),
+        Poll::Ready(Err(e)) => panic!("expected to write 0 bytes, got error {e} instead"),
+        Poll::Pending => panic!("expected to write 0 bytes immediately, but pending instead"),
+    }
+}