praxis-proxy-filter 0.7.5

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

//! Weighted random endpoint selection.

use std::sync::{Arc, atomic::AtomicU64};

use praxis_core::health::ClusterHealthState;

use super::{
    endpoint::WeightedEndpoint,
    is_excluded,
    slow_start::{self, EffectiveWeights},
};

// -----------------------------------------------------------------------------
// Random
// -----------------------------------------------------------------------------

/// Uniform random endpoint selection, weighted by endpoint weight.
///
/// Each endpoint's probability of selection is proportional to its weight
/// relative to the total weight of all (healthy) endpoints. With equal
/// weights this reduces to uniform random selection.
pub(crate) struct Random {
    /// Deduplicated endpoint list with weights.
    endpoints: Vec<WeightedEndpoint>,

    /// Sum of all endpoint weights.
    total_weight: usize,

    /// Deterministic RNG state.
    rng: AtomicU64,
}

impl Random {
    /// Create a random selector from a deduplicated weighted endpoint list.
    pub(crate) fn new(endpoints: Vec<WeightedEndpoint>) -> Self {
        let total_weight: usize = endpoints.iter().map(|ep| ep.weight as usize).sum();
        Self {
            endpoints,
            total_weight,
            rng: AtomicU64::new(1),
        }
    }

    /// Return a randomly selected healthy endpoint address.
    ///
    /// Selection probability is proportional to endpoint weight. Falls back
    /// to the first healthy endpoint when all healthy candidates have zero
    /// weight, or to all endpoints (panic mode) when none are healthy.
    #[inline]
    pub(crate) fn select(&self, health: Option<&ClusterHealthState>, exclude: &[Arc<str>]) -> Option<Arc<str>> {
        self.select_weighted(health, exclude, None)
    }

    /// Pick an endpoint, using scaled ramp weights when `weights` is scaling.
    pub(crate) fn select_weighted(
        &self,
        health: Option<&ClusterHealthState>,
        exclude: &[Arc<str>],
        weights: Option<&EffectiveWeights>,
    ) -> Option<Arc<str>> {
        if self.total_weight == 0 {
            return None;
        }
        if let Some(weights) = slow_start::scaling_weights(weights) {
            return self.select_scaled(health, exclude, weights);
        }

        // Allocation-free two-pass walk (the round-robin shape): pass 1
        // sums candidate weight, pass 2 maps the drawn slot to its
        // cumulative bucket. Collecting candidate SmallVecs heap-allocated
        // twice per request past 8 endpoints.
        if let Some(state) = health {
            let healthy =
                |ep: &WeightedEndpoint| state.is_address_healthy(&ep.address) && !is_excluded(&ep.address, exclude);
            let (first, total) = survey(&self.endpoints, healthy);
            if let Some(first) = first {
                if total > 0 {
                    return pick_where(&self.endpoints, healthy, super::next_random(&self.rng), total);
                }
                return Some(Arc::clone(&first.address));
            }
        }

        let unexcluded = |ep: &WeightedEndpoint| !is_excluded(&ep.address, exclude);
        let (_, total) = survey(&self.endpoints, unexcluded);
        if total == 0 {
            return None;
        }
        pick_where(&self.endpoints, unexcluded, super::next_random(&self.rng), total)
    }

    /// Weighted random over scaled ramp units. A zero unit is not a candidate.
    fn select_scaled(
        &self,
        health: Option<&ClusterHealthState>,
        exclude: &[Arc<str>],
        weights: &EffectiveWeights,
    ) -> Option<Arc<str>> {
        let units = |endpoint: &WeightedEndpoint| {
            slow_start::endpoint_weight(&endpoint.address, endpoint.weight, Some(weights))
        };
        let eligible = |endpoint: &WeightedEndpoint| {
            slow_start::endpoint_weight(&endpoint.address, endpoint.weight, Some(weights)) > 0
                && !is_excluded(&endpoint.address, exclude)
                && health_allows(health, endpoint)
        };
        let (first, total) = survey_units(&self.endpoints, &eligible, &units);
        if total > 0 {
            return pick_units(&self.endpoints, &eligible, &units, super::next_random(&self.rng), total);
        }
        // No positive ramp weight in the healthy set. If nobody is healthy,
        // panic mode still has to serve, using configured weights.
        if health.is_some_and(|state| {
            self.endpoints
                .iter()
                .any(|endpoint| state.is_address_healthy(&endpoint.address))
        }) {
            return first.map(|endpoint| Arc::clone(&endpoint.address));
        }
        self.select(health, exclude)
    }
}

