praxis-proxy-core 0.5.5

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2024 Praxis Contributors

//! Upstream cluster definitions: endpoints, load-balancing strategies, and timeouts.

mod endpoint;
mod health_check;
mod load_balancer_strategy;
mod retry_policy;

use std::sync::Arc;

pub use endpoint::Endpoint;
pub use health_check::{HealthCheckConfig, HealthCheckType};
pub use load_balancer_strategy::{
    ConsistentHashOpts, HashFunction, LoadBalancerStrategy, MaglevOpts, ParameterisedStrategy, PriorityOpts,
    RingHashOpts, SimpleStrategy, SubsetFallbackPolicy, SubsetOpts, ZoneAwareOpts,
};
pub use retry_policy::{
    BackoffConfig, BudgetPercent, DEFAULT_MAX_RETRIES, DEFAULT_RETRY_BODY_LIMIT_BYTES, HttpStatusCode,
    MAX_EFFECTIVE_RETRIES, MAX_RETRY_BODY_LIMIT_BYTES, RetriableCondition, RetryBodyLimit, RetryBudgetConfig,
    RetryPolicy,
};
use serde::{Deserialize, Serialize};

use crate::errors::ProxyError;

// -----------------------------------------------------------------------------
// Cluster
// -----------------------------------------------------------------------------

/// HTTP-specific options for a cluster.
///
/// These apply only when the cluster serves HTTP listeners; TCP load
/// balancers do not process HTTP headers and ignore this block.
#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct ClusterHttpOptions {
    /// Override the upstream HTTP `Host` header.
    ///
    /// When set, the proxy rewrites the `Host` header sent to the
    /// upstream instead of forwarding the downstream value. Upstream
    /// connections use HTTP/1.1; the `:authority` pseudo-header on
    /// the downstream HTTP/2 side is not forwarded upstream. TLS SNI
    /// remains independent — configure `tls.sni` separately when
    /// needed.
    ///
    /// Must be a valid HTTP authority: a hostname with an optional
    /// port, or a bracketed IPv6 address with an optional port. URI
    /// schemes, paths, userinfo, and fragments are rejected.
    #[serde(default)]
    pub authority: Option<Arc<str>>,

    /// Opaque application protocol the upstream cluster expects.
    ///
    /// Declares the wire representation of the request body (for
    /// example `openai_chat_completions` or `openai_responses`). The
    /// value stays opaque to Praxis core — consuming filters interpret
    /// it; Praxis defines no enum of known protocols. Validated as a
    /// bounded, canonical identifier: 1–64 bytes of lowercase ASCII
    /// letters, digits, `.`, `_`, or `-`, starting and ending with a
    /// letter or digit.
    ///
    /// The open string type is deliberate: the protocol set is
    /// open-ended and owned by consuming filters, so keep this a string
    /// — do not convert it to an enum.
    #[serde(default)]
    pub application_protocol: Option<Arc<str>>,

    /// Opaque application provider refining `application_protocol`.
    ///
    /// Distinguishes provider-specific semantics (for example `openai`
    /// or `vllm`) independent of the deployed cluster, whose identity
    /// is already the cluster name. Stays opaque to Praxis core and
    /// follows the same identifier rules — and the same enum-free
    /// rationale — as `application_protocol`.
    #[serde(default)]
    pub application_provider: Option<Arc<str>>,
}

