feat(transport): introduce stream I/O timeouts

crates/quic/src/lib.rs

@@ -18,6 +18,7 @@ pub struct Client(Connection);
 pub struct ConnHandler;
 
 const DONE: VarInt = VarInt::from_u32(1);
+const FAILED: VarInt = VarInt::from_u32(2);
 
 impl wrpc_transport::frame::ConnHandler<RecvStream, SendStream> for ConnHandler {
     async fn on_ingress(mut rx: RecvStream, res: std::io::Result<()>) {
@@ -31,12 +32,15 @@ impl wrpc_transport::frame::ConnHandler<RecvStream, SendStream> for ConnHandler
         }
     }
 
-    async fn on_egress(tx: SendStream, res: std::io::Result<()>) {
+    async fn on_egress(mut tx: SendStream, res: std::io::Result<()>) {
         if let Err(err) = res {
             error!(?err, "egress failed");
-        } else {
-            debug!("egress successfully complete");
+            if let Err(err) = tx.reset(FAILED) {
+                debug!(?err, "failed to reset outgoing stream");
+            }
+            return;
         }
+        debug!("egress successfully complete");
         match tx.stopped().await {
             Ok(None) => {
                 trace!("stream successfully closed");

crates/transport/src/frame/conn/mod.rs

@@ -1,16 +1,18 @@
 use core::fmt::{Debug, Display};
 use core::mem;
 use core::pin::Pin;
-use core::task::{Context, Poll, ready};
+use core::task::{Context, Poll, Waker, ready};
+use core::time::Duration;
 
-use std::sync::Arc;
+use std::sync::{Arc, PoisonError};
 
 use bytes::{Buf as _, Bytes, BytesMut};
 use futures::Sink as _;
 use pin_project_lite::pin_project;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt as _};
 use tokio::sync::mpsc;
 use tokio::task::JoinSet;
+use tokio::time::Sleep;
 use tokio_stream::wrappers::ReceiverStream;
 use tokio_util::codec::Encoder;
 use tokio_util::io::StreamReader;
@@ -355,6 +357,8 @@ pin_project! {
         index: Arc<std::sync::Mutex<IndexTrie>>,
         io: Arc<JoinSet<()>>,
         err: Arc<std::sync::OnceLock<std::io::Error>>,
+        timeout: Option<Duration>,
+        deadline: Option<Pin<Box<Sleep>>>,
     }
 }
 
@@ -401,9 +405,27 @@ impl Incoming {
             index: Arc::clone(&index),
             io: Arc::new(rx_io),
             err,
+            timeout: None,
+            deadline: None,
         }
     }
 
+    /// Sets the read timeout.
+    ///
+    /// Once set, a read that stays pending for longer than `timeout` fails with
+    /// [`std::io::ErrorKind::TimedOut`]. The deadline is armed when a read first returns
+    /// [`Poll::Pending`] and cleared once a read completes or times out. Dropping a pending
+    /// read does not clear it, so the next read fails as soon as the deadline elapses unless
+    /// data arrives first. The timeout is inherited by sub-streams returned by [`Self::index`]
+    /// after this call, including those backing async `stream` and `future` values, which must
+    /// therefore not idle for longer than `timeout`; sub-streams indexed before this call are
+    /// unaffected.
+    #[must_use]
+    pub fn with_timeout(mut self, timeout: Duration) -> Self {
+        self.timeout = Some(timeout);
+        self
+    }
+
     /// Index the incoming stream using a structural `path`, returning a handle to the
     /// multiplexed sub-stream addressed by it.
     #[instrument(level = "trace", skip(self), fields(path = ?self.path))]
@@ -434,6 +456,8 @@ impl Incoming {
             index: Arc::clone(&self.index),
             io: Arc::clone(&self.io),
             err: Arc::clone(&self.err),
+            timeout: self.timeout,
+            deadline: None,
         })
     }
 }