/// Whether `endpoint` passes the health filter. No health state means every endpoint is eligible.
fn health_allows(health: Option<&ClusterHealthState>, endpoint: &WeightedEndpoint) -> bool {
    health.is_none_or(|state| state.is_address_healthy(&endpoint.address))
}

/// First candidate and total candidate weight in one pass.
fn survey(
    endpoints: &[WeightedEndpoint],
    candidate: impl Fn(&WeightedEndpoint) -> bool,
) -> (Option<&WeightedEndpoint>, usize) {
    let mut first = None;
    let mut total = 0_usize;
    for ep in endpoints {
        if candidate(ep) {
            if first.is_none() {
                first = Some(ep);
            }
            total += ep.weight as usize;
        }
    }
    (first, total)
}

/// First candidate and total scaled weight in one pass.
fn survey_units<'a>(
    endpoints: &'a [WeightedEndpoint],
    candidate: &impl Fn(&WeightedEndpoint) -> bool,
    units: &impl Fn(&WeightedEndpoint) -> u64,
) -> (Option<&'a WeightedEndpoint>, u64) {
    let mut first = None;
    let mut total = 0_u64;
    for endpoint in endpoints {
        if candidate(endpoint) {
            if first.is_none() {
                first = Some(endpoint);
            }
            total = total.saturating_add(units(endpoint));
        }
    }
    (first, total)
}

/// Map `random` onto a scaled-weight bucket.
fn pick_units(
    endpoints: &[WeightedEndpoint],
    candidate: &impl Fn(&WeightedEndpoint) -> bool,
    units: &impl Fn(&WeightedEndpoint) -> u64,
    random: u64,
    total_weight: u64,
) -> Option<Arc<str>> {
    let slot = random % total_weight;
    let mut cumulative = 0_u64;
    let mut last = None;
    for endpoint in endpoints {
        if !candidate(endpoint) {
            continue;
        }
        cumulative = cumulative.saturating_add(units(endpoint));
        if slot < cumulative {
            return Some(Arc::clone(&endpoint.address));
        }
        last = Some(endpoint);
    }
    last.map(|endpoint| Arc::clone(&endpoint.address))
}

// -----------------------------------------------------------------------------
// Utilities
// -----------------------------------------------------------------------------

/// Map a random value to a candidate endpoint via cumulative weight
/// buckets, walking only endpoints that pass `candidate`.
///
/// The candidate subsequence keeps endpoint order, so a given draw always
/// lands in the same bucket.
#[expect(clippy::cast_possible_truncation, reason = "modulo total_weight bounds the result")]
fn pick_where(
    endpoints: &[WeightedEndpoint],
    candidate: impl Fn(&WeightedEndpoint) -> bool,
    random: u64,
    total_weight: usize,
) -> Option<Arc<str>> {
    let slot = (random as usize) % total_weight;
    let mut cumulative = 0_usize;
    let mut last = None;
    for ep in endpoints {
        if !candidate(ep) {
            continue;
        }
        cumulative += ep.weight as usize;
        if slot < cumulative {
            return Some(Arc::clone(&ep.address));
        }
        last = Some(ep);
    }
    last.map(|ep| Arc::clone(&ep.address))
}

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

