spacegate_kernel/helper_layers/balancer/
random.rs

1use rand::distr::Distribution;
2
3use super::BalancePolicy;
4/// A policy that selects an instance randomly.
5pub struct Random<I>
6where
7    I: rand::distr::uniform::SampleUniform + std::cmp::PartialOrd,
8{
9    picker: rand::distr::weighted::WeightedIndex<I>,
10}
11
12impl<I> std::fmt::Debug for Random<I>
13where
14    I: rand::distr::uniform::SampleUniform + std::cmp::PartialOrd,
15{
16    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
17        f.debug_struct("Random").finish()
18    }
19}
20
21impl<I> Random<I>
22where
23    I: rand::distr::uniform::SampleUniform + std::cmp::PartialOrd + Clone + Default + for<'a> std::ops::AddAssign<&'a I> + rand::distr::weighted::Weight,
24{
25    pub fn new(weights: impl IntoIterator<Item = I>) -> Self {
26        Self {
27            picker: rand::distr::weighted::WeightedIndex::new(weights).expect("invalid weights"),
28        }
29    }
30}
31
32impl<I, S, R> BalancePolicy<S, R> for Random<I>
33where
34    I: rand::distr::uniform::SampleUniform + std::cmp::PartialOrd,
35{
36    fn pick<'s>(&self, instances: &'s [S], _req: &R) -> Option<&'s S> {
37        if instances.is_empty() {
38            None
39        } else if instances.len() == 1 {
40            instances.first()
41        } else {
42            let index = self.picker.sample(&mut rand::rng());
43            instances.get(index)
44        }
45    }
46}