zlayer_proxy/stream/

registry.rs

//! Stream service registry for L4 routing
//!
//! Maps listen ports to backend services for TCP and UDP proxying.
//! Includes health-aware backend selection: unhealthy backends are
//! skipped during round-robin selection, with a fallback to any
//! backend if all are marked unhealthy.

use dashmap::DashMap;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;

use super::config::{StreamHealthProbe, StreamProxyConfig};

/// Health state of a stream backend
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BackendHealth {
    /// Backend is reachable and accepting connections
    Healthy,
    /// Backend failed the last health probe
    Unhealthy,
    /// Health has not yet been determined (treated as healthy)
    Unknown,
}

impl BackendHealth {
    /// Returns `true` if the backend should be considered usable.
    #[must_use]
    pub fn is_usable(self) -> bool {
        matches!(self, BackendHealth::Healthy | BackendHealth::Unknown)
    }
}

/// A resolved stream service with backend addresses and health state
#[derive(Clone, Debug)]
pub struct StreamService {
    /// Service name (for logging/metrics)
    pub name: String,
    /// Backend addresses for load balancing
    pub backends: Vec<SocketAddr>,
    /// Per-backend health state
    health: Arc<RwLock<HashMap<SocketAddr, BackendHealth>>>,
    /// Round-robin index for backend selection
    rr_index: Arc<AtomicUsize>,
    /// Runtime L4 config (TLS / proxy-protocol / session-timeout / health probe)
    /// translated from the endpoint's `stream:` block. Drives the health
    /// checker's probe selection (see [`StreamRegistry::spawn_health_checker`]).
    pub config: StreamProxyConfig,
}

impl StreamService {
    /// Create a new stream service
    #[must_use]
    pub fn new(name: String, backends: Vec<SocketAddr>) -> Self {
        let health: HashMap<SocketAddr, BackendHealth> = backends
            .iter()
            .map(|addr| (*addr, BackendHealth::Unknown))
            .collect();
        Self {
            name,
            backends,
            health: Arc::new(RwLock::new(health)),
            rr_index: Arc::new(AtomicUsize::new(0)),
            config: StreamProxyConfig::default(),
        }
    }

    /// Attach a runtime [`StreamProxyConfig`] to this service.
    ///
    /// Builder-style; preserves [`StreamService::new`]'s 2-arg arity so the
    /// existing call sites keep compiling. The config's `health_check` drives
    /// what the background health checker probes.
    #[must_use]
    pub fn with_config(mut self, config: StreamProxyConfig) -> Self {
        self.config = config;
        self
    }

    /// Select next backend using round-robin, skipping unhealthy backends.
    ///
    /// Tries up to `backends.len()` candidates. If all backends are unhealthy,
    /// falls back to returning *any* backend (better than nothing).
    #[must_use]
    pub fn select_backend(&self) -> Option<SocketAddr> {
        if self.backends.is_empty() {
            return None;
        }

        let len = self.backends.len();
        let start = self.rr_index.fetch_add(1, Ordering::Relaxed);

        // Try to read health state without blocking; if the lock is held,
        // just fall through to simple round-robin.
        let health_guard = self.health.try_read();

        if let Ok(health) = health_guard {
            // First pass: find a healthy backend
            for i in 0..len {
                let idx = (start + i) % len;
                let addr = self.backends[idx];
                let status = health.get(&addr).copied().unwrap_or(BackendHealth::Unknown);
                if status.is_usable() {
                    return Some(addr);
                }
            }
        }

        // Fallback: all unhealthy or lock contention — use simple round-robin
        Some(self.backends[start % len])
    }

    /// Update backend addresses (for scaling events).
    ///
    /// New backends start with `Unknown` health; removed backends are pruned
    /// from the health map.
    pub fn update_backends(&mut self, backends: Vec<SocketAddr>) {
        // We need to block here since this is called from a &mut self context
        // (inside DashMap::get_mut), so we can use blocking write.
        let mut health = self
            .health
            .try_write()
            .unwrap_or_else(|_| {
                // In the extremely unlikely case of write contention, just proceed
                // with a fresh health map.
                tracing::warn!(service = %self.name, "Health map write contention during backend update");
                // This should never actually happen since update_backends holds &mut self
                unreachable!("update_backends requires exclusive access")
            });

        // Add new backends with Unknown health
        for addr in &backends {
            health.entry(*addr).or_insert(BackendHealth::Unknown);
        }

        // Remove backends that are no longer present
        let backend_set: std::collections::HashSet<SocketAddr> = backends.iter().copied().collect();
        health.retain(|addr, _| backend_set.contains(addr));

        self.backends = backends;
    }

