fix(storage): harden offline drive fail-fast paths (#2564)

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: houseme <4829346+houseme@users.noreply.github.com>

crates/rio/src/http_reader.rs

@@ -27,8 +27,10 @@ use std::ops::Not as _;
 use std::pin::Pin;
 use std::sync::LazyLock;
 use std::task::{Context, Poll};
+use std::time::Duration;
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tokio::sync::mpsc;
+use tokio::time::{self, Sleep};
 use tokio_util::io::StreamReader;
 use tokio_util::sync::PollSender;
 use tracing::error;
@@ -143,6 +145,8 @@ pin_project! {
         method: Method,
         headers: HeaderMap,
         track_internode_metrics: bool,
+        stall_timeout: Option<Duration>,
+        stall_timer: Option<Pin<Box<Sleep>>>,
         #[pin]
         inner: StreamReader<Pin<Box<dyn Stream<Item=std::io::Result<Bytes>>+Send+Sync>>, Bytes>,
     }
@@ -151,8 +155,19 @@ pin_project! {
 impl HttpReader {
     pub async fn new(url: String, method: Method, headers: HeaderMap, body: Option<Vec<u8>>) -> io::Result<Self> {
         // http_log!("[HttpReader::new] url: {url}, method: {method:?}, headers: {headers:?}");
-        Self::with_capacity(url, method, headers, body, 0).await
+        Self::with_capacity_and_stall_timeout(url, method, headers, body, 0, None).await
     }
+
+    pub async fn new_with_stall_timeout(
+        url: String,
+        method: Method,
+        headers: HeaderMap,
+        body: Option<Vec<u8>>,
+        stall_timeout: Option<Duration>,
+    ) -> io::Result<Self> {
+        Self::with_capacity_and_stall_timeout(url, method, headers, body, 0, stall_timeout).await
+    }
+
     /// Create a new HttpReader from a URL. The request is performed immediately.
     pub async fn with_capacity(
         url: String,
@@ -160,6 +175,17 @@ impl HttpReader {
         headers: HeaderMap,
         body: Option<Vec<u8>>,
         _read_buf_size: usize,
+    ) -> io::Result<Self> {
+        Self::with_capacity_and_stall_timeout(url, method, headers, body, _read_buf_size, None).await
+    }
+
+    async fn with_capacity_and_stall_timeout(
+        url: String,
+        method: Method,
+        headers: HeaderMap,
+        body: Option<Vec<u8>>,
+        _read_buf_size: usize,
+        stall_timeout: Option<Duration>,
     ) -> io::Result<Self> {
         let track_internode_metrics = is_internode_rpc_url(&url);
         let client = get_http_client(&url);
@@ -202,6 +228,8 @@ impl HttpReader {
             method,
             headers,
             track_internode_metrics,
+            stall_timer: stall_timeout.map(|timeout| Box::pin(time::sleep(timeout))),
+            stall_timeout,
         })
     }
     pub fn url(&self) -> &str {
@@ -216,16 +244,40 @@ impl HttpReader {
 }
 
 impl AsyncRead for HttpReader {
-    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
+    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
+        let mut this = self.project();
+
         let filled_before = buf.filled().len();
-        match Pin::new(&mut self.inner).poll_read(cx, buf) {
+        match this.inner.as_mut().poll_read(cx, buf) {
             Poll::Ready(Ok(())) => {
                 let bytes_read = buf.filled().len().saturating_sub(filled_before);
-                if self.track_internode_metrics && bytes_read > 0 {
+                if *this.track_internode_metrics && bytes_read > 0 {
                     global_internode_metrics().record_recv_bytes(bytes_read);
                 }
+                if bytes_read > 0 {
+                    if let Some(stall_timeout) = *this.stall_timeout {
+                        *this.stall_timer = Some(Box::pin(time::sleep(stall_timeout)));
+                    }
+                } else {
+                    *this.stall_timer = None;
+                }
                 Poll::Ready(Ok(()))
             }
+            Poll::Pending => {
+                if let Some(timer) = this.stall_timer.as_mut()
+                    && timer.as_mut().poll(cx).is_ready()
+                {
+                    if *this.track_internode_metrics {
+                        global_internode_metrics().record_error();
+                    }
+                    Poll::Ready(Err(Error::new(
+                        io::ErrorKind::TimedOut,
+                        "HttpReader stall timeout: no data received before deadline",
+                    )))
+                } else {
+                    Poll::Pending
+                }
+            }
             other => other,
         }
     }
@@ -570,8 +622,9 @@ impl AsyncWrite for HttpWriter {
 mod tests {
     use super::*;
     use axum::{Router, body::Body, extract::State, http::StatusCode, response::IntoResponse, routing::get};
+    use futures::stream::{self, StreamExt as _};
     use http_body_util::BodyExt as _;
-    use std::io::IoSlice;
+    use std::io::{self, IoSlice};
     use std::sync::{
         Arc,
         atomic::{AtomicUsize, Ordering},
@@ -595,6 +648,12 @@ mod tests {
         (StatusCode::OK, Body::from("hello"))
     }
 
+    async fn get_stalling_stream(State(state): State<TestState>) -> impl IntoResponse {
+        state.get_count.fetch_add(1, Ordering::SeqCst);
+        let body_stream = stream::once(async { Ok::<Bytes, io::Error>(Bytes::from_static(b"hello")) }).chain(stream::pending());
+        (StatusCode::OK, Body::from_stream(body_stream))
+    }
+
     async fn reject_head(State(state): State<TestState>) -> impl IntoResponse {
         state.head_count.fetch_add(1, Ordering::SeqCst);
         StatusCode::METHOD_NOT_ALLOWED
@@ -612,6 +671,7 @@ mod tests {
         let addr = listener.local_addr().unwrap();
         let app = Router::new()
             .route("/stream", get(get_stream).head(reject_head).put(accept_put))
+            .route("/stall", get(get_stalling_stream))
             .with_state(state);
 
         let handle = tokio::spawn(async move {
@@ -637,6 +697,31 @@ mod tests {
         handle.abort();
     }
 
+    #[tokio::test]
+    async fn http_reader_stall_timeout_triggers_after_progress_stops() {
+        let state = TestState::default();
+        let (base_url, handle) = start_test_server(state.clone()).await;
+        let url = base_url.replace("/stream", "/stall");
+
+        let mut reader =
+            HttpReader::new_with_stall_timeout(url, Method::GET, HeaderMap::new(), None, Some(Duration::from_millis(20)))
+                .await
+                .unwrap();
+
+        let mut first = [0u8; 5];
+        reader.read_exact(&mut first).await.unwrap();
+        assert_eq!(&first, b"hello");
+
+        let mut next = [0u8; 1];
+        let err = tokio::time::timeout(Duration::from_secs(1), reader.read(&mut next))
+            .await
+            .expect("stall timeout should wake reader")
+            .expect_err("reader should return a timeout error");
+        assert_eq!(err.kind(), io::ErrorKind::TimedOut);
+
+        handle.abort();
+    }
+
     #[tokio::test]
     async fn http_writer_does_not_send_empty_preflight_put() {
         let state = TestState::default();