#![allow(non_snake_case)]
use crate::core::RiResult;
use std::collections::HashMap as FxHashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Instant;
use tokio::sync::RwLock;
use std::sync::RwLock as StdRwLock;
#[cfg(feature = "gateway")]
use hyper;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub enum RiLoadBalancerStrategy {
RoundRobin,
WeightedRoundRobin,
LeastConnections,
Random,
IpHash,
LeastResponseTime,
ConsistentHash,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiBackendServer {
pub id: String,
pub url: String,
pub weight: u32,
pub max_connections: usize,
pub health_check_path: String,
pub is_healthy: bool,
}
impl RiBackendServer {
pub fn new(id: String, url: String) -> Self {
Self {
id,
url,
weight: 1,
max_connections: 100,
health_check_path: "/health".to_string(),
is_healthy: true,
}
}
}
#[cfg_attr(feature = "pyo3", pyo3::prelude::pymethods)]
impl RiBackendServer {
#[cfg(feature = "pyo3")]
#[new]
fn py_new(id: String, url: String) -> Self {
Self::new(id, url)
}
#[cfg(feature = "pyo3")]
fn set_weight(&mut self, weight: u32) {
self.weight = weight;
}
#[cfg(feature = "pyo3")]
fn set_max_connections(&mut self, max_connections: usize) {
self.max_connections = max_connections;
}
#[cfg(feature = "pyo3")]
fn set_health_check_path(&mut self, path: String) {
self.health_check_path = path;
}
}
#[derive(Debug)]
struct ServerStats {
active_connections: AtomicUsize,
total_requests: AtomicUsize,
failed_requests: AtomicUsize,
response_time_ms: AtomicUsize,
last_used: StdRwLock<Instant>,
}
impl ServerStats {
fn new() -> Self {
Self {
active_connections: AtomicUsize::new(0),
total_requests: AtomicUsize::new(0),
failed_requests: AtomicUsize::new(0),
response_time_ms: AtomicUsize::new(0),
last_used: StdRwLock::new(Instant::now()),
}
}
fn get_active_connections(&self) -> usize {
self.active_connections.load(Ordering::Relaxed)
}
fn increment_connections(&self) {
self.active_connections.fetch_add(1, Ordering::Relaxed);
self.total_requests.fetch_add(1, Ordering::Relaxed);
if let Ok(mut last_used) = self.last_used.write() {
*last_used = Instant::now();
}
}
fn decrement_connections(&self) {
self.active_connections.fetch_sub(1, Ordering::Relaxed);
}
fn record_failure(&self) {
self.failed_requests.fetch_add(1, Ordering::Relaxed);
self.decrement_connections();
}
fn record_response_time(&self, response_time_ms: u64) {
self.response_time_ms.store(response_time_ms as usize, Ordering::Relaxed);
}
fn get_stats(&self) -> RiLoadBalancerServerStats {
RiLoadBalancerServerStats {
active_connections: self.get_active_connections(),
total_requests: self.total_requests.load(Ordering::Relaxed),
failed_requests: self.failed_requests.load(Ordering::Relaxed),
response_time_ms: self.response_time_ms.load(Ordering::Relaxed),
}
}
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiLoadBalancerServerStats {
pub active_connections: usize,
pub total_requests: usize,
pub failed_requests: usize,
pub response_time_ms: usize,
}
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiLoadBalancer {
strategy: RiLoadBalancerStrategy,
servers: RwLock<Vec<RiBackendServer>>,
server_stats: RwLock<FxHashMap<String, Arc<ServerStats>>>,
round_robin_counter: AtomicUsize,
}
impl Clone for RiLoadBalancer {
fn clone(&self) -> Self {
Self {
strategy: self.strategy.clone(),
servers: RwLock::new(Vec::new()),
server_stats: RwLock::new(FxHashMap::default()),
round_robin_counter: AtomicUsize::new(self.round_robin_counter.load(Ordering::Relaxed)),
}
}
}
impl RiLoadBalancer {
pub fn new(strategy: RiLoadBalancerStrategy) -> Self {
Self {
strategy,
servers: RwLock::new(Vec::new()),
server_stats: RwLock::new(FxHashMap::default()),
round_robin_counter: AtomicUsize::new(0),
}
}
pub async fn add_server(&self, server: RiBackendServer) -> RiResult<()> {
self.validate_server_url(&server.url)?;
let mut servers = self.servers.write().await;
let mut stats = self.server_stats.write().await;
servers.push(server.clone());
stats.insert(server.id.clone(), Arc::new(ServerStats::new()));
Ok(())
}
fn validate_server_url(&self, url: &str) -> RiResult<()> {
if url.is_empty() || url.len() > 2048 {
return Err(crate::core::RiError::Other(
"Server URL must be 1-2048 characters".to_string()
));
}
let parsed_url = url::Url::parse(url)
.map_err(|e| crate::core::RiError::Other(
format!("Invalid server URL: {}", e)
))?;
let scheme = parsed_url.scheme();
if scheme != "http" && scheme != "https" {
log::warn!(
"[Ri.LoadBalancer] Blocked non-HTTP(S) URL scheme: {} for URL: {}",
scheme, url
);
return Err(crate::core::RiError::Other(
format!("Invalid URL scheme '{}'. Only HTTP and HTTPS are allowed.", scheme)
));
}
if let Some(host) = parsed_url.host_str() {
if host == "localhost" || host == "127.0.0.1" || host == "::1" {
log::warn!(
"[Ri.LoadBalancer] Server URL points to localhost: {}",
url
);
}
if let Ok(ip) = host.parse::<std::net::IpAddr>() {
let is_private_or_link_local = match ip {
std::net::IpAddr::V4(ipv4) => ipv4.is_private() || ipv4.is_link_local(),
std::net::IpAddr::V6(ipv6) => ipv6.is_unicast_link_local(),
};
if ip.is_loopback() || is_private_or_link_local {
log::warn!(
"[Ri.LoadBalancer] Server URL points to private IP {}: {}",
ip, url
);
}
}
}
if parsed_url.username() != "" || parsed_url.password().is_some() {
log::warn!(
"[Ri.LoadBalancer] Server URL contains credentials (not recommended): {}",
url.split('@').last().unwrap_or(url)
);
}
Ok(())
}
pub async fn remove_server(&self, server_id: &str) -> bool {
let mut servers = self.servers.write().await;
let mut stats = self.server_stats.write().await;
let initial_len = servers.len();
servers.retain(|s| s.id != server_id);
stats.remove(server_id);
servers.len() < initial_len
}
pub async fn get_healthy_servers(&self) -> Vec<RiBackendServer> {
let servers = self.servers.read().await;
servers.iter()
.filter(|s| s.is_healthy)
.cloned()
.collect()
}
pub async fn select_server(&self, client_ip: Option<&str>) -> RiResult<RiBackendServer> {
let healthy_servers = self.get_healthy_servers().await;
if healthy_servers.is_empty() {
return Err(crate::core::RiError::Other("No healthy servers available".to_string()));
}
let stats = self.server_stats.read().await;
let available_servers: Vec<RiBackendServer> = healthy_servers.into_iter()
.filter(|server| {
if let Some(server_stats) = stats.get(&server.id) {
let connections = server_stats.get_active_connections();
connections < server.max_connections
} else {
true }
})
.collect();
if available_servers.is_empty() {
return Err(crate::core::RiError::Other("No servers with available connections".to_string()));
}
let server = match self.strategy {
RiLoadBalancerStrategy::RoundRobin => self.select_round_robin(&available_servers).await,
RiLoadBalancerStrategy::WeightedRoundRobin => self.select_weighted_round_robin(&available_servers).await,
RiLoadBalancerStrategy::LeastConnections => self.select_least_connections(&available_servers).await,
RiLoadBalancerStrategy::Random => self.select_random(&available_servers),
RiLoadBalancerStrategy::IpHash => self.select_ip_hash(&available_servers, client_ip),
RiLoadBalancerStrategy::LeastResponseTime => self.select_least_response_time(&available_servers).await,
RiLoadBalancerStrategy::ConsistentHash => self.select_consistent_hash(&available_servers, client_ip),
};
if let Some(server) = server {
if let Some(stats) = self.server_stats.read().await.get(&server.id) {
stats.increment_connections();
}
Ok(server)
} else {
Err(crate::core::RiError::Other("Failed to select server".to_string()))
}
}
async fn select_round_robin(&self, servers: &[RiBackendServer]) -> Option<RiBackendServer> {
let counter = self.round_robin_counter.fetch_add(1, Ordering::Relaxed);
let index = counter % servers.len();
servers.get(index).cloned()
}
async fn select_weighted_round_robin(&self, servers: &[RiBackendServer]) -> Option<RiBackendServer> {
if servers.is_empty() {
return None;
}
let total_weight: u32 = servers.iter().map(|s| s.weight).sum();
let counter = self.round_robin_counter.fetch_add(1, Ordering::Relaxed);
let weighted_index = counter % total_weight as usize;
let mut accumulated_weight = 0;
for server in servers {
accumulated_weight += server.weight as usize;
if weighted_index < accumulated_weight {
return Some(server.clone());
}
}
servers.first().cloned()
}
async fn select_least_connections(&self, servers: &[RiBackendServer]) -> Option<RiBackendServer> {
let stats = self.server_stats.read().await;
let mut best_server = None;
let mut best_score = f64::MAX;
for server in servers {
if let Some(server_stats) = stats.get(&server.id) {
let connections = server_stats.get_active_connections();
if connections >= server.max_connections {
continue;
}
let weight = server.weight as f64 + 0.001;
let score = connections as f64 / weight;
if score < best_score {
best_score = score;
best_server = Some(server.clone());
}
}
}
best_server.or_else(|| servers.first().cloned())
}
fn select_random(&self, servers: &[RiBackendServer]) -> Option<RiBackendServer> {
use rand::Rng;
let mut rng = rand::thread_rng();
let index = rng.gen_range(0..servers.len());
servers.get(index).cloned()
}
async fn select_least_response_time(&self, servers: &[RiBackendServer]) -> Option<RiBackendServer> {
let stats = self.server_stats.read().await;
let mut best_server = None;
let mut min_response_time = u64::MAX;
for server in servers {
if let Some(server_stats) = stats.get(&server.id) {
let response_time = server_stats.response_time_ms.load(Ordering::Relaxed) as u64;
if response_time < min_response_time {
min_response_time = response_time;
best_server = Some(server.clone());
}
}
}
best_server.or_else(|| servers.first().cloned())
}
fn select_ip_hash(&self, servers: &[RiBackendServer], client_ip: Option<&str>) -> Option<RiBackendServer> {
if let Some(ip) = client_ip {
let hash = self.hash_ip(ip);
let index = hash as usize % servers.len();
servers.get(index).cloned()
} else {
self.select_random(servers)
}
}
fn hash_ip(&self, ip: &str) -> u64 {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let mut hasher = DefaultHasher::new();
ip.hash(&mut hasher);
hasher.finish()
}
fn select_consistent_hash(&self, servers: &[RiBackendServer], client_ip: Option<&str>) -> Option<RiBackendServer> {
if servers.is_empty() {
return None;
}
let key = client_ip.unwrap_or("127.0.0.1");
let mut server_hashes: Vec<(u64, RiBackendServer)> = servers
.iter()
.map(|server| {
let hash = self.hash_ip(&server.id);
(hash, server.clone())
})
.collect();
server_hashes.sort_by(|a, b| a.0.cmp(&b.0));
let key_hash = self.hash_ip(key);
for (server_hash, server) in &server_hashes {
if *server_hash >= key_hash {
return Some(server.clone());
}
}
server_hashes.first().map(|(_, server)| server.clone())
}
pub async fn release_server(&self, server_id: &str) {
if let Some(stats) = self.server_stats.read().await.get(server_id) {
stats.decrement_connections();
}
}
pub async fn record_server_failure(&self, server_id: &str) {
if let Some(stats) = self.server_stats.read().await.get(server_id) {
stats.record_failure();
}
let mut servers = self.servers.write().await;
if let Some(server) = servers.iter_mut().find(|s| s.id == server_id) {
if let Some(stats) = self.server_stats.read().await.get(server_id) {
let total = stats.total_requests.load(Ordering::Relaxed);
let failed = stats.failed_requests.load(Ordering::Relaxed);
if total > 10 && (failed as f64 / total as f64) > 0.5 {
server.is_healthy = false;
}
}
}
}
pub async fn record_response_time(&self, server_id: &str, response_time_ms: u64) {
if let Some(stats) = self.server_stats.read().await.get(server_id) {
stats.record_response_time(response_time_ms);
}
}
pub async fn get_server_stats(&self, server_id: &str) -> Option<RiLoadBalancerServerStats> {
self.server_stats.read().await
.get(server_id)
.map(|stats| stats.get_stats())
}
pub async fn get_all_stats(&self) -> FxHashMap<String, RiLoadBalancerServerStats> {
let stats = self.server_stats.read().await;
let mut result = FxHashMap::default();
for (server_id, server_stats) in stats.iter() {
result.insert(server_id.clone(), server_stats.get_stats());
}
result
}
pub async fn mark_server_healthy(&self, server_id: &str, healthy: bool) {
let mut servers = self.servers.write().await;
if let Some(server) = servers.iter_mut().find(|s| s.id == server_id) {
server.is_healthy = healthy;
}
}
#[cfg(feature = "gateway")]
pub async fn perform_health_check(&self, server_id: &str) -> bool {
let servers = self.servers.read().await;
if let Some(server) = servers.iter().find(|s| s.id == server_id) {
let health_check_url = format!("{}{}", server.url, server.health_check_path);
let uri = match hyper::Uri::from_maybe_shared(health_check_url.clone()) {
Ok(uri) => uri,
Err(e) => {
if let Ok(fs) = crate::fs::RiFileSystem::new_auto_root() {
let logger = crate::log::RiLogger::new(&crate::log::RiLogConfig::default(), fs);
let _ = logger.warn("load_balancer", format!("Invalid health check URL for server {server_id}: {e}"));
}
return false;
}
};
match hyper::Client::new().get(uri).await {
Ok(response) => {
(200..300).contains(&response.status().as_u16())
},
Err(_) => false,
}
} else {
false
}
}
#[cfg(not(feature = "gateway"))]
pub async fn perform_health_check(&self, _server_id: &str) -> bool {
true
}
pub async fn start_health_checks(self: Arc<Self>, interval_secs: u64) -> tokio::task::JoinHandle<()> {
let this = self.clone();
tokio::spawn(async move {
let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(interval_secs));
loop {
interval.tick().await;
let servers = this.servers.read().await;
let server_ids: Vec<String> = servers.iter().map(|s| s.id.clone()).collect();
for server_id in server_ids {
let is_healthy = this.perform_health_check(&server_id).await;
let _ = this.mark_server_healthy(&server_id, is_healthy).await;
if !is_healthy {
let _ = this.record_server_failure(&server_id).await;
}
}
}
})
}
pub fn get_strategy(&self) -> &RiLoadBalancerStrategy {
&self.strategy
}
pub async fn set_strategy(&mut self, strategy: RiLoadBalancerStrategy) {
self.strategy = strategy;
}
pub async fn get_server_count(&self) -> usize {
self.servers.read().await.len()
}
pub async fn get_healthy_server_count(&self) -> usize {
self.get_healthy_servers().await.len()
}
}