foundation_proxy 0.1.2

//! Health probes.
//!
//! WHY: Only healthy backends should receive traffic. Each backend with a
//! configured health check is probed on a fixed interval; consecutive
//! successes/failures drive a hysteresis state machine so a single blip neither
//! ejects a good backend nor readmits a flapping one. An unhealthy backend is
//! removed from rotation (`BackendRuntime::set_healthy(false)`) and readmitted
//! only after it passes `healthy_threshold` consecutive probes.
//!
//! WHAT: [`ProbeState`] (the pure consecutive-count state machine, unit-tested
//! without a network) and [`spawn_service_probes`] which runs one probe thread
//! per backend until a shutdown signal fires.
//!
//! HOW: HTTP/HTTPS backends are probed with a real `GET` to the configured
//! path via the shared client; `tcp://` backends are probed with a TCP connect;
//! `grpc://` backends are probed with a native `grpc.health.v1.Health/Check`
//! over H2 (the standard gRPC health protocol — a `SERVING` verdict is healthy).
//! A 2xx/3xx response (or a successful connect / `SERVING` verdict) is a
//! success; anything else is a failure.

use std::net::TcpStream;
use std::sync::Arc;
use std::thread::JoinHandle;
use std::time::{Duration, Instant};

use bytes::Bytes;
use foundation_core::synca::OnSignal;
use foundation_core::url::Uri;
use foundation_core::valtron::{block_on_future, sleep_async};
use foundation_netio::shared::http::{
    Proto, RequestDescriptor, SimpleHeader, SimpleHeaders, SimpleMethod, SimpleUrl,
};

use foundation_connectrpc::{
    CancelSignal, ClientExchange, ClientExchangeConfig, GrpcClient, H2Transport, ProtocolClient,
    Spec, Transport,
};
use futures_util::future::{select, Either};
use futures_util::FutureExt;

use crate::config::{BackendProtocol, HealthCheckConfig};
use crate::forward::SharedHttpClient;
use crate::runtime::{BackendRuntime, ServiceRuntime};

/// Consecutive-count health state machine (Decision 14 `HealthCheckConfig`).
///
/// Tracks runs of successes and failures and flips the health verdict only when
/// a run reaches its threshold. `record` returns `Some(new_health)` exactly on
/// the probe that causes a transition, and `None` otherwise.
#[derive(Debug, Clone)]
pub struct ProbeState {
    healthy_threshold: u32,
    unhealthy_threshold: u32,
    consecutive_successes: u32,
    consecutive_failures: u32,
    healthy: bool,
}

impl ProbeState {
    /// Create a state machine. `initial_healthy` seeds the verdict — backends
    /// begin healthy so a service with a slow first probe still serves.
    #[must_use]
    pub fn new(healthy_threshold: u32, unhealthy_threshold: u32, initial_healthy: bool) -> Self {
        Self {
            healthy_threshold: healthy_threshold.max(1),
            unhealthy_threshold: unhealthy_threshold.max(1),
            consecutive_successes: 0,
            consecutive_failures: 0,
            healthy: initial_healthy,
        }
    }

    #[must_use]
    pub fn is_healthy(&self) -> bool {
        self.healthy
    }

    /// Record one probe outcome. Returns `Some(new_health)` on a transition.
    ///
    /// A success resets the failure run and, if healthy is not already set,
    /// flips to healthy once `healthy_threshold` consecutive successes land. A
    /// failure is symmetric against `unhealthy_threshold`.
    pub fn record(&mut self, success: bool) -> Option<bool> {
        if success {
            self.consecutive_failures = 0;
            self.consecutive_successes = self.consecutive_successes.saturating_add(1);
            if !self.healthy && self.consecutive_successes >= self.healthy_threshold {
                self.healthy = true;
                return Some(true);
            }
        } else {
            self.consecutive_successes = 0;
            self.consecutive_failures = self.consecutive_failures.saturating_add(1);
            if self.healthy && self.consecutive_failures >= self.unhealthy_threshold {
                self.healthy = false;
                return Some(false);
            }
        }
        None
    }
}

/// Handle to a service's running probe threads; joins them on shutdown.
#[derive(Debug)]
pub struct HealthMonitor {
    handles: Vec<JoinHandle<()>>,
}