@@ -454,7 +478,24 @@ impl AsyncRead for Incoming {
             trace!("reader is closed");
             return Poll::Ready(Ok(()));
         };
-        ready!(rx.poll_read(cx, buf))?;
+        if let Poll::Ready(res) = rx.poll_read(cx, buf) {
+            *this.deadline = None;
+            res?;
+        } else {
+            let Some(timeout) = *this.timeout else {
+                return Poll::Pending;
+            };
+            let sleep = this
+                .deadline
+                .get_or_insert_with(|| Box::pin(tokio::time::sleep(timeout)));
+            ready!(sleep.as_mut().poll(cx));
+            *this.deadline = None;
+            trace!("read timed out");
+            return Poll::Ready(Err(std::io::Error::new(
+                std::io::ErrorKind::TimedOut,
+                "read timed out",
+            )));
+        }
         trace!(buf = ?buf.filled(), "read buffer");
         if buf.filled().is_empty() {
             self.rx.take();
@@ -474,9 +515,17 @@ pin_project! {
         tx: PollSender<(Bytes, Bytes)>,
         path: Arc<[usize]>,
         path_buf: Bytes,
+        timeout: Arc<std::sync::Mutex<Timeout>>,
+        err: Arc<std::sync::OnceLock<std::io::Error>>,
     }
 }
 
+#[derive(Default)]
+struct Timeout {
+    duration: Option<Duration>,
+    waker: Option<Waker>,
+}
+
 impl Outgoing {
     /// Creates a new [Outgoing] given an [`AsyncWrite`].
     pub fn new<T, Fut>(
@@ -488,21 +537,63 @@ impl Outgoing {
         Fut: Future<Output = ()> + Send,
     {
         let span = Span::current();
-        let (tx_tx, tx_rx) = mpsc::channel(128);
-        tokio::spawn(
-            async {
-                let res = egress(&mut tx, tx_rx).await;
+        let (tx_tx, mut tx_rx) = mpsc::channel(128);
+        let timeout = Arc::new(std::sync::Mutex::default());
+        let err = Arc::new(std::sync::OnceLock::default());
+        tokio::spawn({
+            let timeout = Arc::clone(&timeout);
+            let err = Arc::clone(&err);
+            async move {
+                let res = egress(
+                    TimeoutWriter {
+                        inner: &mut tx,
+                        timeout,
+                        duration: None,
+                        deadline: None,
+                    },
+                    &mut tx_rx,
+                )
+                .await;
+                if let Err(e) = &res {
+                    _ = err.set(copy_io_error(e));
+                }
+                drop(tx_rx);
                 on_egress(tx, res).await;
             }
-            .instrument(span.clone()),
-        );
+            .instrument(span.clone())
+        });
         Self {
             tx: PollSender::new(tx_tx),
             path: Arc::from([]),
             path_buf: Bytes::from_static(&[0]),
+            timeout,
+            err,
         }
     }
 
+    /// Sets the write timeout.
+    ///
+    /// Writes are buffered and written to the underlying connection by a background task.
+    /// Once set, that task fails with [`std::io::ErrorKind::TimedOut`] if the connection does
+    /// not accept any buffered data within `timeout`, after which the result is passed to
+    /// `on_egress` and all subsequent writes, flushes and shutdowns on this stream fail with
+    /// the same error. Flushes and shutdowns do not wait for the background task, so they only
+    /// report a timeout that has already occurred and data buffered before it is lost. Other
+    /// background write errors are only reported by subsequent writes. The timeout is shared
+    /// by all sub-streams of this connection, including those returned by [`Self::index`].
+    #[must_use]
+    pub fn with_timeout(self, timeout: Duration) -> Self {
+        let waker = {
+            let mut t = self.timeout.lock().unwrap_or_else(PoisonError::into_inner);
+            t.duration = Some(timeout);
+            t.waker.take()
+        };
+        if let Some(waker) = waker {
+            waker.wake();
+        }
+        self
+    }
+
     /// Index the outgoing stream using a structural `path`, returning a handle that writes
     /// to the multiplexed sub-stream addressed by it.
     #[instrument(level = "trace", skip(self), fields(path = ?self.path))]
@@ -533,6 +624,8 @@ impl Outgoing {
             tx: self.tx.clone(),
             path,
             path_buf: buf.freeze(),
+            timeout: Arc::clone(&self.timeout),
+            err: Arc::clone(&self.err),
         })
     }
 }
@@ -546,22 +639,114 @@ impl AsyncWrite for Outgoing {
     ) -> Poll<std::io::Result<usize>> {
         trace!("writing outgoing chunk");
         let mut this = self.project();
-        ready!(this.tx.as_mut().poll_ready(cx))
-            .map_err(|err| std::io::Error::new(std::io::ErrorKind::BrokenPipe, err))?;
+        ready!(this.tx.as_mut().poll_ready(cx)).map_err(|e| egress_error(this.err, e))?;
         this.tx
             .start_send((this.path_buf.clone(), Bytes::copy_from_slice(buf)))
-            .map_err(|err| std::io::Error::new(std::io::ErrorKind::BrokenPipe, err))?;
+            .map_err(|e| egress_error(this.err, e))?;
         Poll::Ready(Ok(buf.len()))
     }
 
     #[instrument(level = "trace", skip_all, fields(path = ?self.path), ret(level = "trace"))]
     fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        Poll::Ready(Ok(()))
+        Poll::Ready(timeout_error(&self.err).map_or(Ok(()), Err))
     }
 
     #[instrument(level = "trace", skip_all, fields(path = ?self.path), ret(level = "trace"))]
     fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        Poll::Ready(Ok(()))
