loadwise-core 0.1.0

Core traits, strategies, and in-memory stores for loadwise
Documentation
use rand::Rng;

use super::{SelectionContext, Strategy};
use crate::RateMetric;

/// Selects the node with the most remaining capacity, with epsilon tie-breaking.
///
/// Scoring: each candidate's score is its [`remaining_ratio()`](RateMetric::remaining_ratio).
/// When multiple candidates score within `epsilon` of the best,
/// one is chosen randomly to avoid thundering-herd on the single "most available" node.
///
/// # Panics
///
/// Panics if `epsilon` is negative or NaN.
///
/// # Examples
///
/// ```
/// # extern crate loadwise_core as loadwise;
/// use loadwise::{Strategy, RateMetric, SelectionContext};
/// use loadwise::strategy::MostAvailable;
///
/// struct Key { remaining: f64 }
/// impl RateMetric for Key {
///     fn remaining_ratio(&self) -> f64 { self.remaining }
/// }
///
/// let strategy = MostAvailable::default();
/// let keys = [Key { remaining: 0.2 }, Key { remaining: 0.8 }, Key { remaining: 0.75 }];
/// let ctx = SelectionContext::default();
///
/// // Picks index 1 or 2 (both within epsilon=0.05 of best 0.8)
/// let idx = strategy.select(&keys, &ctx).unwrap();
/// assert!(idx == 1 || idx == 2);
/// ```
#[derive(Debug, bon::Builder)]
pub struct MostAvailable {
    /// Two candidates are considered equivalent if their remaining_ratio
    /// differs by less than this value. Prevents all traffic from hitting
    /// a single "most available" node.
    /// **Default: 0.05 (5%)**
    #[builder(default = 0.05)]
    epsilon: f64,
}

impl Default for MostAvailable {
    fn default() -> Self {
        Self::builder().build()
    }
}

impl<N: RateMetric> Strategy<N> for MostAvailable {
    fn select(&self, candidates: &[N], ctx: &SelectionContext) -> Option<usize> {
        assert!(
            self.epsilon.is_finite() && self.epsilon >= 0.0,
            "epsilon must be finite and non-negative"
        );

        if candidates.is_empty() {
            return None;
        }

        // Single pass: collect eligible (index, score) pairs and track best
        let mut best_score = f64::NEG_INFINITY;
        let eligible: Vec<(usize, f64)> = candidates
            .iter()
            .enumerate()
            .filter(|(i, _)| !ctx.is_excluded(*i))
            .map(|(i, node)| {
                let score = node.remaining_ratio();
                if score > best_score {
                    best_score = score;
                }
                (i, score)
            })
            .collect();

        if eligible.is_empty() {
            return None;
        }

        // Keep candidates within epsilon of best
        let threshold = best_score - self.epsilon;
        let tier: Vec<usize> = eligible
            .into_iter()
            .filter(|(_, score)| *score >= threshold)
            .map(|(i, _)| i)
            .collect();

        match tier.len() {
            0 => None,
            1 => Some(tier[0]),
            n => Some(tier[rand::rng().random_range(0..n)]),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    struct R(f64);
    impl RateMetric for R {
        fn remaining_ratio(&self) -> f64 {
            self.0
        }
    }

    #[test]
    fn picks_highest_remaining() {
        let s = MostAvailable::builder().epsilon(0.0).build();
        let nodes = [R(0.3), R(0.9), R(0.5)];
        let ctx = SelectionContext::default();
        assert_eq!(s.select(&nodes, &ctx), Some(1));
    }

    #[test]
    fn groups_within_epsilon() {
        let s = MostAvailable::builder().epsilon(0.1).build();
        // 0.85 and 0.9 are within epsilon=0.1, 0.3 is not
        let nodes = [R(0.3), R(0.9), R(0.85)];
        let ctx = SelectionContext::default();

        let mut seen = std::collections::HashSet::new();
        for _ in 0..100 {
            seen.insert(s.select(&nodes, &ctx).unwrap());
        }
        // Should pick from {1, 2} but never 0
        assert!(seen.contains(&1));
        assert!(seen.contains(&2));
        assert!(!seen.contains(&0));
    }

    #[test]
    fn respects_exclude() {
        let s = MostAvailable::default();
        let nodes = [R(0.9), R(0.5)];
        let ctx = SelectionContext::builder().exclude(vec![0]).build();
        assert_eq!(s.select(&nodes, &ctx), Some(1));
    }

    #[test]
    fn empty_returns_none() {
        let s = MostAvailable::default();
        let nodes: [R; 0] = [];
        assert_eq!(s.select(&nodes, &SelectionContext::default()), None);
    }

    #[test]
    fn all_excluded_returns_none() {
        let s = MostAvailable::default();
        let nodes = [R(0.5)];
        let ctx = SelectionContext::builder().exclude(vec![0]).build();
        assert_eq!(s.select(&nodes, &ctx), None);
    }

    #[test]
    #[should_panic(expected = "epsilon must be finite and non-negative")]
    fn negative_epsilon_panics() {
        let s = MostAvailable::builder().epsilon(-0.1).build();
        let nodes = [R(0.5)];
        s.select(&nodes, &SelectionContext::default());
    }
}