jjanella/ningy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

rustfmt, fixed ip hash issue

Browse Source
This commit is contained in:
psun256
2025-12-10 03:03:10 -05:00
parent 90d326ba33
commit 9fb423b949
11 changed files with 150 additions and 239 deletions
Download Patch File Download Diff File Expand all files Collapse all files

173
src/balancer/ip_hashing.rs
View File

@@ -1,11 +1,11 @@
use crate::backend::{Backend, BackendPool, ServerHealth};
use crate::balancer::{Balancer, CURRENT_CONNECTION_INFO, ConnectionInfo};
use std::hash::{Hasher, DefaultHasher, Hash};
use crate::backend::{Backend, BackendPool};
use crate::balancer::{Balancer, ConnectionInfo};
use std::hash::{DefaultHasher, Hash, Hasher};
use std::sync::{Arc, RwLock};
#[derive(Debug)]
pub struct SourceIPHash {
pool : BackendPool,
pool: BackendPool,
}
impl SourceIPHash {
@@ -15,165 +15,84 @@ impl SourceIPHash {
}
impl Balancer for SourceIPHash {
fn choose_backend(&mut self) -> Option<Arc<Backend>>{
let client_ip = CURRENT_CONNECTION_INFO.with(|info| {
info.borrow().as_ref().map(|c| c.client_ip.clone())
});
let client_ip = match client_ip {
Some(ip) => ip,
None => return None, // no client info available
};
fn choose_backend(&mut self, ctx: ConnectionInfo) -> Option<Arc<Backend>> {
let client_ip = ctx.client_ip;
let mut hasher = DefaultHasher::new();
client_ip.hash(&mut hasher);
let hash = hasher.finish();
let idx = (hash as usize) % self.pool.backends.len();
return Some(self.pool.backends[idx].clone());
Some(self.pool.backends[idx].clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::net::IpAddr;
use crate::backend::ServerMetrics;
fn create_dummy_backends(count: usize) -> BackendPool {
let mut backends = Vec::new();
for i in 1..=count {
backends.push(Arc::new(Backend::new(
format!("backend {}", i),
format!("127.0.0.1:808{}", i).parse().unwrap(),
Arc::new(RwLock::new(ServerMetrics::default())),
)));
}
BackendPool::new(backends)
}
#[test]
fn test_same_ip_always_selects_same_backend() {
let backends = vec![
Arc::new(Backend::new(
"backend 1".into(),
"127.0.0.1:8081".parse().unwrap(),
Arc::new(RwLock::new(ServerHealth::default())),
)),
Arc::new(Backend::new(
"backend 2".into(),
"127.0.0.1:8082".parse().unwrap(),
Arc::new(RwLock::new(ServerHealth::default())),
)),
Arc::new(Backend::new(
"backend 3".into(),
"127.0.0.1:8083".parse().unwrap(),
Arc::new(RwLock::new(ServerHealth::default())),
)),
];
let backends = create_dummy_backends(3);
let mut balancer = SourceIPHash::new(backends);
let mut balancer = SourceIPHash::new(BackendPool::new(backends));
let client_ip = "192.168.1.100:54321".parse().unwrap();
let client_ip: IpAddr = "192.168.1.100".parse().unwrap();
CURRENT_CONNECTION_INFO.with(|info| {
*info.borrow_mut() = Some(ConnectionInfo { client_ip });
});
let first_choice = balancer.choose_backend();
CURRENT_CONNECTION_INFO.with(|info| {
*info.borrow_mut() = Some(ConnectionInfo { client_ip });
});
let second_choice = balancer.choose_backend();
let first_choice = balancer.choose_backend(ConnectionInfo { client_ip });
let second_choice = balancer.choose_backend(ConnectionInfo { client_ip });
assert!(first_choice.is_some());
assert!(second_choice.is_some());
let first = first_choice.unwrap();
let second = second_choice.unwrap();
assert_eq!(first.id, second.id);
// Cleanup
CURRENT_CONNECTION_INFO.with(|info| {
*info.borrow_mut() = None;
});
assert_eq!(first.id, second.id);
}
#[test]
fn test_different_ips_may_select_different_backends() {
let backends = vec![
Arc::new(Backend::new(
"backend 1".into(),
"127.0.0.1:8081".parse().unwrap(),
Arc::new(RwLock::new(ServerHealth::default())),
)),
Arc::new(Backend::new(
"backend 2".into(),
"127.0.0.1:8082".parse().unwrap(),
Arc::new(RwLock::new(ServerHealth::default())),
)),
];
let backends = create_dummy_backends(2);
let mut balancer = SourceIPHash::new(backends);
let mut balancer = SourceIPHash::new(BackendPool::new(backends));
let ip1: IpAddr = "192.168.1.100".parse().unwrap();
let choice1 = balancer.choose_backend(ConnectionInfo { client_ip: ip1 });
let ip1 = "192.168.1.100:54321".parse().unwrap();
CURRENT_CONNECTION_INFO.with(|info| {
*info.borrow_mut() = Some(ConnectionInfo { client_ip: ip1 });
});
let choice1 = balancer.choose_backend();
let ip2 = "192.168.1.101:54322".parse().unwrap();
CURRENT_CONNECTION_INFO.with(|info| {
*info.borrow_mut() = Some(ConnectionInfo { client_ip: ip2 });
});
let choice2 = balancer.choose_backend();
let ip2: IpAddr = "192.168.1.101".parse().unwrap();
let choice2 = balancer.choose_backend(ConnectionInfo { client_ip: ip2 });
assert!(choice1.is_some());
assert!(choice2.is_some());
// Note: choice1 and choice2 might be equal by chance, but statistically should differ
CURRENT_CONNECTION_INFO.with(|info| {
*info.borrow_mut() = None;
});
}
#[test]
fn test_returns_none_when_no_connection_info() {
let backends = vec![Arc::new(Backend::new(
"backend 1".into(),
"127.0.0.1:8081".parse().unwrap(),
Arc::new(RwLock::new(ServerHealth::default())),
))];
let mut balancer = SourceIPHash::new(BackendPool::new(backends));
// Don't set any connection info
CURRENT_CONNECTION_INFO.with(|info| {
*info.borrow_mut() = None;
});
let choice = balancer.choose_backend();
assert!(choice.is_none());
}
#[test]
fn test_hash_distribution_across_backends() {
let backends = vec![
Arc::new(Backend::new(
"backend 1".into(),
"127.0.0.1:8081".parse().unwrap(),
Arc::new(RwLock::new(ServerHealth::default())),
)),
Arc::new(Backend::new(
"backend 2".into(),
"127.0.0.1:8082".parse().unwrap(),
Arc::new(RwLock::new(ServerHealth::default())),
)),
Arc::new(Backend::new(
"backend 3".into(),
"127.0.0.1:8083".parse().unwrap(),
Arc::new(RwLock::new(ServerHealth::default())),
)),
];
let pool = create_dummy_backends(3);
let backends_ref = pool.backends.clone();
let mut balancer = SourceIPHash::new(BackendPool::new(backends.clone()));
let mut balancer = SourceIPHash::new(pool);
let mut distribution = [0, 0, 0];
// Test 30 different IPs to see if they distribute across backends
// Test 30 different IPs
for i in 0..30 {
let client_ip = format!("192.168.1.{}:54321", 100 + i).parse().unwrap();
CURRENT_CONNECTION_INFO.with(|info| {
*info.borrow_mut() = Some(ConnectionInfo { client_ip });
});
let client_ip: IpAddr = format!("192.168.1.{}", 100 + i).parse().unwrap();
if let Some(backend) = balancer.choose_backend() {
for (idx, b) in backends.iter().enumerate() {
if let Some(backend) = balancer.choose_backend(ConnectionInfo { client_ip }) {
for (idx, b) in backends_ref.iter().enumerate() {
if backend.id == b.id && backend.address == b.address {
distribution[idx] += 1;
break;
@@ -182,12 +101,8 @@ mod tests {
}
}
assert!(distribution[0] > 0);
assert!(distribution[1] > 0);
assert!(distribution[2] > 0);
CURRENT_CONNECTION_INFO.with(|info| {
*info.borrow_mut() = None;
});
assert!(distribution[0] > 0, "Backend 0 received no traffic");
assert!(distribution[1] > 0, "Backend 1 received no traffic");
assert!(distribution[2] > 0, "Backend 2 received no traffic");
}
}