loadwise_core/strategy/
p2c.rs

1use rand::Rng;
2
3use super::{SelectionContext, Strategy};
4use crate::LoadMetric;
5
6/// Power of Two Choices: randomly pick two nodes, select the one with lower load.
7///
8/// Provides near-optimal load distribution with O(1) selection time,
9/// a significant improvement over pure random while avoiding the overhead
10/// of scanning all nodes like [`LeastLoad`](super::LeastLoad).
11///
12/// When fewer than two candidates are eligible, the sole eligible node is
13/// returned directly.
14#[derive(Debug)]
15pub struct PowerOfTwoChoices;
16
17impl PowerOfTwoChoices {
18    pub fn new() -> Self {
19        Self
20    }
21}
22
23impl Default for PowerOfTwoChoices {
24    fn default() -> Self {
25        Self
26    }
27}
28
29impl<N: LoadMetric> Strategy<N> for PowerOfTwoChoices {
30    fn select(&self, candidates: &[N], ctx: &SelectionContext) -> Option<usize> {
31        let eligible: Vec<usize> = (0..candidates.len())
32            .filter(|i| !ctx.is_excluded(*i))
33            .collect();
34        match eligible.len() {
35            0 => None,
36            1 => Some(eligible[0]),
37            len => {
38                let mut rng = rand::rng();
39                let ai = rng.random_range(0..len);
40                let mut bi = rng.random_range(0..len - 1);
41                if bi >= ai {
42                    bi += 1;
43                }
44                let a = eligible[ai];
45                let b = eligible[bi];
46                if candidates[a].load_score() <= candidates[b].load_score() {
47                    Some(a)
48                } else {
49                    Some(b)
50                }
51            }
52        }
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use super::*;
59
60    struct L(f64);
61    impl LoadMetric for L {
62        fn load_score(&self) -> f64 {
63            self.0
64        }
65    }
66
67    #[test]
68    fn picks_lower_of_two() {
69        let p2c = PowerOfTwoChoices::new();
70        // With only 2 candidates, always picks the lower-load one
71        let nodes = [L(10.0), L(1.0)];
72        let ctx = SelectionContext::default();
73        assert_eq!(p2c.select(&nodes, &ctx), Some(1));
74    }
75
76    #[test]
77    fn skips_excluded() {
78        let p2c = PowerOfTwoChoices::new();
79        let nodes = [L(1.0), L(2.0), L(3.0)];
80        let ctx = SelectionContext::builder().exclude(vec![0]).build();
81        // Best node (0) is excluded, should pick from {1, 2}
82        let idx = p2c.select(&nodes, &ctx).unwrap();
83        assert!(idx == 1 || idx == 2);
84    }
85
86    #[test]
87    fn all_excluded_returns_none() {
88        let p2c = PowerOfTwoChoices::new();
89        let nodes = [L(1.0)];
90        let ctx = SelectionContext::builder().exclude(vec![0]).build();
91        assert_eq!(p2c.select(&nodes, &ctx), None);
92    }
93}
loadwise_core/strategy/p2c.rs

loadwise_core/strategy/
p2c.rs