impl HealthMonitor {
    /// Join every probe thread. Callers turn the shared shutdown signal on
    /// first, which breaks each probe loop at its next checkpoint.
    pub fn join(self) {
        for handle in self.handles {
            let _ = handle.join();
        }
    }
}

/// Spawn one probe thread per backend of `service` that has a health check.
///
/// Returns a [`HealthMonitor`]; the threads run until `shutdown` is turned on.
/// A service without a `health_check` config spawns nothing — its backends stay
/// at their initial healthy state.
#[must_use]
pub fn spawn_service_probes(
    service: &Arc<ServiceRuntime>,
    client: SharedHttpClient,
    shutdown: Arc<OnSignal>,
) -> HealthMonitor {
    let Some(config) = service.config().health_check.clone() else {
        return HealthMonitor { handles: Vec::new() };
    };

    let mut handles = Vec::new();
    for backend in service.backends() {
        let backend = Arc::clone(backend);
        let config = config.clone();
        let client = client.clone();
        let shutdown = Arc::clone(&shutdown);
        let handle = std::thread::spawn(move || {
            probe_loop(&backend, &config, &client, &shutdown);
        });
        handles.push(handle);
    }
    HealthMonitor { handles }
}

/// Run the probe loop for one backend until shutdown.
fn probe_loop(
    backend: &Arc<BackendRuntime>,
    config: &HealthCheckConfig,
    client: &SharedHttpClient,
    shutdown: &Arc<OnSignal>,
) {
    let mut state = ProbeState::new(
        config.healthy_threshold,
        config.unhealthy_threshold,
        backend.is_healthy(),
    );

    loop {
        if shutdown.probe() {
            return;
        }

        let success = probe_once(backend, config, client);
        if let Some(now_healthy) = state.record(success) {
            backend.set_healthy(now_healthy);
            if now_healthy {
                tracing::info!(url = %backend.url(), "backend became healthy");
            } else {
                tracing::warn!(url = %backend.url(), "backend became unhealthy");
            }
        }

        // Sleep the interval in small slices so shutdown is responsive.
        if !sleep_interruptible(config.interval, shutdown) {
            return;
        }
    }
}

/// Perform a single probe. `true` = success.
fn probe_once(
    backend: &Arc<BackendRuntime>,
    config: &HealthCheckConfig,
    client: &SharedHttpClient,
) -> bool {
    match backend.protocol() {
        BackendProtocol::Tcp => probe_tcp(&backend.target().authority(), config.timeout),
        BackendProtocol::Udp => probe_udp(&backend.target().authority(), config.timeout),
        // `grpc://` backends answer the standard gRPC health protocol over H2,
        // not an HTTP GET path.
        BackendProtocol::Grpc => probe_grpc_health(backend.url(), config.timeout),
        // `connect` / `grpc-web` are HTTP-family; probe them with the HTTP path
        // probe (a Connect backend has no gRPC health endpoint).
        BackendProtocol::Http
        | BackendProtocol::Https
        | BackendProtocol::Connect
        | BackendProtocol::GrpcWeb => {
            probe_http(backend.url(), &config.path, config.timeout, client)
        }
    }
}

/// TCP connect probe with a timeout.
fn probe_tcp(authority: &str, timeout: Duration) -> bool {
    match authority.to_socket_addrs_first() {
        Some(addr) => TcpStream::connect_timeout(&addr, timeout).is_ok(),
        None => TcpStream::connect(authority).is_ok(),
    }
}

/// UDP echo probe: send a 1-byte probe and return `true` if a response arrives
/// within the timeout. Uses a non-empty payload because `socat EXEC:cat` and
/// similar UDP echoers produce no response for an empty datagram.
fn probe_udp(authority: &str, timeout: Duration) -> bool {
    use std::net::UdpSocket;
    let Ok(socket) = UdpSocket::bind("0.0.0.0:0") else {
        return false;
    };
    let _ = socket.set_read_timeout(Some(timeout));
    if socket.send_to(&[1u8], authority).is_err() {
        return false;
    }
    let mut buf = [0u8; 16];
    socket.recv_from(&mut buf).is_ok()
}