/// A named group of upstream endpoints.
///
/// ```
/// # use praxis_core::config::Cluster;
/// let yaml = r#"
/// name: "backend"
/// endpoints: ["10.0.0.1:8080"]
/// connection_timeout_ms: 5000
/// idle_timeout_ms: 30000
/// "#;
/// let cluster: Cluster = serde_yaml::from_str(yaml).unwrap();
/// assert_eq!(cluster.connection_timeout_ms, Some(5000));
/// assert_eq!(cluster.idle_timeout_ms, Some(30000));
/// assert!(cluster.read_timeout_ms.is_none());
/// assert!(cluster.tls.is_none());
/// ```
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct Cluster {
    /// Unique name for the cluster.
    pub name: Arc<str>,

    /// HTTP-specific cluster options.
    ///
    /// Grouped under `http:` so the shared cluster/upstream transport
    /// types stay protocol-agnostic; TCP load balancers ignore this
    /// block entirely.
    ///
    /// ```
    /// # use praxis_core::config::Cluster;
    /// let yaml = r#"
    /// name: "api"
    /// endpoints: ["10.0.0.1:443"]
    /// http:
    ///   authority: "api.example.com"
    /// tls:
    ///   sni: "api.example.com"
    /// "#;
    /// let cluster: Cluster = serde_yaml::from_str(yaml).unwrap();
    /// assert_eq!(cluster.http.authority.as_deref(), Some("api.example.com"));
    /// ```
    #[serde(default)]
    pub http: ClusterHttpOptions,

    /// TCP connection timeout in milliseconds.
    ///
    /// Applies to the TCP handshake only (before TLS). When
    /// exceeded, the connection attempt fails and the load
    /// balancer may retry on the next endpoint. `None` (the
    /// default) uses Pingora's built-in timeout.
    #[serde(default)]
    pub connection_timeout_ms: Option<u64>,

    /// List of endpoints for the cluster. Each entry is either a plain
    /// `"host:port"` string or a `{ address, weight }` object.
    pub endpoints: Vec<Endpoint>,

    /// Active health check configuration for this cluster.
    #[serde(default)]
    pub health_check: Option<HealthCheckConfig>,

    /// Idle connection timeout in milliseconds.
    ///
    /// Closes pooled upstream connections that have been idle
    /// longer than this duration. `None` uses Pingora's default.
    #[serde(default)]
    pub idle_timeout_ms: Option<u64>,

    /// Load-balancing algorithm for this cluster. Defaults to `round_robin`.
    #[serde(default)]
    pub load_balancer_strategy: LoadBalancerStrategy,

    /// Maximum concurrent in-flight requests to this cluster.
    ///
    /// When set, excess requests receive 503. Prevents a single
    /// slow upstream from consuming all available capacity.
    ///
    /// ```
    /// # use praxis_core::config::Cluster;
    /// let yaml = r#"
    /// name: backend
    /// endpoints: ["10.0.0.1:80"]
    /// max_connections: 100
    /// "#;
    /// let cluster: Cluster = serde_yaml::from_str(yaml).unwrap();
    /// assert_eq!(cluster.max_connections, Some(100));
    /// ```
    #[serde(default)]
    pub max_connections: Option<u32>,

    /// Per-read timeout in milliseconds.
    ///
    /// Applies to each individual read operation on an
    /// established upstream connection. A timeout fires a 502
    /// response to the client. Use [`total_connection_timeout_ms`]
    /// to bound the entire exchange instead.
    ///
    /// [`total_connection_timeout_ms`]: Cluster::total_connection_timeout_ms
    #[serde(default)]
    pub read_timeout_ms: Option<u64>,

    /// TLS settings for upstream connections.
    ///
    /// Presence implies TLS is enabled. Omit for plaintext HTTP.
    #[serde(default)]
    pub tls: Option<praxis_tls::ClusterTls>,

    /// Total connection timeout in milliseconds (TCP + TLS).
    ///
    /// Bounds the combined TCP handshake and TLS negotiation.
    /// When exceeded, the connection attempt fails with a 502
    /// response. Prefer this over [`connection_timeout_ms`] for
    /// TLS-enabled clusters where the handshake dominates latency.
    ///
    /// [`connection_timeout_ms`]: Cluster::connection_timeout_ms
    #[serde(default)]
    pub total_connection_timeout_ms: Option<u64>,

    /// Per-write timeout in milliseconds.
    ///
    /// Applies to each individual write operation on an
    /// established upstream connection. A timeout fires a 502
    /// response to the client.
    #[serde(default)]
    pub write_timeout_ms: Option<u64>,

    /// Optional retry policy for this cluster.
    ///
    /// When unset, the proxy retains the legacy connect-failure
    /// retry behavior (3 attempts, idempotent methods, 64 `KiB` body).
    #[serde(default)]
    pub retry_policy: Option<RetryPolicy>,
}

