wake up the receiver on shutdown

Signed-off-by: ADD-SP <qiqi.zhang@konghq.com>

Author: ADD-SP

tokio-util/src/io/simplex.rs

@@ -217,8 +217,15 @@ impl AsyncWrite for Sender {
     /// The [`Receiver`] can still be used to read remaining data
     /// until all bytes have been consumed.
     fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<IoResult<()>> {
-        let mut inner = self.inner.lock().unwrap();
-        inner.close_sender();
+        let maybe_waker = {
+            let mut inner = self.inner.lock().unwrap();
+            inner.close_sender()
+        };
+
+        if let Some(waker) = maybe_waker {
+            waker.wake();
+        }
+
         Poll::Ready(Ok(()))
     }
 }

tokio-util/tests/io_simplex.rs

@@ -3,6 +3,7 @@ use futures_test::task::noop_context;
 use std::{future::Future, task::Poll};
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
 use tokio_test::{assert_pending, assert_ready_err};
+use tokio_test::task::spawn;
 use tokio_util::io::simplex;
 
 /// Sanity check for single-threaded operation.
@@ -139,6 +140,26 @@ async fn shutdown_sender_1() {
     tx.shutdown().await.unwrap();
 }
 
+/// The `Sender::poll_shutdown` should wake up the `Receiver`
+#[tokio::test]
+async fn shutdown_sender_2() {
+    let (mut tx, mut rx) = simplex::new(32);
+
+    let mut buf = vec![];
+    let mut read_task = spawn(rx.read_to_end(&mut buf));
+    assert_pending!(read_task.poll());
+
+    tx.write_u8(1).await.unwrap();
+    assert_pending!(read_task.poll());
+
+    assert!(!read_task.is_woken());
+    tx.shutdown().await.unwrap();
+    assert!(read_task.is_woken());
+
+    read_task.await.unwrap();
+    assert_eq!(buf, vec![1]);
+}
+
 /// Both `Sender` and `Receiver` should yield periodically
 /// in a tight-loop.
 #[tokio::test]