io: retry ErrorKind::Interrupted in read_exact (#8417)

Co-authored-by: philphauler <[email protected]>

tokio/src/io/util/read_exact.rs

@@ -57,7 +57,11 @@ where
             // if our buffer is empty, then we need to read some data to continue.
             let rem = me.buf.remaining();
             if rem != 0 {
-                ready!(Pin::new(&mut *me.reader).poll_read(cx, me.buf))?;
+                match ready!(Pin::new(&mut *me.reader).poll_read(cx, me.buf)) {
+                    Ok(()) => {}
+                    Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
+                    Err(e) => return Err(e).into(),
+                }
                 if me.buf.remaining() == rem {
                     return Err(eof()).into();
                 }

tokio/tests/io_read_exact.rs

@@ -9,7 +9,10 @@
     )
 ))]
 
-use tokio::io::AsyncReadExt;
+use std::io;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use tokio::io::{AsyncRead, AsyncReadExt, ReadBuf};
 use tokio_test::assert_ok;
 
 #[tokio::test]
@@ -21,3 +24,37 @@ async fn read_exact() {
     assert_eq!(n, 8);
     assert_eq!(buf[..], b"hello wo"[..]);
 }
+
+struct InterruptThenRead {
+    interrupted: bool,
+    data: &'static [u8],
+}
+
+impl AsyncRead for InterruptThenRead {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        if !self.interrupted {
+            self.interrupted = true;
+            return Poll::Ready(Err(io::Error::from(io::ErrorKind::Interrupted)));
+        }
+        let n = std::cmp::min(self.data.len(), buf.remaining());
+        buf.put_slice(&self.data[..n]);
+        self.data = &self.data[n..];
+        Poll::Ready(Ok(()))
+    }
+}
+
+#[tokio::test]
+async fn read_exact_retries_interrupted() {
+    let mut reader = InterruptThenRead {
+        interrupted: false,
+        data: b"hello",
+    };
+    let mut buf = [0u8; 5];
+    let n = reader.read_exact(&mut buf).await.unwrap();
+    assert_eq!(n, 5);
+    assert_eq!(&buf, b"hello");
+}