+        Poll::Ready(timeout_error(&self.err).map_or(Ok(()), Err))
+    }
+}
+
+fn timeout_error(err: &std::sync::OnceLock<std::io::Error>) -> Option<std::io::Error> {
+    err.get()
+        .filter(|err| err.kind() == std::io::ErrorKind::TimedOut)
+        .map(copy_io_error)
+}
+
+fn egress_error(
+    err: &std::sync::OnceLock<std::io::Error>,
+    e: impl Into<Box<dyn core::error::Error + Send + Sync>>,
+) -> std::io::Error {
+    err.get().map_or_else(
+        || std::io::Error::new(std::io::ErrorKind::BrokenPipe, e),
+        copy_io_error,
+    )
+}
+
+pin_project! {
+    struct TimeoutWriter<T> {
+        #[pin]
+        inner: T,
+        timeout: Arc<std::sync::Mutex<Timeout>>,
+        duration: Option<Duration>,
+        deadline: Option<Pin<Box<Sleep>>>,
+    }
+}
+
+impl<T: AsyncWrite> TimeoutWriter<T> {
+    fn poll_timeout<R>(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        poll: impl FnOnce(Pin<&mut T>, &mut Context<'_>) -> Poll<std::io::Result<R>>,
+    ) -> Poll<std::io::Result<R>> {
+        let this = self.project();
+        if let Poll::Ready(res) = poll(this.inner, cx) {
+            *this.deadline = None;
+            return Poll::Ready(res);
+        }
+        let duration = {
+            let mut t = this.timeout.lock().unwrap_or_else(PoisonError::into_inner);
+            if !t.waker.as_ref().is_some_and(|w| w.will_wake(cx.waker())) {
+                t.waker = Some(cx.waker().clone());
+            }
+            t.duration
+        };
+        if duration != *this.duration {
+            *this.duration = duration;
+            *this.deadline = None;
+        }
+        let Some(duration) = duration else {
+            return Poll::Pending;
+        };
+        let sleep = this
+            .deadline
+            .get_or_insert_with(|| Box::pin(tokio::time::sleep(duration)));
+        ready!(sleep.as_mut().poll(cx));
+        trace!("connection write timed out");
+        Poll::Ready(Err(std::io::Error::new(
+            std::io::ErrorKind::TimedOut,
+            "connection write timed out",
+        )))
+    }
+}
+
+impl<T: AsyncWrite> AsyncWrite for TimeoutWriter<T> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<std::io::Result<usize>> {
+        self.poll_timeout(cx, |inner, cx| inner.poll_write(cx, buf))
+    }
+
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[std::io::IoSlice<'_>],
+    ) -> Poll<std::io::Result<usize>> {
+        self.poll_timeout(cx, |inner, cx| inner.poll_write_vectored(cx, bufs))
+    }
+
+    fn is_write_vectored(&self) -> bool {
+        self.inner.is_write_vectored()
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        self.poll_timeout(cx, AsyncWrite::poll_flush)
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        self.poll_timeout(cx, AsyncWrite::poll_shutdown)
     }
 }
 
