apigate_core/balancing/
least_request.rs

1use std::sync::OnceLock;
2use std::sync::atomic::{AtomicUsize, Ordering};
3
4use super::{BalanceCtx, Balancer, ResultEvent, StartEvent};
5
6pub struct LeastRequest {
7    in_flight: OnceLock<Box<[AtomicUsize]>>,
8    offset: AtomicUsize,
9}
10
11impl LeastRequest {
12    pub fn new() -> Self {
13        Self {
14            in_flight: OnceLock::new(),
15            offset: AtomicUsize::new(0),
16        }
17    }
18
19    fn counters(&self, pool_len: usize) -> &[AtomicUsize] {
20        self.in_flight.get_or_init(|| {
21            (0..pool_len)
22                .map(|_| AtomicUsize::new(0))
23                .collect::<Vec<_>>()
24                .into_boxed_slice()
25        })
26    }
27}
28
29impl Default for LeastRequest {
30    fn default() -> Self {
31        Self::new()
32    }
33}
34
35impl Balancer for LeastRequest {
36    fn pick(&self, ctx: &BalanceCtx) -> Option<usize> {
37        let len = ctx.candidate_len();
38        if len == 0 {
39            return None;
40        }
41
42        let counters = self.counters(ctx.pool.len());
43        let offset = self.offset.fetch_add(1, Ordering::Relaxed);
44        let mut best_index = None;
45        let mut best_count = usize::MAX;
46
47        for i in 0..len {
48            let nth = (offset + i) % len;
49            if let Some(idx) = ctx.candidate_index(nth) {
50                let count = counters
51                    .get(idx)
52                    .map(|a| a.load(Ordering::Relaxed))
53                    .unwrap_or(0);
54                if count < best_count {
55                    best_count = count;
56                    best_index = Some(idx);
57                }
58            }
59        }
60
61        best_index
62    }
63
64    fn on_start(&self, event: &StartEvent) {
65        if let Some(counters) = self.in_flight.get() {
66            if let Some(counter) = counters.get(event.backend_index) {
67                counter.fetch_add(1, Ordering::Relaxed);
68            }
69        }
70    }
71
72    fn on_result(&self, event: &ResultEvent) {
73        if let Some(counters) = self.in_flight.get() {
74            if let Some(counter) = counters.get(event.backend_index) {
75                counter.fetch_sub(1, Ordering::Relaxed);
76            }
77        }
78    }
79}
apigate_core/balancing/least_request.rs

apigate_core/balancing/
least_request.rs
