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 Balancer for LeastRequest {
30    fn pick(&self, ctx: &BalanceCtx) -> Option<usize> {
31        let len = ctx.candidate_len();
32        if len == 0 {
33            return None;
34        }
35
36        let counters = self.counters(ctx.pool.len());
37        let offset = self.offset.fetch_add(1, Ordering::Relaxed);
38        let mut best_index = None;
39        let mut best_count = usize::MAX;
40
41        for i in 0..len {
42            let nth = (offset + i) % len;
43            if let Some(idx) = ctx.candidate_index(nth) {
44                let count = counters
45                    .get(idx)
46                    .map(|a| a.load(Ordering::Relaxed))
47                    .unwrap_or(0);
48                if count < best_count {
49                    best_count = count;
50                    best_index = Some(idx);
51                }
52            }
53        }
54
55        best_index
56    }
57
58    fn on_start(&self, event: &StartEvent) {
59        if let Some(counters) = self.in_flight.get() {
60            if let Some(counter) = counters.get(event.backend_index) {
61                counter.fetch_add(1, Ordering::Relaxed);
62            }
63        }
64    }
65
66    fn on_result(&self, event: &ResultEvent) {
67        if let Some(counters) = self.in_flight.get() {
68            if let Some(counter) = counters.get(event.backend_index) {
69                counter.fetch_sub(1, Ordering::Relaxed);
70            }
71        }
72    }
73}
apigate_core/balancing/least_request.rs

apigate_core/balancing/
least_request.rs
