apigate_core/balancing/
mod.rs

1mod consistent_hash;
2mod least_request;
3mod least_time;
4mod round_robin;
5
6use std::time::Duration;
7
8use crate::backend::{Backend, BackendPool};
9use crate::routing::{AffinityKey, CandidateSet};
10
11pub use consistent_hash::ConsistentHash;
12pub use least_request::LeastRequest;
13pub use least_time::LeastTime;
14pub use round_robin::RoundRobin;
15
16pub struct BackendRef<'a> {
17    pub index: usize,
18    pub backend: &'a Backend,
19}
20
21pub struct BalanceCtx<'a> {
22    pub service: &'a str,
23    pub affinity: Option<&'a AffinityKey<'a>>,
24    pub pool: &'a BackendPool,
25    pub candidates: CandidateSet<'a>,
26}
27
28impl<'a> BalanceCtx<'a> {
29    pub fn candidate_len(&self) -> usize {
30        match self.candidates {
31            CandidateSet::All => self.pool.len(),
32            CandidateSet::Indices(indices) => indices.len(),
33        }
34    }
35
36    pub fn candidate_index(&self, nth: usize) -> Option<usize> {
37        match self.candidates {
38            CandidateSet::All => {
39                if nth < self.pool.len() {
40                    Some(nth)
41                } else {
42                    None
43                }
44            }
45            CandidateSet::Indices(indices) => indices.get(nth).copied(),
46        }
47    }
48
49    pub fn candidate_backend(&self, nth: usize) -> Option<BackendRef<'a>> {
50        let index = self.candidate_index(nth)?;
51        let backend = self.pool.get(index)?;
52        Some(BackendRef { index, backend })
53    }
54
55    pub fn is_candidate(&self, backend_idx: usize) -> bool {
56        match self.candidates {
57            CandidateSet::All => backend_idx < self.pool.len(),
58            CandidateSet::Indices(indices) => indices.contains(&backend_idx),
59        }
60    }
61}
62
63#[derive(Debug, Clone, Copy)]
64pub enum ProxyErrorKind {
65    InvalidUpstreamUri,
66    UpstreamRequest,
67    NoBackends,
68    Timeout,
69}
70
71pub struct StartEvent<'a> {
72    pub service: &'a str,
73    pub backend_index: usize,
74}
75
76pub struct ResultEvent<'a> {
77    pub service: &'a str,
78    pub backend_index: usize,
79    pub status: Option<http::StatusCode>,
80    pub error: Option<ProxyErrorKind>,
81    pub head_latency: Duration,
82}
83
84pub trait Balancer: Send + Sync + 'static {
85    fn pick<'a>(&self, ctx: &'a BalanceCtx<'a>) -> Option<usize>;
86
87    fn on_start(&self, _event: &StartEvent<'_>) {}
88
89    fn on_result(&self, _event: &ResultEvent<'_>) {}
90}
apigate_core/balancing/mod.rs

apigate_core/balancing/
mod.rs