impl Cluster {
    /// Validate the optional upstream HTTP authority override.
    ///
    /// # Errors
    ///
    /// Returns [`ProxyError::Config`] when the authority is not a
    /// supported hostname with an optional port or a bracketed IPv6
    /// address with an optional port.
    pub fn validate_authority(&self) -> Result<(), ProxyError> {
        let Some(authority) = self.http.authority.as_deref() else {
            return Ok(());
        };
        super::validate::cluster::validate_authority(authority, &self.name)
    }

    /// Build a cluster with only a name and endpoints; all other
    /// fields use their defaults (no timeouts, no TLS, no health
    /// check, `round_robin` strategy).
    ///
    /// ```
    /// use praxis_core::config::Cluster;
    /// use praxis_tls::ClusterTls;
    ///
    /// let c = Cluster {
    ///     tls: Some(ClusterTls::default()),
    ///     ..Cluster::with_defaults("backend", vec!["10.0.0.1:443".into()])
    /// };
    /// assert_eq!(&*c.name, "backend");
    /// assert!(c.tls.is_some());
    /// assert!(c.tls.as_ref().unwrap().verify);
    /// ```
    pub fn with_defaults(name: &str, endpoints: Vec<Endpoint>) -> Self {
        Self {
            name: Arc::from(name),
            http: ClusterHttpOptions::default(),
            connection_timeout_ms: None,
            endpoints,
            health_check: None,
            idle_timeout_ms: None,
            load_balancer_strategy: LoadBalancerStrategy::default(),
            max_connections: None,
            read_timeout_ms: None,
            tls: None,
            total_connection_timeout_ms: None,
            write_timeout_ms: None,
            retry_policy: None,
        }
    }
}

// -----------------------------------------------------------------------------
// Tests
// -----------------------------------------------------------------------------

#[cfg(test)]
#[expect(clippy::allow_attributes, reason = "blanket test suppressions")]
#[allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::indexing_slicing,
    clippy::needless_raw_strings,
    clippy::needless_raw_string_hashes,
    reason = "tests use unwrap/expect/indexing/raw strings for brevity"
)]
mod tests {
    use super::*;

    #[test]
    fn parse_cluster_minimal() {
        let yaml = r#"
name: "backend"
endpoints: ["10.0.0.1:8080"]
"#;
        let cluster: Cluster = serde_yaml::from_str(yaml).unwrap();
        assert_eq!(&*cluster.name, "backend", "cluster name mismatch");
        assert_eq!(
            cluster.endpoints[0].address(),
            "10.0.0.1:8080",
            "endpoint address mismatch"
        );
        assert_eq!(cluster.endpoints[0].weight(), 1, "default weight should be 1");
        assert_eq!(
            cluster.load_balancer_strategy,
            LoadBalancerStrategy::default(),
            "strategy should default"
        );
        assert!(
            cluster.connection_timeout_ms.is_none(),
            "connection_timeout should default to None"
        );
    }

    #[test]
    fn parse_cluster_with_weights() {
        let yaml = r#"
name: "backend"
endpoints:
  - "10.0.0.1:8080"
  - address: "10.0.0.2:8080"
    weight: 3
"#;
        let cluster: Cluster = serde_yaml::from_str(yaml).unwrap();
        assert_eq!(cluster.endpoints.len(), 2, "should parse two endpoints");
        assert_eq!(cluster.endpoints[0].weight(), 1, "simple endpoint weight should be 1");
        assert_eq!(cluster.endpoints[1].weight(), 3, "weighted endpoint weight should be 3");
    }