    /// Set the health status of a specific backend
    pub async fn set_backend_health(&self, addr: SocketAddr, status: BackendHealth) {
        let mut health = self.health.write().await;
        if let Some(h) = health.get_mut(&addr) {
            *h = status;
        }
    }

    /// Get the health status of a specific backend
    pub async fn get_backend_health(&self, addr: SocketAddr) -> BackendHealth {
        let health = self.health.read().await;
        health.get(&addr).copied().unwrap_or(BackendHealth::Unknown)
    }

    /// Get current backend count
    #[must_use]
    pub fn backend_count(&self) -> usize {
        self.backends.len()
    }

    /// Get count of healthy (usable) backends
    pub async fn healthy_count(&self) -> usize {
        let health = self.health.read().await;
        self.backends
            .iter()
            .filter(|addr| {
                health
                    .get(addr)
                    .copied()
                    .unwrap_or(BackendHealth::Unknown)
                    .is_usable()
            })
            .count()
    }
}

/// Registry for L4 stream services
///
/// Maps listen ports to services for both TCP and UDP protocols.
#[derive(Default)]
pub struct StreamRegistry {
    /// TCP services by listen port
    tcp_services: DashMap<u16, StreamService>,
    /// UDP services by listen port
    udp_services: DashMap<u16, StreamService>,
}

impl StreamRegistry {
    /// Create a new empty registry
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a TCP service for a port
    pub fn register_tcp(&self, port: u16, service: StreamService) {
        tracing::debug!(
            port = port,
            service = %service.name,
            backends = service.backend_count(),
            "Registered TCP stream service"
        );
        self.tcp_services.insert(port, service);
    }

    /// Register a UDP service for a port
    pub fn register_udp(&self, port: u16, service: StreamService) {
        tracing::debug!(
            port = port,
            service = %service.name,
            backends = service.backend_count(),
            "Registered UDP stream service"
        );
        self.udp_services.insert(port, service);
    }

    /// Resolve TCP service for a port
    #[must_use]
    pub fn resolve_tcp(&self, port: u16) -> Option<StreamService> {
        self.tcp_services.get(&port).map(|s| s.clone())
    }

    /// Resolve UDP service for a port
    #[must_use]
    pub fn resolve_udp(&self, port: u16) -> Option<StreamService> {
        self.udp_services.get(&port).map(|s| s.clone())
    }

    /// Apply a runtime [`StreamProxyConfig`] to the TCP service on `port`.
    ///
    /// No-op when no TCP service is registered for that port. Used by the agent
    /// to attach the endpoint's translated `stream:` settings (notably the
    /// health probe) to a service whose backends were registered out-of-band.
    pub fn set_tcp_config(&self, port: u16, config: StreamProxyConfig) {
        if let Some(mut service) = self.tcp_services.get_mut(&port) {
            service.config = config;
        }
    }

    /// Apply a runtime [`StreamProxyConfig`] to the UDP service on `port`.
    ///
    /// No-op when no UDP service is registered for that port.
    pub fn set_udp_config(&self, port: u16, config: StreamProxyConfig) {
        if let Some(mut service) = self.udp_services.get_mut(&port) {
            service.config = config;
        }
    }

    /// Update backends for a TCP service
    pub fn update_tcp_backends(&self, port: u16, backends: Vec<SocketAddr>) {
        if let Some(mut service) = self.tcp_services.get_mut(&port) {
            tracing::debug!(
                port = port,
                service = %service.name,
                old_count = service.backend_count(),
                new_count = backends.len(),
                "Updating TCP backends"
            );
            service.update_backends(backends);
        }
    }

    /// Update backends for a UDP service
    pub fn update_udp_backends(&self, port: u16, backends: Vec<SocketAddr>) {
        if let Some(mut service) = self.udp_services.get_mut(&port) {
            tracing::debug!(
                port = port,
                service = %service.name,
                old_count = service.backend_count(),
                new_count = backends.len(),
                "Updating UDP backends"
            );
            service.update_backends(backends);
        }
    }

    /// Remove a TCP service
    #[must_use]
    pub fn unregister_tcp(&self, port: u16) -> Option<StreamService> {
        self.tcp_services.remove(&port).map(|(_, s)| s)
    }

    /// Remove a UDP service
    #[must_use]
    pub fn unregister_udp(&self, port: u16) -> Option<StreamService> {
        self.udp_services.remove(&port).map(|(_, s)| s)
    }

    /// Get count of registered TCP services
    #[must_use]
    pub fn tcp_count(&self) -> usize {
        self.tcp_services.len()
    }

    /// Get count of registered UDP services
    #[must_use]
    pub fn udp_count(&self) -> usize {
        self.udp_services.len()
    }