/// HTTP GET probe; success is any 2xx or 3xx response.
fn probe_http(base_url: &str, path: &str, timeout: Duration, client: &SharedHttpClient) -> bool {
    let url = join_url(base_url, path);
    let probe_client = client
        .clone()
        .connect_timeout(timeout)
        .read_timeout(timeout);
    let Ok(builder) = probe_client.get(&url) else {
        return false;
    };
    let Ok(request) = probe_client.request(builder) else {
        return false;
    };
    match request.send() {
        Ok(resp) => {
            let code = resp.get_status().into_usize();
            (200..400).contains(&code)
        }
        Err(err) => {
            tracing::trace!(%url, ?err, "health probe request failed");
            false
        }
    }
}

// ── gRPC health (grpc.health.v1.Health/Check) ──────────────────────────────

/// `ServingStatus.SERVING` — the only value that counts as healthy.
const SERVING: u64 = 0;

/// Decode `HealthCheckResponse.status` (field 1, varint) from its protobuf.
/// `SERVING` (0) is the proto3 default, so an omitted field means the same as an
/// empty message. Unknown trailing fields are skipped; a malformed message
/// degrades to `SERVING` so the verdict comes from `grpc-status` instead.
fn decode_health_status(msg: &[u8]) -> u64 {
    let mut i = 0;
    while i < msg.len() {
        let tag = varint(msg, &mut i);
        match tag & 0x07 {
            0 => {
                let value = varint(msg, &mut i);
                if tag >> 3 == 1 {
                    return value;
                }
            }
            2 => {
                let len = varint(msg, &mut i) as usize;
                i = i.saturating_add(len);
            }
            5 => i = i.saturating_add(4), // 32-bit
            1 => i = i.saturating_add(8), // 64-bit
            _ => return SERVING,
        }
    }
    SERVING
}

/// Decode one base-128 varint from `buf` starting at `*i`, advancing `*i`.
fn varint(buf: &[u8], i: &mut usize) -> u64 {
    let mut result = 0u64;
    let mut shift = 0u32;
    loop {
        let Some(&byte) = buf.get(*i) else {
            break;
        };
        *i += 1;
        result |= u64::from(byte & 0x7f) << shift;
        if byte & 0x80 == 0 {
            break;
        }
        shift += 7;
        if shift >= 64 {
            break;
        }
    }
    result
}

/// Drive one `grpc.health.v1.Health/Check` round trip over H2. The empty service
/// string (`HealthCheckRequest { service: "" }`) probes overall health; the
/// single response message is decoded to its `ServingStatus`, cross-checked
/// against the trailing `grpc-status`. `timeout` bounds the whole exchange: the
/// pump tasks are raced against a timer, and a win by the timer yields `Err`.
async fn grpc_health_check(url: String, timeout: Duration) -> Result<bool, String> {
    let uri = Uri::parse(&url).map_err(|e| format!("parse health uri {url}: {e}"))?;
    let mut headers = SimpleHeaders::new();
    headers.insert(
        SimpleHeader::CONTENT_TYPE,
        vec!["application/grpc".to_string()],
    );
    headers.insert(
        SimpleHeader::from("te".to_string()),
        vec!["trailers".to_string()],
    );

    let out_request = RequestDescriptor {
        proto: Proto::HTTP20,
        request_url: SimpleUrl::url_only(&url),
        request_uri: uri,
        headers,
        method: SimpleMethod::POST,
    };
    let stream = H2Transport::new()
        .open(out_request)
        .map_err(|e| format!("open h2 health upstream {url}: {e}"))?;
    let outbound = GrpcClient
        .new_conn(
            &Spec::empty(true),
            SimpleHeaders::new(),
            stream,
            ClientExchangeConfig {
                cancel: CancelSignal::new(),
                read_max_bytes: 0,
            },
        )
        .map_err(|e| format!("build grpc health exchange: {e}"))?;
    let ClientExchange {
        conn,
        reader_task,
        writer_task,
    } = outbound;
    let (mut out_send, mut out_recv) = conn.split();

    // The receiver stashes its outcome in a shared slot so the timeout race can
    // still surface it after the pump `join_all` completes.
    let result: Arc<std::sync::Mutex<Option<Result<bool, String>>>> =
        Arc::new(std::sync::Mutex::new(None));

    let sender = async move {
        // `HealthCheckRequest { service: "" }` serializes to zero bytes.
        let _ = out_send.send(Bytes::new()).await;
        let _ = out_send.close_send().await;
    };

    let result_ref = Arc::clone(&result);
    let receiver = async move {
        let outcome = async {
            let mut msg: Option<Bytes> = None;
            loop {
                match out_recv.receive().await {
                    Ok(Some(bytes)) => msg = Some(bytes),
                    Ok(None) => break,
                    Err(e) => return Err(format!("health receive failed: {e}")),
                }
            }
            let trailers = out_recv
                .response_trailers()
                .await
                .map(|t| t.clone())
                .unwrap_or_default();
            let grpc_status = trailers
                .get(&SimpleHeader::GRPC_STATUS)
                .and_then(|v| v.first())
                .map(|s| s.as_str());
            if grpc_status != Some("0") {
                return Err(format!("health check grpc-status {grpc_status:?}"));
            }
            Ok(decode_health_status(msg.as_deref().unwrap_or(b"")) == SERVING)
        }
        .await;
        *result_ref.lock().unwrap() = Some(outcome);
    };

    let work = futures_util::future::join_all(vec![
        sender.boxed(),
        receiver.boxed(),
        reader_task.map(|_| ()).boxed(),
        writer_task.map(|_| ()).boxed(),
    ])
    .boxed();
    let timer = async move { sleep_async(timeout).await }.boxed();

    match select(work, timer).await {
        Either::Left(_) => result
            .lock()
            .unwrap()
            .take()
            .unwrap_or_else(|| Err("gRPC health check produced no result".to_string())),
        Either::Right(_) => Err("gRPC health check timed out".to_string()),
    }
}

