proxy.rs

//! Reverse-proxy axum handler.
//!
//! Routing rules:
//! 1. Pick an `Active` backend (round-robin).
//! 2. If none → wait up to `hold_timeout` for one (this is the "rolling
//!    restart invisible to clients" knob).
//! 3. Forward the request. If it fails with a transient error or the backend
//!    returns 503-with-Draining, retry against another backend up to
//!    `max_retries` times.
//! 4. On final failure surface 503 to the client.

use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Instant;

use axum::body::Body;
use axum::extract::{Request, State};
use axum::http::{HeaderName, StatusCode};
use axum::response::{IntoResponse, Response};
use bytes::Bytes;
use futures::Stream;
use futures::StreamExt;
use futures::ready;
use http_body_util::BodyExt;

use crate::backend::InflightGuard;
use crate::pool::BackendPool;

/// Headers that hop-by-hop semantics (RFC 7230 §6.1) say we must not forward.
const HOP_BY_HOP: &[&str] = &[
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailer",
    "transfer-encoding",
    "upgrade",
    // `host` is rewritten when we re-target the request.
    "host",
];

#[derive(Clone)]
pub struct ProxyState {
    pub pool: Arc<BackendPool>,
    pub client: reqwest::Client,
}

impl ProxyState {
    pub fn new(pool: Arc<BackendPool>) -> anyhow::Result<Self> {
        let cfg = pool.config();
        let mut client_builder = reqwest::Client::builder()
            .timeout(cfg.request_timeout)
            // Per-host pool sized to absorb a moderate burst without
            // re-handshaking; keep_alive_while_idle keeps connections warm
            // across the typical inter-request gap of a Grafana refresh.
            .pool_max_idle_per_host(64)
            .pool_idle_timeout(std::time::Duration::from_secs(90))
            .tcp_nodelay(true);
        if cfg.http2_prior_knowledge {
            client_builder = client_builder.http2_prior_knowledge();
        }
        let client = client_builder.build()?;
        Ok(Self { pool, client })
    }
}

/// The single fallback handler that captures every URI/method.
pub async fn handle(State(state): State<ProxyState>, req: Request) -> Response {
    let started = Instant::now();
    let cfg = state.pool.config();

    // Buffer the request body once. We may need to send it more than once if
    // a backend returns a transient failure mid-restart. Wrapped in `Arc` so
    // retries clone a pointer, not the payload.
    let (parts, body) = req.into_parts();
    let body_bytes = match body.collect().await {
        Ok(c) => Arc::new(c.to_bytes()),
        Err(e) => {
            tracing::warn!(error = %e, "failed to read incoming request body");
            return error_response(StatusCode::BAD_REQUEST, "could not read request body");
        }
    };

    let path_query = parts
        .uri
        .path_and_query()
        .map(|pq| pq.as_str().to_string())
        .unwrap_or_else(|| parts.uri.path().to_string());

    // 1. Wait for at least one active backend (cheap if one already exists).
    let first = match state.pool.wait_for_active(cfg.hold_timeout).await {
        Some(b) => b,
        None => {
            tracing::warn!(
                hold_timeout_secs = cfg.hold_timeout.as_secs(),
                "no active backend available; held until timeout"
            );
            metrics::counter!("hyperbytedb_proxy_no_backend_total").increment(1);
            return error_response(
                StatusCode::SERVICE_UNAVAILABLE,
                "no healthy hyperbytedb backend available",
            );
        }
    };

    // 2. Try the first backend, then up to `max_retries` more.
    // Initial `None` here is intentionally overwritten by the Retryable arm
    // before the format! at the bottom of the loop ever reads it.
    #[allow(unused_assignments)]
    let mut last_status: Option<StatusCode> = None;
    #[allow(unused_assignments)]
    let mut last_err: Option<String> = None;
    let mut attempt: u32 = 0;
    let mut current = first;

    loop {
        let guard = current.enter();
        let url = format!("{}{}", current.origin, path_query);
        tracing::debug!(
            attempt,
            backend = %current.addr,
            method = %parts.method,
            path = %path_query,
            "forwarding"
        );

        let outcome = forward_once(
            &state.client,
            &url,
            &parts.method,
            &parts.headers,
            Arc::clone(&body_bytes),
        )
        .await;

        match outcome {
            ForwardOutcome::Ok(resp) => {
                let status = resp.status();
                metrics::counter!(
                    "hyperbytedb_proxy_requests_total",
                    "outcome" => "ok",
                    "status" => status.as_u16().to_string(),
                )
                .increment(1);
                metrics::histogram!("hyperbytedb_proxy_request_duration_seconds")
                    .record(started.elapsed().as_secs_f64());
                return attach_inflight_guard(resp, guard);
            }
            ForwardOutcome::Retryable { status, msg } => {
                tracing::info!(
                    attempt,
                    backend = %current.addr,
                    ?status,
                    err = msg.as_deref().unwrap_or("(none)"),
                    "retryable failure, will pick another backend"
                );
                last_status = status;
                last_err = msg;
                // Attempt failed before a response was handed to the client.
                drop(guard);
            }
            ForwardOutcome::Fatal(resp) => {
                metrics::counter!(
                    "hyperbytedb_proxy_requests_total",
                    "outcome" => "fatal",
                )
                .increment(1);
                return attach_inflight_guard(resp, guard);
            }
        }

        if attempt >= cfg.max_retries {
            break;
        }
        attempt += 1;

        // Pick a different backend; if there isn't one, hold briefly.
        current = match state.pool.pick_active_excluding(&current).await {
            Some(b) => b,
            None => match state.pool.wait_for_active(cfg.hold_timeout).await {
                Some(b) => b,
                None => break,
            },
        };
    }

    metrics::counter!(
        "hyperbytedb_proxy_requests_total",
        "outcome" => "exhausted",
    )
    .increment(1);
    metrics::histogram!("hyperbytedb_proxy_request_duration_seconds")
        .record(started.elapsed().as_secs_f64());

    let body = format!(
        r#"{{"status":"fail","message":"all backends exhausted (last status {:?}, err {:?})"}}"#,
        last_status.map(|s| s.as_u16()),
        last_err.as_deref().unwrap_or("none")
    );
    (
        StatusCode::SERVICE_UNAVAILABLE,
        [("content-type", "application/json")],
        body,
    )
        .into_response()
}