    /// List all registered TCP ports
    #[must_use]
    pub fn tcp_ports(&self) -> Vec<u16> {
        self.tcp_services.iter().map(|e| *e.key()).collect()
    }

    /// List all registered UDP ports
    #[must_use]
    pub fn udp_ports(&self) -> Vec<u16> {
        self.udp_services.iter().map(|e| *e.key()).collect()
    }

    /// List all registered TCP services with their listen ports.
    #[must_use]
    pub fn list_tcp_services(&self) -> Vec<(u16, StreamService)> {
        self.tcp_services
            .iter()
            .map(|e| (*e.key(), e.value().clone()))
            .collect()
    }

    /// List all registered UDP services with their listen ports.
    #[must_use]
    pub fn list_udp_services(&self) -> Vec<(u16, StreamService)> {
        self.udp_services
            .iter()
            .map(|e| (*e.key(), e.value().clone()))
            .collect()
    }

    /// Spawn a background health checker that periodically probes registered
    /// stream backends.
    ///
    /// TCP services are probed with a connect-only check (matching the
    /// `TcpConnect` health-check type — the only supported TCP probe).
    ///
    /// UDP services are probed only when their [`StreamProxyConfig::health_check`]
    /// is `Some(StreamHealthProbe::UdpProbe { .. })`: the configured request is
    /// sent to each backend and the backend is marked `Healthy` iff a reply
    /// arrives (and matches `expect` by byte-substring when set). UDP services
    /// without a configured probe remain `Unknown` (always usable) — preserving
    /// the previous "never probe UDP" behavior.
    ///
    /// The task runs every `interval` and uses `timeout` for each probe.
    /// Returns a `JoinHandle` that can be used to cancel the checker.
    #[must_use]
    pub fn spawn_health_checker(
        self: &Arc<Self>,
        interval: Duration,
        timeout: Duration,
    ) -> tokio::task::JoinHandle<()> {
        let registry = Arc::clone(self);

        tokio::spawn(async move {
            let mut ticker = tokio::time::interval(interval);
            // Skip the first immediate tick
            ticker.tick().await;

            loop {
                ticker.tick().await;

                // Iterate all TCP services and probe each backend
                for entry in &registry.tcp_services {
                    let service = entry.value().clone();
                    let backends = service.backends.clone();

                    for addr in backends {
                        let svc = service.clone();
                        let probe_timeout = timeout;

                        // Probe each backend concurrently
                        tokio::spawn(async move {
                            let result = tokio::time::timeout(
                                probe_timeout,
                                tokio::net::TcpStream::connect(addr),
                            )
                            .await;

                            let health = match result {
                                Ok(Ok(_stream)) => BackendHealth::Healthy,
                                Ok(Err(e)) => {
                                    tracing::debug!(
                                        service = %svc.name,
                                        backend = %addr,
                                        error = %e,
                                        "TCP health check failed (connect error)"
                                    );
                                    BackendHealth::Unhealthy
                                }
                                Err(_) => {
                                    tracing::debug!(
                                        service = %svc.name,
                                        backend = %addr,
                                        "TCP health check failed (timeout)"
                                    );
                                    BackendHealth::Unhealthy
                                }
                            };

                            svc.set_backend_health(addr, health).await;
                        });
                    }
                }

                // Iterate all UDP services and probe each backend when a UDP
                // probe is configured. Services without one are left untouched.
                for entry in &registry.udp_services {
                    let service = entry.value().clone();
                    let Some(StreamHealthProbe::UdpProbe { request, expect }) =
                        service.config.health_check.clone()
                    else {
                        continue;
                    };
                    let backends = service.backends.clone();

                    for addr in backends {
                        let svc = service.clone();
                        let probe_timeout = timeout;
                        let request = request.clone();
                        let expect = expect.clone();

                        tokio::spawn(async move {
                            let health = match probe_udp_backend(
                                addr,
                                &request,
                                expect.as_deref(),
                                probe_timeout,
                            )
                            .await
                            {
                                Ok(true) => BackendHealth::Healthy,
                                Ok(false) => {
                                    tracing::debug!(
                                        service = %svc.name,
                                        backend = %addr,
                                        "UDP health check failed (reply did not match expect)"
                                    );
                                    BackendHealth::Unhealthy
                                }
                                Err(e) => {
                                    tracing::debug!(
                                        service = %svc.name,
                                        backend = %addr,
                                        error = %e,
                                        "UDP health check failed"
                                    );
                                    BackendHealth::Unhealthy
                                }
                            };

                            svc.set_backend_health(addr, health).await;
                        });
                    }
                }
            }
        })
    }
}

