featherbit/
balancer.rs

1//! Shared upstream load balancer.
2//!
3//! A [`Balancer`] owns a pool of backend [`Target`]s and picks one per
4//! request/connection according to a [`Strategy`] (round-robin,
5//! least-connections, or client-IP hash). It is deliberately transport-neutral:
6//! the HTTP `upstream` node and the L4 (TCP/UDP) stream proxy both select
7//! targets through it, so the load-balancing logic lives in exactly one place.
8//!
9//! `least_connections` is backed by per-target in-flight counters. Callers bump
10//! a counter for the life of a request/connection by holding a guard:
11//! [`ConnGuard`] (borrowed — held on the stack for a single request) or
12//! [`OwnedConnGuard`] (owns an `Arc<Balancer>` — for a long-lived connection
13//! stored outside the borrowing scope, e.g. a UDP session in a map). Both
14//! decrement on drop.
15
16use std::sync::atomic::{AtomicUsize, Ordering};
17use std::sync::Arc;
18
19use serde::Deserialize;
20
21/// A single backend address (`host:port`) connections can be forwarded to.
22#[derive(Debug, Clone, Deserialize)]
23pub struct Target {
24    pub host: String,
25    pub port: u16,
26}
27
28/// Strategy for picking a target from the pool.
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
30pub enum Strategy {
31    /// Cycles through targets in order (the default).
32    #[default]
33    RoundRobin,
34    /// Picks the target with the fewest in-flight connections.
35    LeastConnections,
36    /// Hashes the client IP (port stripped) so a client sticks to one target.
37    IpHash,
38}
39
40impl Strategy {
41    /// Accepts spec spelling (`round_robin`, `least_connections`, `ip_hash`)
42    /// plus the hyphenated/short variants older UI configs saved
43    /// (`round-robin`, `least-conn`).
44    pub fn parse(value: &str) -> Result<Self, String> {
45        match value.to_lowercase().replace('-', "_").as_str() {
46            "round_robin" => Ok(Self::RoundRobin),
47            "least_connections" | "least_conn" => Ok(Self::LeastConnections),
48            "ip_hash" => Ok(Self::IpHash),
49            other => Err(format!(
50                "Unknown load_balancing '{}' — supported: round_robin, least_connections, ip_hash",
51                other
52            )),
53        }
54    }
55}
56
57/// A pool of backend targets plus the state driving target selection.
58pub struct Balancer {
59    targets: Vec<Target>,
60    strategy: Strategy,
61    /// Monotonic pick counter driving round-robin selection.
62    counter: AtomicUsize,
63    /// Per-target in-flight counts (parallel to `targets`), used by
64    /// least-connections selection and kept current via the guards.
65    in_flight: Vec<AtomicUsize>,
66}
67
68impl Balancer {
69    /// Builds a balancer over `targets`. Errors if the pool is empty.
70    pub fn new(targets: Vec<Target>, strategy: Strategy) -> Result<Self, String> {
71        if targets.is_empty() {
72            return Err("load balancer requires at least one target".to_string());
73        }
74        let in_flight = targets.iter().map(|_| AtomicUsize::new(0)).collect();
75        Ok(Self {
76            targets,
77            strategy,
78            counter: AtomicUsize::new(0),
79            in_flight,
80        })
81    }
82
83    /// Picks the index of the target to use; `remote_addr` (client `ip:port`)
84    /// is only consulted for `ip_hash`.
85    pub fn select(&self, remote_addr: &str) -> usize {
86        match self.strategy {
87            Strategy::RoundRobin => {
88                self.counter.fetch_add(1, Ordering::Relaxed) % self.targets.len()
89            }
90            Strategy::LeastConnections => self
91                .in_flight
92                .iter()
93                .enumerate()
94                .min_by_key(|(_, c)| c.load(Ordering::Relaxed))
95                .map(|(i, _)| i)
96                .unwrap_or(0),
97            Strategy::IpHash => {
98                use std::hash::{Hash, Hasher};
99                // Hash only the IP so all connections from one client stick to
100                // the same target regardless of ephemeral port.
101                let ip = remote_addr
102                    .rsplit_once(':')
103                    .map_or(remote_addr, |(ip, _)| ip);
104                let mut hasher = std::collections::hash_map::DefaultHasher::new();
105                ip.hash(&mut hasher);
106                (hasher.finish() as usize) % self.targets.len()
107            }
108        }
109    }
110
111    /// The target at `idx` (from a prior [`select`](Self::select)).
112    pub fn target(&self, idx: usize) -> &Target {
113        &self.targets[idx]
114    }
115
116    /// Number of targets in the pool.
117    // Pool accessors: `in_flight_count`/`strategy` are exercised by tests;
118    // `len`/`is_empty` round out the API.
119    #[allow(dead_code)]
120    pub fn len(&self) -> usize {
121        self.targets.len()
122    }
123
124    /// Whether the pool is empty (always false for a constructed `Balancer`).
125    #[allow(dead_code)]
126    pub fn is_empty(&self) -> bool {
127        self.targets.is_empty()
128    }
129
130    /// Current in-flight count for a target (for tests and future metrics).
131    #[allow(dead_code)]
132    pub fn in_flight_count(&self, idx: usize) -> usize {
133        self.in_flight[idx].load(Ordering::Relaxed)
134    }
135
136    /// The configured strategy.
137    #[allow(dead_code)]
138    pub fn strategy(&self) -> Strategy {
139        self.strategy
140    }
141
142    /// Increments the in-flight count for `idx`, returning a borrowed guard that
143    /// decrements it on drop. For a request/connection whose lifetime is a
144    /// single stack scope (HTTP request, TCP connection).
145    pub fn acquire(&self, idx: usize) -> ConnGuard<'_> {
146        self.in_flight[idx].fetch_add(1, Ordering::Relaxed);
147        ConnGuard(&self.in_flight[idx])
148    }
149
150    /// Like [`acquire`](Self::acquire) but owns an `Arc<Balancer>`, so the guard
151    /// can be stored past the borrowing scope (e.g. a UDP session in a map).
152    pub fn owned_acquire(self: &Arc<Self>, idx: usize) -> OwnedConnGuard {
153        self.in_flight[idx].fetch_add(1, Ordering::Relaxed);
154        OwnedConnGuard {
155            balancer: self.clone(),
156            idx,
157        }
158    }
159}
160
161/// Decrements a target's in-flight counter when a request/connection completes.
162pub struct ConnGuard<'a>(&'a AtomicUsize);
163
164impl Drop for ConnGuard<'_> {
165    fn drop(&mut self) {
166        self.0.fetch_sub(1, Ordering::Relaxed);
167    }
168}
169
170/// Owned counterpart to [`ConnGuard`] for connections tracked outside the
171/// borrowing scope; decrements the target's in-flight counter on drop.
172pub struct OwnedConnGuard {
173    balancer: Arc<Balancer>,
174    idx: usize,
175}
176
177impl Drop for OwnedConnGuard {
178    fn drop(&mut self) {
179        self.balancer.in_flight[self.idx].fetch_sub(1, Ordering::Relaxed);
180    }
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186
187    fn targets(n: usize) -> Vec<Target> {
188        (0..n)
189            .map(|i| Target {
190                host: format!("backend-{}", i),
191                port: 3000,
192            })
193            .collect()
194    }
195
196    #[test]
197    fn test_strategy_parse_aliases() {
198        assert_eq!(
199            Strategy::parse("round_robin").unwrap(),
200            Strategy::RoundRobin
201        );
202        assert_eq!(
203            Strategy::parse("round-robin").unwrap(),
204            Strategy::RoundRobin
205        );
206        assert_eq!(
207            Strategy::parse("least_connections").unwrap(),
208            Strategy::LeastConnections
209        );
210        assert_eq!(
211            Strategy::parse("least-conn").unwrap(),
212            Strategy::LeastConnections
213        );
214        assert_eq!(Strategy::parse("IP_HASH").unwrap(), Strategy::IpHash);
215        assert!(Strategy::parse("random").is_err());
216    }
217
218    #[test]
219    fn test_new_rejects_empty_pool() {
220        assert!(Balancer::new(vec![], Strategy::RoundRobin).is_err());
221    }
222
223    #[test]
224    fn test_round_robin_cycles() {
225        let b = Balancer::new(targets(3), Strategy::RoundRobin).unwrap();
226        let picks: Vec<usize> = (0..6).map(|_| b.select("1.2.3.4:555")).collect();
227        assert_eq!(picks, vec![0, 1, 2, 0, 1, 2]);
228    }
229
230    #[test]
231    fn test_ip_hash_is_sticky_per_ip() {
232        let b = Balancer::new(targets(3), Strategy::IpHash).unwrap();
233        let a = b.select("10.0.0.1:1111");
234        // same IP, different ephemeral port -> same target
235        assert_eq!(a, b.select("10.0.0.1:2222"));
236        assert_eq!(a, b.select("10.0.0.1:3333"));
237    }
238
239    #[test]
240    fn test_least_connections_picks_idle_target_via_guards() {
241        let b = Balancer::new(targets(3), Strategy::LeastConnections).unwrap();
242        // Load targets 0 and 2 heavily, leave 1 idle.
243        let _g0a = b.acquire(0);
244        let _g0b = b.acquire(0);
245        let _g2 = b.acquire(2);
246        assert_eq!(b.select("1.2.3.4:555"), 1);
247        assert_eq!(b.in_flight_count(0), 2);
248        assert_eq!(b.in_flight_count(2), 1);
249    }
250
251    #[test]
252    fn test_guard_decrements_on_drop() {
253        let b = Balancer::new(targets(2), Strategy::LeastConnections).unwrap();
254        {
255            let _g = b.acquire(0);
256            assert_eq!(b.in_flight_count(0), 1);
257        }
258        assert_eq!(b.in_flight_count(0), 0);
259    }
260
261    #[test]
262    fn test_owned_guard_decrements_on_drop() {
263        let b = Arc::new(Balancer::new(targets(2), Strategy::LeastConnections).unwrap());
264        let g = b.owned_acquire(1);
265        assert_eq!(b.in_flight_count(1), 1);
266        drop(g);
267        assert_eq!(b.in_flight_count(1), 0);
268    }
269}
featherbit/balancer.rs

featherbit/
balancer.rs