enum ForwardOutcome {
    Ok(Response),
    /// Pickable for another backend.
    Retryable {
        status: Option<StatusCode>,
        msg: Option<String>,
    },
    /// Don't retry: the failure is the client's fault (4xx) and shouldn't be
    /// blamed on backend health.
    Fatal(Response),
}

async fn forward_once(
    client: &reqwest::Client,
    url: &str,
    method: &http::Method,
    headers: &http::HeaderMap,
    body: Arc<Bytes>,
) -> ForwardOutcome {
    let mut rb = client.request(method.clone(), url);
    for (name, value) in headers {
        if HOP_BY_HOP
            .iter()
            .any(|h| name.as_str().eq_ignore_ascii_case(h))
        {
            continue;
        }
        rb = rb.header(name, value);
    }
    if !body.is_empty() {
        rb = rb.body(body.as_ref().clone());
    }

    let resp = match rb.send().await {
        Ok(r) => r,
        Err(e) => {
            // Transport error → backend candidate failed; retry elsewhere.
            return ForwardOutcome::Retryable {
                status: None,
                msg: Some(e.to_string()),
            };
        }
    };

    let status = resp.status();
    let upstream_status =
        StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);

    // 502/504 from the backend itself = transient infra problem upstream;
    // try another node. Body is not needed for the retry decision.
    if matches!(upstream_status.as_u16(), 502 | 504,) {
        return ForwardOutcome::Retryable {
            status: Some(upstream_status),
            msg: Some("backend returned bad-gateway/timeout".into()),
        };
    }

    // 503 with the well-known draining marker → another backend may serve us.
    // These envelopes are small JSON; buffer to inspect before deciding.
    if upstream_status == StatusCode::SERVICE_UNAVAILABLE {
        let resp_headers = resp.headers().clone();
        let body_bytes = match resp.bytes().await {
            Ok(b) => b,
            Err(e) => {
                return ForwardOutcome::Retryable {
                    status: Some(upstream_status),
                    msg: Some(format!("response body read: {e}")),
                };
            }
        };
        if looks_like_drain(&body_bytes) {
            return ForwardOutcome::Retryable {
                status: Some(upstream_status),
                msg: Some("backend reports draining/syncing".into()),
            };
        }
        let resp = build_buffered_response(upstream_status, &resp_headers, body_bytes);
        return if upstream_status.is_client_error() {
            ForwardOutcome::Fatal(resp)
        } else {
            ForwardOutcome::Ok(resp)
        };
    }

    if upstream_status.is_client_error() {
        let resp_headers = resp.headers().clone();
        let body_bytes = match resp.bytes().await {
            Ok(b) => b,
            Err(e) => {
                return ForwardOutcome::Retryable {
                    status: Some(upstream_status),
                    msg: Some(format!("response body read: {e}")),
                };
            }
        };
        return ForwardOutcome::Fatal(build_buffered_response(
            upstream_status,
            &resp_headers,
            body_bytes,
        ));
    }

    // Success and other non-retryable responses: stream upstream body through
    // without buffering the full payload in proxy memory.
    ForwardOutcome::Ok(build_streaming_response(upstream_status, resp))
}