/// Probe a single UDP backend by sending `request` and waiting (up to
/// `timeout`) for any reply.
///
/// Returns `Ok(true)` when a reply arrives that satisfies `expect` (any reply
/// when `expect` is `None`; a reply containing `expect` as a byte-substring
/// otherwise). Returns `Ok(false)` when a reply arrives but does not contain
/// `expect`. Returns `Err(..)` on socket errors or recv timeout.
///
/// # Errors
///
/// Returns an error if binding/connecting the probe socket fails, if the send
/// fails, or if no reply arrives before `timeout` elapses.
pub async fn probe_udp_backend(
    addr: SocketAddr,
    request: &[u8],
    expect: Option<&[u8]>,
    timeout: Duration,
) -> std::result::Result<bool, std::io::Error> {
    let socket = tokio::net::UdpSocket::bind("0.0.0.0:0").await?;
    socket.connect(addr).await?;
    socket.send(request).await?;

    let mut buf = vec![0u8; 65535];
    let len = tokio::time::timeout(timeout, socket.recv(&mut buf))
        .await
        .map_err(|_| {
            std::io::Error::new(std::io::ErrorKind::TimedOut, "UDP health probe timed out")
        })??;

    let reply = &buf[..len];
    match expect {
        Some(pat) => Ok(byte_contains(reply, pat)),
        None => Ok(true),
    }
}

/// `true` iff `haystack` contains `needle` as a contiguous byte substring.
/// An empty needle always matches.
#[must_use]
fn byte_contains(haystack: &[u8], needle: &[u8]) -> bool {
    if needle.is_empty() {
        return true;
    }
    if needle.len() > haystack.len() {
        return false;
    }
    haystack.windows(needle.len()).any(|w| w == needle)
}

#[cfg(test)]
mod health_probe_tests {
    use super::*;
    use std::time::Duration;
    use tokio::net::UdpSocket;

    #[test]
    fn byte_contains_matches() {
        assert!(byte_contains(b"hello world", b"world"));
        assert!(byte_contains(b"hello world", b"hello"));
        assert!(byte_contains(b"\xFF\x00\xAB", b"\x00\xAB"));
        assert!(byte_contains(b"anything", b"")); // empty needle always matches
    }

    #[test]
    fn byte_contains_rejects() {
        assert!(!byte_contains(b"hello", b"world"));
        assert!(!byte_contains(b"abc", b"abcd")); // needle longer than haystack
        assert!(!byte_contains(b"", b"x"));
    }

    #[tokio::test]
    async fn udp_probe_healthy_against_echo() {
        // Spawn a tiny echo server.
        let echo = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let echo_addr = echo.local_addr().unwrap();
        tokio::spawn(async move {
            let mut buf = vec![0u8; 1500];
            if let Ok((n, peer)) = echo.recv_from(&mut buf).await {
                let _ = echo.send_to(&buf[..n], peer).await;
            }
        });

        // Any reply (no expect) -> healthy.
        let ok = probe_udp_backend(echo_addr, b"ping", None, Duration::from_secs(2))
            .await
            .unwrap();
        assert!(ok, "echo reply with no expect must be healthy");
    }

    #[tokio::test]
    async fn udp_probe_expect_substring() {
        let echo = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let echo_addr = echo.local_addr().unwrap();
        tokio::spawn(async move {
            let mut buf = vec![0u8; 1500];
            for _ in 0..2 {
                if let Ok((n, peer)) = echo.recv_from(&mut buf).await {
                    let _ = echo.send_to(&buf[..n], peer).await;
                }
            }
        });

        // Reply contains expect substring -> healthy.
        let ok = probe_udp_backend(
            echo_addr,
            b"PONG-token",
            Some(b"token"),
            Duration::from_secs(2),
        )
        .await
        .unwrap();
        assert!(ok, "reply containing expect substring must be healthy");

        // Reply does NOT contain expect -> unhealthy (Ok(false)).
        let not_matched =
            probe_udp_backend(echo_addr, b"abc", Some(b"zzz"), Duration::from_secs(2))
                .await
                .unwrap();
        assert!(
            !not_matched,
            "reply missing expect substring must be unhealthy"
        );
    }

    #[tokio::test]
    async fn udp_probe_dead_port_times_out() {
        // Bind a socket to grab a free port, then drop it so nothing listens.
        let dead = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let dead_addr = dead.local_addr().unwrap();
        drop(dead);

        // No listener -> no reply -> timeout error (treated as Unhealthy by caller).
        let res = probe_udp_backend(dead_addr, b"ping", None, Duration::from_millis(300)).await;
        assert!(res.is_err(), "probe to dead UDP port must error (timeout)");
    }
}