/// gRPC health probe: `true` = the backend reported `SERVING`.
///
/// The check is driven on a dedicated thread and bounded by a hard
/// [`std::sync::mpsc::Receiver::recv_timeout`]. `grpc_health_check` already races
/// its pump against an internal timer, but a stuck timer (or an in-flight probe
/// racing proxy shutdown) must not be allowed to hang this probe thread — a hung
/// probe would block [`HealthMonitor::join`] and stall `ProxyServer::shutdown`.
/// The inner future still owns the normal path (and, on timeout, drops the H2
/// pump so the connection is closed); the outer deadline is a backstop only.
fn probe_grpc_health(base_url: &str, timeout: Duration) -> bool {
    let url = join_url(base_url, "/grpc.health.v1.Health/Check");
    let (tx, rx) = std::sync::mpsc::channel();
    std::thread::spawn(move || {
        let _ = tx.send(block_on_future(grpc_health_check(url, timeout)));
    });
    // A margin over `timeout` so the inner timer + H2 open settle before the
    // outer deadline fires; the outer deadline only trips on a stuck inner timer.
    rx.recv_timeout(timeout + Duration::from_millis(250))
        .ok()
        .and_then(std::result::Result::ok)
        .unwrap_or(false)
}

/// Join a base URL and a health path, avoiding a doubled slash.
fn join_url(base_url: &str, path: &str) -> String {
    let base = base_url.strip_suffix('/').unwrap_or(base_url);
    if path.starts_with('/') {
        format!("{base}{path}")
    } else {
        format!("{base}/{path}")
    }
}

/// Sleep `total` in short slices, returning `false` if shutdown fired.
fn sleep_interruptible(total: Duration, shutdown: &Arc<OnSignal>) -> bool {
    let slice = Duration::from_millis(100);
    let start = Instant::now();
    while start.elapsed() < total {
        if shutdown.probe() {
            return false;
        }
        std::thread::sleep(slice.min(total.saturating_sub(start.elapsed())));
    }
    true
}

/// Minimal helper: resolve the first socket address for `host:port`.
trait FirstSocketAddr {
    fn to_socket_addrs_first(&self) -> Option<std::net::SocketAddr>;
}

impl FirstSocketAddr for str {
    fn to_socket_addrs_first(&self) -> Option<std::net::SocketAddr> {
        use std::net::ToSocketAddrs;
        self.to_socket_addrs().ok().and_then(|mut it| it.next())
    }
}