Keep every acknowledged write through a rolling restart

hyperbytedb-proxy/src/proxy.rs

@@ -24,7 +24,7 @@ use futures::StreamExt;
 use futures::ready;
 use http_body_util::BodyExt;
 
-use crate::backend::InflightGuard;
+use crate::backend::{Health, InflightGuard};
 use crate::pool::BackendPool;
 
 /// Headers that hop-by-hop semantics (RFC 7230 §6.1) say we must not forward.
@@ -148,7 +148,22 @@ pub async fn handle(State(state): State<ProxyState>, req: Request) -> Response {
                     .record(started.elapsed().as_secs_f64());
                 return attach_inflight_guard(resp, guard);
             }
-            ForwardOutcome::Retryable { status, msg } => {
+            ForwardOutcome::Retryable { status, msg, mark } => {
+                // Take it out of rotation now. Its health probe only notices
+                // on its next run, about a second later, and until then the
+                // pick below handed the same dead backend straight back
+                // (always, in a one-backend pool): every retry failed within
+                // milliseconds and the client got 503 instead of being held
+                // for `hold_timeout`. The probe restores it once it answers.
+                if let Some(health) = mark
+                    && current.set_health(health).is_some()
+                {
+                    tracing::info!(
+                        backend = %current.addr,
+                        to = health.as_str(),
+                        "backend taken out of rotation after a failed forward"
+                    );
+                }
                 tracing::info!(
                     attempt,
                     backend = %current.addr,
@@ -213,6 +228,9 @@ enum ForwardOutcome {
     Retryable {
         status: Option<StatusCode>,
         msg: Option<String>,
+        /// What the failure says about the backend's health, applied at once
+        /// rather than at its next probe.
+        mark: Option<Health>,
     },
     /// Don't retry: the failure is the client's fault (4xx) and shouldn't be
     /// blamed on backend health.
@@ -243,9 +261,12 @@ async fn forward_once(
     let resp = match rb.send().await {
         Ok(r) => r,
         Err(e) => {
-            // Transport error → backend candidate failed; retry elsewhere.
+            // Transport error → backend candidate failed; retry elsewhere. A
+            // refused or reset connection means the backend is gone (a pod
+            // stopping); a timeout may be a slow request on a live one.
             return ForwardOutcome::Retryable {
                 status: None,
+                mark: (!e.is_timeout()).then_some(Health::Down),
                 msg: Some(e.to_string()),
             };
         }
@@ -260,6 +281,7 @@ async fn forward_once(
     if matches!(upstream_status.as_u16(), 502 | 504,) {
         return ForwardOutcome::Retryable {
             status: Some(upstream_status),
+            mark: None,
             msg: Some("backend returned bad-gateway/timeout".into()),
         };
     }
@@ -273,13 +295,15 @@ async fn forward_once(
             Err(e) => {
                 return ForwardOutcome::Retryable {
                     status: Some(upstream_status),
+                    mark: None,
                     msg: Some(format!("response body read: {e}")),
                 };
             }
         };
         if looks_like_drain(&body_bytes) {
             return ForwardOutcome::Retryable {
                 status: Some(upstream_status),
+                mark: Some(Health::Draining),
                 msg: Some("backend reports draining/syncing".into()),
             };
         }
@@ -298,6 +322,7 @@ async fn forward_once(
             Err(e) => {
                 return ForwardOutcome::Retryable {
                     status: Some(upstream_status),
+                    mark: None,
                     msg: Some(format!("response body read: {e}")),
                 };
             }

hyperbytedb-proxy/tests/hold_on_failed_backend.rs

@@ -0,0 +1,88 @@
+//! A request whose backend has just gone down is held until a backend is
+//! back, not failed with 503.
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use axum::Router;
+use axum::http::StatusCode;
+use axum::routing::{any, get};
+use hyperbytedb_proxy::backend::{Backend, Health};
+use hyperbytedb_proxy::config::ProxyConfig;
+use hyperbytedb_proxy::pool::BackendPool;
+use hyperbytedb_proxy::proxy::{ProxyState, handle};
+
+fn config(port: u16) -> ProxyConfig {
+    ProxyConfig {
+        listen_addr: "127.0.0.1:0".into(),
+        admin_listen_addr: "127.0.0.1:0".into(),
+        backend_service: "127.0.0.1".into(),
+        backend_port: port,
+        discovery_interval: Duration::from_secs(3600),
+        health_interval: Duration::from_millis(100),
+        health_path: "/health".into(),
+        health_timeout: Duration::from_millis(200),
+        request_timeout: Duration::from_secs(5),
+        hold_timeout: Duration::from_secs(5),
+        max_retries: 3,
+        shutdown_grace: Duration::from_secs(1),
+        self_ip: None,
+        http2_prior_knowledge: false,
+    }
+}
+
+/// The backend went down a moment ago and its last probe still says Active
+/// (probes run about every second): a pool of one, as for a single-node pool
+/// in a rolling restart. The request used to fail at once with 503, every
+/// retry handed the same dead backend; now the failed forward takes it out
+/// of rotation and the request waits for it to come back.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn a_request_to_a_backend_that_just_went_down_is_held() {
+    // A port with nothing listening yet.
+    let port = {
+        let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+        l.local_addr().unwrap().port()
+    };
+    let pool = BackendPool::new(config(port)).unwrap();
+    let backend = Arc::new(Backend::new("127.0.0.1".parse().unwrap(), port));
+    backend.set_health(Health::Active);
+    pool.set_backends_for_test(vec![Arc::clone(&backend)]).await;
+
+    let proxy = Router::new()
+        .route("/write", any(handle))
+        .with_state(ProxyState::new(Arc::clone(&pool)).unwrap());
+    let proxy_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let proxy_url = format!("http://{}", proxy_listener.local_addr().unwrap());
+    tokio::spawn(async move { axum::serve(proxy_listener, proxy).await.unwrap() });
+
+    let request = tokio::spawn(async move {
+        reqwest::Client::new()
+            .post(format!("{proxy_url}/write?db=d"))
+            .body("m v=1")
+            .send()
+            .await
+            .unwrap()
+            .status()
+    });
+
+    // The backend comes back 500 ms later, and its probes start.
+    tokio::time::sleep(Duration::from_millis(500)).await;
+    assert!(
+        !request.is_finished(),
+        "held while the backend is down, not failed"
+    );
+    let app = Router::new()
+        .route("/write", any(|| async { StatusCode::NO_CONTENT }))
+        .route("/health", get(|| async { StatusCode::OK }));
+    let listener = tokio::net::TcpListener::bind(("127.0.0.1", port))
+        .await
+        .unwrap();
+    tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+    tokio::spawn(Arc::clone(&pool).run_health());
+
+    let status = tokio::time::timeout(Duration::from_secs(5), request)
+        .await
+        .expect("answered once the backend is back")
+        .unwrap();
+    assert_eq!(status, reqwest::StatusCode::NO_CONTENT);
+}