Skip to main content

loadwise_core/strategy/
most_available.rs

1use rand::Rng;
2
3use super::{SelectionContext, Strategy};
4use crate::RateMetric;
5
6/// Selects the node with the most remaining capacity, with epsilon tie-breaking.
7///
8/// Scoring: each candidate's score is its [`remaining_ratio()`](RateMetric::remaining_ratio).
9/// When multiple candidates score within `epsilon` of the best,
10/// one is chosen randomly to avoid thundering-herd on the single "most available" node.
11///
12/// # Panics
13///
14/// Panics if `epsilon` is negative or NaN.
15///
16/// # Examples
17///
18/// ```
19/// # extern crate loadwise_core as loadwise;
20/// use loadwise::{Strategy, RateMetric, SelectionContext};
21/// use loadwise::strategy::MostAvailable;
22///
23/// struct Key { remaining: f64 }
24/// impl RateMetric for Key {
25///     fn remaining_ratio(&self) -> f64 { self.remaining }
26/// }
27///
28/// let strategy = MostAvailable::default();
29/// let keys = [Key { remaining: 0.2 }, Key { remaining: 0.8 }, Key { remaining: 0.75 }];
30/// let ctx = SelectionContext::default();
31///
32/// // Picks index 1 or 2 (both within epsilon=0.05 of best 0.8)
33/// let idx = strategy.select(&keys, &ctx).unwrap();
34/// assert!(idx == 1 || idx == 2);
35/// ```
36#[derive(Debug, bon::Builder)]
37pub struct MostAvailable {
38    /// Two candidates are considered equivalent if their remaining_ratio
39    /// differs by less than this value. Prevents all traffic from hitting
40    /// a single "most available" node.
41    /// **Default: 0.05 (5%)**
42    #[builder(default = 0.05)]
43    epsilon: f64,
44}
45
46impl Default for MostAvailable {
47    fn default() -> Self {
48        Self::builder().build()
49    }
50}
51
52impl<N: RateMetric> Strategy<N> for MostAvailable {
53    fn select(&self, candidates: &[N], ctx: &SelectionContext) -> Option<usize> {
54        assert!(
55            self.epsilon.is_finite() && self.epsilon >= 0.0,
56            "epsilon must be finite and non-negative"
57        );
58
59        if candidates.is_empty() {
60            return None;
61        }
62
63        // Single pass: collect eligible (index, score) pairs and track best
64        let mut best_score = f64::NEG_INFINITY;
65        let eligible: Vec<(usize, f64)> = candidates
66            .iter()
67            .enumerate()
68            .filter(|(i, _)| !ctx.is_excluded(*i))
69            .map(|(i, node)| {
70                let score = node.remaining_ratio();
71                if score > best_score {
72                    best_score = score;
73                }
74                (i, score)
75            })
76            .collect();
77
78        if eligible.is_empty() {
79            return None;
80        }
81
82        // Keep candidates within epsilon of best
83        let threshold = best_score - self.epsilon;
84        let tier: Vec<usize> = eligible
85            .into_iter()
86            .filter(|(_, score)| *score >= threshold)
87            .map(|(i, _)| i)
88            .collect();
89
90        match tier.len() {
91            0 => None,
92            1 => Some(tier[0]),
93            n => Some(tier[rand::rng().random_range(0..n)]),
94        }
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101
102    struct R(f64);
103    impl RateMetric for R {
104        fn remaining_ratio(&self) -> f64 {
105            self.0
106        }
107    }
108
109    #[test]
110    fn picks_highest_remaining() {
111        let s = MostAvailable::builder().epsilon(0.0).build();
112        let nodes = [R(0.3), R(0.9), R(0.5)];
113        let ctx = SelectionContext::default();
114        assert_eq!(s.select(&nodes, &ctx), Some(1));
115    }
116
117    #[test]
118    fn groups_within_epsilon() {
119        let s = MostAvailable::builder().epsilon(0.1).build();
120        // 0.85 and 0.9 are within epsilon=0.1, 0.3 is not
121        let nodes = [R(0.3), R(0.9), R(0.85)];
122        let ctx = SelectionContext::default();
123
124        let mut seen = std::collections::HashSet::new();
125        for _ in 0..100 {
126            seen.insert(s.select(&nodes, &ctx).unwrap());
127        }
128        // Should pick from {1, 2} but never 0
129        assert!(seen.contains(&1));
130        assert!(seen.contains(&2));
131        assert!(!seen.contains(&0));
132    }
133
134    #[test]
135    fn respects_exclude() {
136        let s = MostAvailable::default();
137        let nodes = [R(0.9), R(0.5)];
138        let ctx = SelectionContext::builder().exclude(vec![0]).build();
139        assert_eq!(s.select(&nodes, &ctx), Some(1));
140    }
141
142    #[test]
143    fn empty_returns_none() {
144        let s = MostAvailable::default();
145        let nodes: [R; 0] = [];
146        assert_eq!(s.select(&nodes, &SelectionContext::default()), None);
147    }
148
149    #[test]
150    fn all_excluded_returns_none() {
151        let s = MostAvailable::default();
152        let nodes = [R(0.5)];
153        let ctx = SelectionContext::builder().exclude(vec![0]).build();
154        assert_eq!(s.select(&nodes, &ctx), None);
155    }
156
157    #[test]
158    #[should_panic(expected = "epsilon must be finite and non-negative")]
159    fn negative_epsilon_panics() {
160        let s = MostAvailable::builder().epsilon(-0.1).build();
161        let nodes = [R(0.5)];
162        s.select(&nodes, &SelectionContext::default());
163    }
164}