From 3b96043dc670fd85273607eb586565b3beed1b3b Mon Sep 17 00:00:00 2001 From: Jeremy Janella Date: Wed, 10 Dec 2025 17:25:16 -0500 Subject: [PATCH 1/3] feat: adaptive weight tests --- src/balancer/adaptive_weight.rs | 123 ++++++++++++++++++++++++++++++++ 1 file changed, 123 insertions(+) diff --git a/src/balancer/adaptive_weight.rs b/src/balancer/adaptive_weight.rs index abcd486..c4c1284 100644 --- a/src/balancer/adaptive_weight.rs +++ b/src/balancer/adaptive_weight.rs @@ -142,3 +142,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 1 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-1"); + } + + #[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"); + } +} From a9db727bdeef847ef8185608c67a1293c9b3f469 Mon Sep 17 00:00:00 2001 From: nnhphong Date: Wed, 10 Dec 2025 18:08:54 -0500 Subject: [PATCH 2/3] fix the test, all pass now --- src/balancer/adaptive_weight.rs | 52 ++++++++++++--------------------- 1 file changed, 18 insertions(+), 34 deletions(-) diff --git a/src/balancer/adaptive_weight.rs b/src/balancer/adaptive_weight.rs index c4c1284..49148ca 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()); } @@ -172,14 +171,10 @@ mod tests { 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); - } + 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 @@ -203,26 +198,22 @@ mod tests { 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 1 ratio = 0 / 1 = 0 <= 50 so it should be chosen immediately. + // 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 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-1"); + assert_eq!(chosen.id, "server-0"); } #[test] @@ -237,19 +228,12 @@ mod tests { ]); // 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); - } + 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); From 7a68e4b17bccb7fca5b4288cea7bd1e0f1bdadc7 Mon Sep 17 00:00:00 2001 From: nnhphong Date: Wed, 10 Dec 2025 18:26:52 -0500 Subject: [PATCH 3/3] fix the test --- src/balancer/adaptive_weight.rs | 43 ++++++++++++++++++++++----------- 1 file changed, 29 insertions(+), 14 deletions(-) diff --git a/src/balancer/adaptive_weight.rs b/src/balancer/adaptive_weight.rs index 49148ca..e1570d8 100644 --- a/src/balancer/adaptive_weight.rs +++ b/src/balancer/adaptive_weight.rs @@ -171,10 +171,14 @@ mod tests { 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); + { + 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 @@ -205,10 +209,14 @@ mod tests { ]); 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 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()) @@ -228,12 +236,19 @@ mod tests { ]); // 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); + + { + 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);