fn build_buffered_response(
    status: StatusCode,
    resp_headers: &reqwest::header::HeaderMap,
    body_bytes: Bytes,
) -> Response {
    let mut out = Response::builder().status(status);
    copy_upstream_headers(
        resp_headers,
        out.headers_mut().expect("response builder has headers map"),
    );
    out.body(Body::from(body_bytes))
        .expect("axum response body construction is infallible")
}

fn build_streaming_response(status: StatusCode, resp: reqwest::Response) -> Response {
    let resp_headers = resp.headers().clone();
    let mut out = Response::builder().status(status);
    copy_upstream_headers(
        &resp_headers,
        out.headers_mut().expect("response builder has headers map"),
    );
    let stream = resp
        .bytes_stream()
        .map(|result| result.map_err(std::io::Error::other));
    out.body(Body::from_stream(stream))
        .expect("axum response body construction is infallible")
}

fn copy_upstream_headers(upstream: &reqwest::header::HeaderMap, out_headers: &mut http::HeaderMap) {
    for (name, value) in upstream.iter() {
        if HOP_BY_HOP
            .iter()
            .any(|h| name.as_str().eq_ignore_ascii_case(h))
        {
            continue;
        }
        match HeaderName::from_bytes(name.as_ref()) {
            Ok(header_name) => {
                out_headers.insert(header_name, value.clone());
            }
            Err(_) => {
                tracing::warn!(header = ?name, "skipping upstream header with invalid name");
            }
        }
    }
}

/// Keep the backend inflight counter elevated until the response body is fully
/// consumed (including streamed query results).
fn attach_inflight_guard(resp: Response, guard: InflightGuard) -> Response {
    let (parts, body) = resp.into_parts();
    let stream = body
        .into_data_stream()
        .map(|result| result.map_err(std::io::Error::other));
    let guarded = GuardedBodyStream {
        inner: Box::pin(stream),
        guard: Some(guard),
    };
    Response::from_parts(parts, Body::from_stream(guarded))
}

struct GuardedBodyStream {
    inner: Pin<Box<dyn Stream<Item = Result<Bytes, std::io::Error>> + Send>>,
    guard: Option<InflightGuard>,
}

impl Stream for GuardedBodyStream {
    type Item = Result<Bytes, std::io::Error>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.as_mut().get_mut();
        match ready!(this.inner.as_mut().poll_next(cx)) {
            Some(item) => Poll::Ready(Some(item)),
            None => {
                this.guard.take();
                Poll::Ready(None)
            }
        }
    }
}

/// Loose match against hyperbytedb not-ready JSON envelopes (503 /health).
pub fn looks_like_drain(body: &Bytes) -> bool {
    // Very loose match against the JSON envelopes hyperbytedb uses for
    // not-ready states. Avoids parsing JSON for every response.
    let s = match std::str::from_utf8(body) {
        Ok(s) => s,
        Err(_) => return false,
    };
    s.contains("\"status\":\"warn\"")
        || s.contains("Draining")
        || s.contains("Syncing")
        || s.contains("Joining")
        || s.contains("Leaving")
        || s.contains("draining")
}

fn error_response(status: StatusCode, msg: &str) -> Response {
    let body = format!(
        r#"{{"status":"fail","message":"{}"}}"#,
        msg.replace('"', "\\\"")
    );
    (status, [("content-type", "application/json")], body).into_response()
}

/// Reject paths outside the public InfluxDB v1 surface.
pub async fn not_found() -> Response {
    error_response(
        StatusCode::NOT_FOUND,
        "only /write and /query are exposed on the public listener",
    )
}