    #[test]
    fn parse_cluster_with_timeouts() {
        let yaml = r#"
name: "backend"
endpoints: ["10.0.0.1:8080"]
connection_timeout_ms: 5000
idle_timeout_ms: 30000
read_timeout_ms: 10000
write_timeout_ms: 10000
"#;
        let cluster: Cluster = serde_yaml::from_str(yaml).unwrap();
        assert_eq!(
            cluster.connection_timeout_ms,
            Some(5000),
            "connection_timeout_ms mismatch"
        );
        assert_eq!(cluster.idle_timeout_ms, Some(30000), "idle_timeout_ms mismatch");
        assert_eq!(cluster.read_timeout_ms, Some(10000), "read_timeout_ms mismatch");
        assert_eq!(cluster.write_timeout_ms, Some(10000), "write_timeout_ms mismatch");
    }

    #[test]
    fn cluster_roundtrips_via_serde() {
        let cluster = Cluster {
            connection_timeout_ms: Some(1000),
            ..Cluster::with_defaults("web", vec!["10.0.0.1:80".into()])
        };
        let value = serde_yaml::to_value(&cluster).unwrap();
        let back: Cluster = serde_yaml::from_value(value).unwrap();
        assert_eq!(back.name, cluster.name, "name should roundtrip");
        assert_eq!(back.endpoints, cluster.endpoints, "endpoints should roundtrip");
        assert_eq!(
            back.connection_timeout_ms, cluster.connection_timeout_ms,
            "timeout should roundtrip"
        );
    }

    #[test]
    fn tls_and_sni_parse_correctly() {
        let yaml = r#"
name: "backend"
endpoints: ["10.0.0.1:443"]
tls:
  sni: "api.example.com"
"#;
        let cluster: Cluster = serde_yaml::from_str(yaml).unwrap();
        assert!(cluster.tls.is_some(), "tls should be present");
        assert_eq!(
            cluster.tls.as_ref().unwrap().sni.as_deref(),
            Some("api.example.com"),
            "sni mismatch"
        );
    }

    #[test]
    fn tls_verify_defaults_to_true() {
        let yaml = r#"
name: "backend"
endpoints: ["10.0.0.1:443"]
tls: {}
"#;
        let cluster: Cluster = serde_yaml::from_str(yaml).unwrap();
        assert!(cluster.tls.as_ref().unwrap().verify, "verify should default to true");
    }

    #[test]
    fn tls_verify_can_be_disabled() {
        let yaml = r#"
name: "backend"
endpoints: ["10.0.0.1:443"]
tls:
  verify: false
"#;
        let cluster: Cluster = serde_yaml::from_str(yaml).unwrap();
        assert!(
            !cluster.tls.as_ref().unwrap().verify,
            "verify should be false when explicitly set"
        );
    }

    #[test]
    fn no_tls_by_default() {
        let cluster = Cluster::with_defaults("web", vec!["10.0.0.1:80".into()]);
        assert!(cluster.tls.is_none(), "tls should be None by default");
    }

    #[test]
    fn parse_cluster_application_metadata() {
        let yaml = r#"
name: "backend"
endpoints: ["10.0.0.1:8080"]
http:
  application_protocol: openai_chat_completions
  application_provider: vllm
"#;
        let cluster: Cluster = serde_yaml::from_str(yaml).unwrap();
        assert_eq!(
            cluster.http.application_protocol.as_deref(),
            Some("openai_chat_completions"),
            "application_protocol mismatch"
        );
        assert_eq!(
            cluster.http.application_provider.as_deref(),
            Some("vllm"),
            "application_provider mismatch"
        );
    }

    #[test]
    fn application_metadata_defaults_to_none() {
        let cluster = Cluster::with_defaults("web", vec!["10.0.0.1:80".into()]);
        assert!(
            cluster.http.application_protocol.is_none(),
            "application_protocol should default to None"
        );
        assert!(
            cluster.http.application_provider.is_none(),
            "application_provider should default to None"
        );
    }
}