reqwest_lb/lb/
policy.rs

1use crate::lb::weight::WeightProvider;
2use crate::lb::Statistic;
3use crate::with::With;
4use http::Extensions;
5use rand::Rng;
6use std::fmt::{Debug, Formatter};
7use std::sync::atomic::Ordering;
8use std::sync::Arc;
9
10#[derive(Default)]
11pub enum LoadBalancerPolicy<I> {
12    #[default]
13    RoundRobin,
14    Random,
15    First,
16    Last,
17    Weight(Arc<dyn WeightProvider<I> + Send + Sync>),
18    Dynamic(Arc<dyn LoadBalancerPolicyTrait<I> + Send + Sync>),
19}
20
21impl<I> Debug for LoadBalancerPolicy<I> {
22    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
23        match self {
24            LoadBalancerPolicy::RoundRobin => f.write_str("RoundRobin"),
25            LoadBalancerPolicy::Random => f.write_str("Random"),
26            LoadBalancerPolicy::First => f.write_str("First"),
27            LoadBalancerPolicy::Last => f.write_str("Last"),
28            LoadBalancerPolicy::Weight(_) => f.write_str("Weight(f)"),
29            LoadBalancerPolicy::Dynamic(_) => f.write_str("Dynamic(f)"),
30        }
31    }
32}
33
34impl<I> Clone for LoadBalancerPolicy<I> {
35    fn clone(&self) -> Self {
36        match self {
37            LoadBalancerPolicy::RoundRobin => LoadBalancerPolicy::RoundRobin,
38            LoadBalancerPolicy::Random => LoadBalancerPolicy::Random,
39            LoadBalancerPolicy::First => LoadBalancerPolicy::First,
40            LoadBalancerPolicy::Last => LoadBalancerPolicy::Last,
41            LoadBalancerPolicy::Weight(f) => LoadBalancerPolicy::Weight(f.clone()),
42            LoadBalancerPolicy::Dynamic(f) => LoadBalancerPolicy::Dynamic(f.clone()),
43        }
44    }
45}
46
47impl<I> LoadBalancerPolicy<I> {
48    pub fn weight<F: Fn(&I) -> usize + Send + Sync + 'static>(f: F) -> Self {
49        Self::Weight(Arc::new(f))
50    }
51
52    pub fn dynamic<F: Fn(&[I], &Extensions) -> usize + Send + Sync + 'static>(f: F) -> Self {
53        Self::Dynamic(Arc::new(f))
54    }
55}
56
57pub trait LoadBalancerPolicyTrait<I>: sealed::Sealed<I> {
58    fn choose(&self, items: &[I], extensions: &mut Extensions) -> usize;
59}
60
61impl<I> sealed::Sealed<I> for LoadBalancerPolicy<I> {}
62
63impl<I> LoadBalancerPolicyTrait<I> for LoadBalancerPolicy<I> {
64    fn choose(&self, items: &[I], extensions: &mut Extensions) -> usize {
65        let len = items.len();
66        assert!(len > 1);
67        match self {
68            LoadBalancerPolicy::RoundRobin => match extensions.get::<Statistic>() {
69                Some(statistic) => {
70                    let count = statistic.count.load(Ordering::Relaxed).saturating_sub(1);
71                    (count % (len as u64)) as usize
72                }
73                None => 0,
74            },
75            LoadBalancerPolicy::Random => rand::thread_rng().gen_range(0..len),
76            LoadBalancerPolicy::First => 0,
77            LoadBalancerPolicy::Last => items.len() - 1,
78            LoadBalancerPolicy::Weight(f) => {
79                let indexes = items
80                    .iter()
81                    .enumerate()
82                    .map(|(index, item)| (index, f.weight(item)))
83                    .flat_map(|(index, len)| {
84                        Vec::with_capacity(len).with(|c| {
85                            for _ in 0..len {
86                                c.push(index);
87                            }
88                        })
89                    })
90                    .collect::<Vec<_>>();
91                let index = rand::thread_rng().gen_range(0..indexes.len());
92                indexes[index]
93            }
94            LoadBalancerPolicy::Dynamic(f) => f.choose(items, extensions),
95        }
96    }
97}
98
99impl<I, F> sealed::Sealed<I> for F where F: Fn(&[I], &Extensions) -> usize {}
100
101impl<I, F> LoadBalancerPolicyTrait<I> for F
102where
103    F: Fn(&[I], &Extensions) -> usize,
104{
105    fn choose(&self, items: &[I], extensions: &mut Extensions) -> usize {
106        self(items, extensions)
107    }
108}
109
110mod sealed {
111    pub trait Sealed<I> {}
112}
reqwest_lb/lb/policy.rs

reqwest_lb/lb/
policy.rs