@@ -638,7 +823,7 @@ async fn ingress(
 #[instrument(level = "trace", skip_all)]
 async fn egress(
     mut tx: impl AsyncWrite + Unpin,
-    mut rx: mpsc::Receiver<(Bytes, Bytes)>,
+    rx: &mut mpsc::Receiver<(Bytes, Bytes)>,
 ) -> std::io::Result<()> {
     let mut buf = BytesMut::with_capacity(5);
     trace!("waiting for next frame");
@@ -700,6 +885,52 @@ mod tests {
         assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
     }
 
+    #[test_log::test(tokio::test)]
+    async fn incoming_read_after_timeout() {
+        let (mut tx, rx) = tokio::io::duplex(16);
+        let mut rx =
+            Incoming::new(rx, [[]; 0], |_, _| async {}).with_timeout(Duration::from_millis(50));
+        let mut buf = [0; 1];
+        let err = rx.read_exact(&mut buf).await.unwrap_err();
+        assert_eq!(err.kind(), std::io::ErrorKind::TimedOut);
+        tokio::spawn(async move {
+            tokio::time::sleep(Duration::from_millis(10)).await;
+            tx.write_all(&[0x00, 0x01, 0x2a]).await.unwrap();
+            tx
+        });
+        rx.read_exact(&mut buf).await.unwrap();
+        assert_eq!(buf, [0x2a]);
+    }
+
+    #[test_log::test(tokio::test)]
+    async fn outgoing_timeout_after_stall() {
+        let (tx, _rx) = tokio::io::duplex(1);
+        let (res_tx, res_rx) = tokio::sync::oneshot::channel();
+        let mut out = Outgoing::new(tx, |_, res| async {
+            _ = res_tx.send(res);
+        });
+        out.write_all(&[0; 16]).await.unwrap();
+        tokio::time::sleep(Duration::from_millis(10)).await;
+        let mut out = out.with_timeout(Duration::from_millis(10));
+        let err = tokio::time::timeout(Duration::from_secs(5), res_rx)
+            .await
+            .unwrap()
+            .unwrap()
+            .unwrap_err();
+        assert_eq!(err.kind(), std::io::ErrorKind::TimedOut);
+        let err = out.write_all(&[0]).await.unwrap_err();
+        assert_eq!(err.kind(), std::io::ErrorKind::TimedOut);
+    }
+
+    #[test_log::test(tokio::test)]
+    async fn outgoing_max_timeout() {
+        let (tx, _rx) = tokio::io::duplex(1);
+        let mut out = Outgoing::new(tx, |_, _| async {}).with_timeout(Duration::MAX);
+        out.write_all(&[0; 16]).await.unwrap();
+        tokio::time::sleep(Duration::from_millis(10)).await;
+        assert!(out.err.get().is_none());
+    }
+
     #[test_log::test(tokio::test)]
     async fn ingress_chunked_data() {
         let index = std::sync::Mutex::new(IndexTrie::Empty);