#[cfg(test)]
#[expect(clippy::allow_attributes, reason = "blanket test suppressions")]
#[allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::indexing_slicing,
    clippy::panic,
    clippy::too_many_lines,
    reason = "tests"
)]
mod tests {
    use praxis_core::health::{ClusterHealthEntry, EndpointHealth};

    use super::*;

    #[test]
    fn scaled_random_skips_a_zero_ramp() {
        let random = Random::new(vec![ep("10.0.0.1:80", 1), ep("10.0.0.2:80", 1)]);
        let weights = EffectiveWeights::from_scaled(&[("10.0.0.1:80", 1_000_000), ("10.0.0.2:80", 0)]);
        for _ in 0..16 {
            let picked = random
                .select_weighted(None, &[], Some(&weights))
                .expect("the settled endpoint");
            assert_eq!(&*picked, "10.0.0.1:80", "an endpoint at zero stays out of the draw");
        }
    }

    #[test]
    fn single_endpoint_always_selected() {
        let r = Random::new(vec![ep("10.0.0.1:80", 1)]);
        for _ in 0..10 {
            assert_eq!(
                &*r.select(None, &[]).unwrap(),
                "10.0.0.1:80",
                "single endpoint must always be returned"
            );
        }
    }

    #[test]
    fn distributes_across_endpoints() {
        let r = Random::new(vec![ep("10.0.0.1:80", 1), ep("10.0.0.2:80", 1), ep("10.0.0.3:80", 1)]);

        let mut counts = std::collections::HashMap::new();
        for _ in 0..300 {
            *counts.entry(r.select(None, &[]).unwrap()).or_insert(0_u32) += 1;
        }

        assert_eq!(counts.len(), 3, "random should use all 3 endpoints");
        for (addr, count) in &counts {
            assert!((50..=200).contains(count), "expected ~100 for {addr}, got {count}");
        }
    }

    #[test]
    fn two_equal_endpoints_do_not_strictly_alternate() {
        let r = Random::new(vec![ep("10.0.0.1:80", 1), ep("10.0.0.2:80", 1)]);

        let picks: Vec<Arc<str>> = std::iter::repeat_with(|| r.select(None, &[]).unwrap())
            .take(8)
            .collect();

        assert!(
            picks.windows(2).any(|pair| pair[0] == pair[1]),
            "the raw LCG low bit alternates A,B,A,B; mixed output must repeat an endpoint within 8 draws: {picks:?}"
        );
    }

    #[test]
    fn weighted_bias() {
        let r = Random::new(vec![ep("10.0.0.1:80", 1), ep("10.0.0.2:80", 9)]);

        let mut counts = std::collections::HashMap::new();
        for _ in 0..1000 {
            *counts.entry(r.select(None, &[]).unwrap()).or_insert(0_u32) += 1;
        }

        let heavy = counts.get("10.0.0.2:80").copied().unwrap_or(0);
        assert!(
            heavy > 700,
            "weight-9 endpoint should get ~90% of traffic: heavy={heavy}"
        );
    }

    #[test]
    fn skips_unhealthy() {
        let r = Random::new(vec![ep("10.0.0.1:80", 1), ep("10.0.0.2:80", 1)]);
        let state = health_state(2);
        state.endpoints()[0].mark_unhealthy();

        for _ in 0..10 {
            assert_eq!(
                &*r.select(Some(&state), &[]).unwrap(),
                "10.0.0.2:80",
                "should skip unhealthy endpoint"
            );
        }
    }

    #[test]
    fn panic_mode_when_all_unhealthy() {
        let r = Random::new(vec![ep("10.0.0.1:80", 1), ep("10.0.0.2:80", 1)]);
        let state = health_state(2);
        state.endpoints()[0].mark_unhealthy();
        state.endpoints()[1].mark_unhealthy();

        let addr = r.select(Some(&state), &[]).unwrap();
        assert!(
            &*addr == "10.0.0.1:80" || &*addr == "10.0.0.2:80",
            "panic mode should still return an endpoint"
        );
    }

