rskit_discovery/
balancer.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3use dashmap::DashMap;
4
5use crate::instance::ServiceInstance;
6
7/// Strategy for picking one instance from a set of candidates.
8pub trait LoadBalancer: Send + Sync {
9    /// Select an instance, or `None` if the slice is empty.
10    fn pick<'a>(&self, instances: &'a [ServiceInstance]) -> Option<&'a ServiceInstance>;
11}
12
13/// Cycles through instances in order.
14pub struct RoundRobin {
15    counter: AtomicUsize,
16}
17
18impl RoundRobin {
19    /// Create a new round-robin balancer.
20    pub fn new() -> Self {
21        Self {
22            counter: AtomicUsize::new(0),
23        }
24    }
25}
26
27impl Default for RoundRobin {
28    fn default() -> Self {
29        Self::new()
30    }
31}
32
33impl LoadBalancer for RoundRobin {
34    fn pick<'a>(&self, instances: &'a [ServiceInstance]) -> Option<&'a ServiceInstance> {
35        if instances.is_empty() {
36            return None;
37        }
38        let idx = self.counter.fetch_add(1, Ordering::Relaxed) % instances.len();
39        Some(&instances[idx])
40    }
41}
42
43/// Picks a random instance each time.
44pub struct Random;
45
46impl LoadBalancer for Random {
47    fn pick<'a>(&self, instances: &'a [ServiceInstance]) -> Option<&'a ServiceInstance> {
48        if instances.is_empty() {
49            return None;
50        }
51        use rand::Rng;
52        let idx = rand::rng().random_range(0..instances.len());
53        Some(&instances[idx])
54    }
55}
56
57/// Picks an instance randomly according to each instance's relative weight.
58#[derive(Debug, Clone, Copy, Default)]
59pub struct Weighted;
60
61impl LoadBalancer for Weighted {
62    fn pick<'a>(&self, instances: &'a [ServiceInstance]) -> Option<&'a ServiceInstance> {
63        if instances.is_empty() {
64            return None;
65        }
66
67        let total = instances
68            .iter()
69            .map(|instance| u64::from(instance.weight.max(1)))
70            .sum::<u64>();
71        use rand::Rng;
72        let mut slot = rand::rng().random_range(0..total);
73        for instance in instances {
74            let weight = u64::from(instance.weight.max(1));
75            if slot < weight {
76                return Some(instance);
77            }
78            slot -= weight;
79        }
80        instances.last()
81    }
82}
83
84/// Tracks in-flight requests per instance and picks the one with the fewest.
85pub struct LeastConnections {
86    in_flight: DashMap<String, AtomicUsize>,
87}
88
89impl LeastConnections {
90    /// Create a new least-connections balancer.
91    pub fn new() -> Self {
92        Self {
93            in_flight: DashMap::new(),
94        }
95    }
96
97    /// Increment the in-flight count for an instance.
98    pub fn acquire(&self, id: &str) {
99        self.in_flight
100            .entry(id.to_owned())
101            .or_insert_with(|| AtomicUsize::new(0))
102            .fetch_add(1, Ordering::Relaxed);
103    }
104
105    /// Decrement the in-flight count for an instance.
106    pub fn release(&self, id: &str) {
107        if let Some(counter) = self.in_flight.get(id) {
108            counter.fetch_sub(1, Ordering::Relaxed);
109        }
110    }
111}
112
113impl Default for LeastConnections {
114    fn default() -> Self {
115        Self::new()
116    }
117}
118
119impl LoadBalancer for LeastConnections {
120    fn pick<'a>(&self, instances: &'a [ServiceInstance]) -> Option<&'a ServiceInstance> {
121        if instances.is_empty() {
122            return None;
123        }
124        instances.iter().min_by_key(|inst| {
125            self.in_flight
126                .get(&inst.id)
127                .map(|c| c.load(Ordering::Relaxed))
128                .unwrap_or(0)
129        })
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136    use std::collections::HashMap;
137
138    fn instance(id: &str) -> ServiceInstance {
139        ServiceInstance {
140            id: id.into(),
141            name: "svc".into(),
142            address: "127.0.0.1".into(),
143            port: 8080,
144            healthy: true,
145            weight: 1,
146            tags: vec![],
147            metadata: HashMap::new(),
148        }
149    }
150
151    #[test]
152    fn round_robin_cycles() {
153        let rr = RoundRobin::new();
154        let instances = vec![instance("a"), instance("b"), instance("c")];
155
156        assert_eq!(rr.pick(&instances).unwrap().id, "a");
157        assert_eq!(rr.pick(&instances).unwrap().id, "b");
158        assert_eq!(rr.pick(&instances).unwrap().id, "c");
159        assert_eq!(rr.pick(&instances).unwrap().id, "a");
160    }
161
162    #[test]
163    fn round_robin_empty() {
164        let rr = RoundRobin::new();
165        assert!(rr.pick(&[]).is_none());
166    }
167
168    #[test]
169    fn random_non_empty() {
170        let instances = vec![instance("a"), instance("b")];
171        let picked = Random.pick(&instances);
172        assert!(picked.is_some());
173    }
174
175    #[test]
176    fn weighted_non_empty() {
177        let instances = vec![instance("a"), instance("b")];
178        let picked = Weighted.pick(&instances);
179        assert!(picked.is_some());
180    }
181
182    #[test]
183    fn least_connections_prefers_idle() {
184        let lc = LeastConnections::new();
185        let instances = vec![instance("a"), instance("b")];
186
187        lc.acquire("a");
188        lc.acquire("a");
189        lc.acquire("b");
190
191        assert_eq!(lc.pick(&instances).unwrap().id, "b");
192    }
193}
rskit_discovery/balancer.rs

rskit_discovery/
balancer.rs
