diff --git a/src/balancer/adaptive_weight.rs b/src/balancer/adaptive_weight.rs index abcd486..e1570d8 100644 --- a/src/balancer/adaptive_weight.rs +++ b/src/balancer/adaptive_weight.rs @@ -27,7 +27,7 @@ impl AdaptiveWeightBalancer { .iter() .map(|b| AdaptiveNode { backend: b.clone(), - weight: 0f64, + weight: 1f64, }) .collect(); @@ -71,7 +71,7 @@ impl Balancer for AdaptiveWeightBalancer { let safe_w_sum = w_sum.max(1e-12); let threshold = self.alpha * (r_sum / safe_w_sum); - + for idx in 0..self.pool.len() { let node = &self.pool[idx]; @@ -85,7 +85,6 @@ impl Balancer for AdaptiveWeightBalancer { }; let ratio = risk / node.weight; - if ratio <= threshold { return Some(node.backend.clone()); } @@ -142,3 +141,126 @@ impl Balancer for AdaptiveWeightBalancer { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::backend::Backend; + use std::net::SocketAddr; + + fn backend_factory(id: &str, ip: &str, port: u16) -> Arc { + Arc::new(Backend::new( + id.to_string(), + SocketAddr::new(ip.parse().unwrap(), port), + Arc::new(RwLock::new(ServerMetrics::default())), + )) + } + + fn unused_ctx() -> ConnectionInfo { + ConnectionInfo { + client_ip: ("0.0.0.0".parse().unwrap()), + } + } + + #[test] + fn basic_weight_update_and_choose() { + let backends = BackendPool::new(vec![ + backend_factory("server-0", "127.0.0.1", 3000), + backend_factory("server-1", "127.0.0.1", 3001), + ]); + let mut b = AdaptiveWeightBalancer::new(backends.clone(), [0.5, 0.2, 0.2, 0.1], 0.5); + // initially equal weights + // update one backend to be heavily loaded + { + let mut sm0_guard = backends.backends.get(0).unwrap().metrics.write().unwrap(); + sm0_guard.update(90.0, 80.0, 10.0, 5.0); + } + { + let mut sm1_guard = backends.backends.get(1).unwrap().metrics.write().unwrap(); + sm1_guard.update(10.0, 5.0, 1.0, 1.0); + } + + // Choose backend: should pick the less loaded host server1 + let chosen = b + .choose_backend(unused_ctx()) + .expect("should choose a backend"); + + let sm0: &ServerMetrics = &backends.backends.get(0).unwrap().metrics.read().unwrap(); + let sm1: &ServerMetrics = &backends.backends.get(1).unwrap().metrics.read().unwrap(); + println!("{:?}, {:?}", sm0, sm1); + assert_eq!(chosen.id, "server-1"); + } + + #[test] + fn choose_none_when_empty() { + let mut b = + AdaptiveWeightBalancer::new(BackendPool::new(vec![]), [0.5, 0.2, 0.2, 0.1], 0.5); + assert!(b.choose_backend(unused_ctx()).is_none()); + } + + #[test] + fn ratio_triggers_immediate_selection() { + // Arrange two servers where server 1 has composite load 0 and server 2 has composite load 100. + // With alpha = 1.0 and two servers, threshold = 1.0 * (r_sum / w_sum) = 1.0 * (100 / 2) = 50. + // Server 0 ratio = 0 / 1 = 0 <= 50 so it should be chosen immediately. + let backends = BackendPool::new(vec![ + backend_factory("server-0", "127.0.0.1", 3000), + backend_factory("server-1", "127.0.0.1", 3001), + ]); + let mut b = AdaptiveWeightBalancer::new(backends.clone(), [0.25, 0.25, 0.25, 0.25], 1.0); + + { + let mut sm0_guard = backends.backends.get(0).unwrap().metrics.write().unwrap(); + sm0_guard.update(0.0, 0.0, 0.0, 0.0); + } + { + let mut sm1_guard = backends.backends.get(1).unwrap().metrics.write().unwrap(); + sm1_guard.update(100.0, 100.0, 100.0, 100.0); + } + + let chosen = b + .choose_backend(unused_ctx()) + .expect("should choose a backend"); + assert_eq!(chosen.id, "server-0"); + } + + #[test] + fn choose_min_current_load_when_no_ratio() { + // Arrange three servers with identical composite loads so no server satisfies Ri/Wi <= threshold + // (set alpha < 1 so threshold < ratio). The implementation then falls back to picking the + // server with minimum current_load + let backends = BackendPool::new(vec![ + backend_factory("server-0", "127.0.0.1", 3000), + backend_factory("server-1", "127.0.0.1", 3001), + backend_factory("server-2", "127.0.0.1", 3002), + ]); + + // set current_loads (field expected to be public) + + { + let mut sm0_guard = backends.backends.get(0).unwrap().metrics.write().unwrap(); + sm0_guard.update(10.0, 10.0, 10.0, 10.0); + } + { + let mut sm1_guard = backends.backends.get(1).unwrap().metrics.write().unwrap(); + sm1_guard.update(5.0, 5.0, 5.0, 5.0); + } + { + let mut sm2_guard = backends.backends.get(2).unwrap().metrics.write().unwrap(); + sm2_guard.update(20.0, 20.0, 20.0, 20.0); + } + + // Use coeffs that only consider CPU so composite load is easy to reason about. + let mut bal = AdaptiveWeightBalancer::new(backends.clone(), [1.0, 0.0, 0.0, 0.0], 0.5); + + // set identical composite loads > 0 for all so ratio = x and threshold = alpha * x < x + // you will have threshold = 25 for all 3 backend servers and ratio = 50 + // so that forces to choose the smallest current load backend + + let chosen = bal + .choose_backend(unused_ctx()) + .expect("should choose a backend"); + // expect server with smallest current_load server-1 + assert_eq!(chosen.id, "server-1"); + } +}