    #[test]
    fn empty_endpoints_returns_none() {
        let r = Random::new(vec![]);
        assert!(r.select(None, &[]).is_none(), "empty endpoint list should return None");
    }

    #[test]
    fn empty_endpoints_with_health_returns_none() {
        let r = Random::new(vec![]);
        let state: ClusterHealthState = Arc::new(ClusterHealthEntry::new(vec![], vec![], None, None));
        assert!(
            r.select(Some(&state), &[]).is_none(),
            "empty endpoint list with health state should return None"
        );
    }

    #[test]
    fn all_zero_weight_returns_none() {
        let r = Random::new(vec![ep("10.0.0.1:80", 0), ep("10.0.0.2:80", 0)]);
        assert!(
            r.select(None, &[]).is_none(),
            "all-zero-weight endpoints should return None"
        );
    }

    #[test]
    fn zero_weight_healthy_returns_first_healthy() {
        let r = Random::new(vec![ep("10.0.0.1:80", 0), ep("10.0.0.2:80", 5)]);
        let state = health_state(2);
        state.endpoints()[1].mark_unhealthy();

        let addr = r.select(Some(&state), &[]).unwrap();
        assert_eq!(
            &*addr, "10.0.0.1:80",
            "should return first healthy endpoint when healthy candidates have zero total weight"
        );
    }

    #[test]
    fn pick_exact_bucket_boundaries() {
        let endpoints = vec![ep("A", 1), ep("B", 3), ep("C", 1)];
        let all = |_: &WeightedEndpoint| true;
        assert_eq!(&*pick_where(&endpoints, all, 0, 5).unwrap(), "A", "slot 0 → A");
        assert_eq!(&*pick_where(&endpoints, all, 1, 5).unwrap(), "B", "slot 1 → B");
        assert_eq!(&*pick_where(&endpoints, all, 2, 5).unwrap(), "B", "slot 2 → B");
        assert_eq!(&*pick_where(&endpoints, all, 3, 5).unwrap(), "B", "slot 3 → B");
        assert_eq!(&*pick_where(&endpoints, all, 4, 5).unwrap(), "C", "slot 4 → C");
        assert_eq!(
            &*pick_where(&endpoints, all, 5, 5).unwrap(),
            "A",
            "slot 5 wraps to 0 → A"
        );
        assert_eq!(
            &*pick_where(&endpoints, all, 9, 5).unwrap(),
            "C",
            "slot 9 wraps to 4 → C"
        );
    }

    #[test]
    fn pick_where_skips_non_candidates() {
        let endpoints = [ep("A", 2), ep("B", 2), ep("C", 2)];
        let skip_b = |ep: &WeightedEndpoint| &*ep.address != "B";
        assert_eq!(&*pick_where(&endpoints, skip_b, 0, 4).unwrap(), "A", "slot 0 → A");
        assert_eq!(&*pick_where(&endpoints, skip_b, 1, 4).unwrap(), "A", "slot 1 → A");
        assert_eq!(&*pick_where(&endpoints, skip_b, 2, 4).unwrap(), "C", "slot 2 → C");
        assert_eq!(&*pick_where(&endpoints, skip_b, 3, 4).unwrap(), "C", "slot 3 → C");
    }

    // -------------------------------------------------------------------------
    // Test Utilities
    // -------------------------------------------------------------------------

    fn ep(addr: &str, weight: u32) -> WeightedEndpoint {
        WeightedEndpoint::simple(Arc::from(addr), weight)
    }

    fn health_state(n: usize) -> ClusterHealthState {
        let healths: Vec<_> = std::iter::repeat_with(EndpointHealth::new).take(n).collect();
        let addrs: Vec<_> = (0..n)
            .map(|i| Arc::from(format!("10.0.0.{}:80", i + 1).as_str()))
            .collect();
        Arc::new(ClusterHealthEntry::new(healths, addrs, None, None))
    }
}