Skip to main content

ri/gateway/
load_balancer.rs

1//! Copyright © 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18#![allow(non_snake_case)]
19
20//! # Load Balancer Module
21//! 
22//! This module provides a robust load balancer implementation for distributing incoming requests
23//! across multiple backend servers. It supports various load balancing strategies and includes
24//! health checking, connection management, and detailed statistics.
25//! 
26//! ## Key Components
27//! 
28//! - **RiLoadBalancerStrategy**: Enum representing different load balancing algorithms
29//! - **RiBackendServer**: Represents a backend server with configuration and health status
30//! - **RiLoadBalancer**: Main load balancer implementation
31//! - **RiLoadBalancerServerStats**: Metrics for monitoring server performance
32//! 
33//! ## Design Principles
34//! 
35//! 1. **Multiple Strategies**: Supports RoundRobin, WeightedRoundRobin, LeastConnections, Random, and IpHash
36//! 2. **Health Checking**: Automatic periodic health checks to ensure traffic is only sent to healthy servers
37//! 3. **Connection Management**: Tracks active connections and enforces max connections per server
38//! 4. **Detailed Statistics**: Collects metrics on requests, failures, and response times
39//! 5. **Thread Safety**: Uses Arc and RwLock for safe operation in multi-threaded environments
40//! 6. **Scalability**: Designed to handle large numbers of servers and requests
41//! 7. **Configurable**: Allows fine-tuning of server weights, max connections, and health check paths
42//! 8. **Async Compatibility**: Built with async/await patterns for modern Rust applications
43//! 
44//! ## Usage
45//! 
46//! ```rust
47//! use ri::prelude::*;
48//! use std::sync::Arc;
49//! 
50//! async fn example() -> RiResult<()> {
51//!     // Create a load balancer with Round Robin strategy
52//!     let lb = Arc::new(RiLoadBalancer::new(RiLoadBalancerStrategy::RoundRobin));
53//!     
54//!     // Add backend servers
55//!     lb.add_server(RiBackendServer::new("server1".to_string(), "http://localhost:8081".to_string())
56//!         .with_weight(2)
57//!         .with_max_connections(200))
58//!         .await;
59//!     
60//!     lb.add_server(RiBackendServer::new("server2".to_string(), "http://localhost:8082".to_string())
61//!         .with_weight(1)
62//!         .with_max_connections(100))
63//!         .await;
64//!     
65//!     // Start periodic health checks every 30 seconds
66//!     lb.clone().start_health_checks(30).await;
67//!     
68//!     // Select a server for a client request
69//!     let server = lb.select_server(Some("192.168.1.1")).await?;
70//!     println!("Selected server: {}", server.url);
71//!     
72//!     // Record response time when done
73//!     lb.record_response_time(&server.id, 150).await;
74//!     
75//!     // Release the server when the request is complete
76//!     lb.release_server(&server.id).await;
77//!     
78//!     // Get server statistics
79//!     let stats = lb.get_all_stats().await;
80//!     println!("Server stats: {:?}", stats);
81//!     
82//!     Ok(())
83//! }
84//! ```
85
86use crate::core::RiResult;
87use std::collections::HashMap as FxHashMap;
88use std::sync::Arc;
89use std::sync::atomic::{AtomicUsize, Ordering};
90use std::time::Instant;
91use tokio::sync::RwLock;
92use std::sync::RwLock as StdRwLock;
93
94#[cfg(feature = "gateway")]
95use hyper;
96
97/// Load balancing strategies supported by Ri.
98/// 
99/// These strategies determine how the load balancer selects which backend server
100/// to route incoming requests to.
101#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
102#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
103pub enum RiLoadBalancerStrategy {
104    /// **Round Robin**: Sequentially selects the next available server in rotation.
105    /// 
106    /// Simple and fair distribution, ideal for servers with similar capabilities.
107    RoundRobin,
108    
109    /// **Weighted Round Robin**: Selects servers based on assigned weights.
110    /// 
111    /// Allows more powerful servers to handle a larger share of traffic.
112    WeightedRoundRobin,
113    
114    /// **Least Connections**: Selects the server with the fewest active connections.
115    /// 
116    /// Ideal for handling varying request durations, ensuring balanced load.
117    LeastConnections,
118    
119    /// **Random**: Randomly selects an available server.
120    /// 
121    /// Simple implementation with good distribution characteristics.
122    Random,
123    
124    /// **IP Hash**: Uses client IP address to consistently route to the same server.
125    /// 
126    /// Maintains session persistence by mapping clients to specific servers.
127    IpHash,
128    
129    /// **Least Response Time**: Selects the server with the lowest average response time.
130    /// 
131    /// Ideal for optimizing user experience by directing traffic to the fastest servers.
132    LeastResponseTime,
133    
134    /// **Consistent Hash**: Uses a consistent hashing algorithm for server selection.
135    /// 
136    /// Provides stable mapping between requests and servers, minimizing disruption when servers are added or removed.
137    ConsistentHash,
138}
139
140/// Represents a backend server in the load balancer.
141/// 
142/// This struct contains all the configuration and state information for a backend server,
143/// including its ID, URL, weight, max connections, health check path, and health status.
144#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
145#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
146pub struct RiBackendServer {
147    /// Unique identifier for the server
148    pub id: String,
149    
150    /// Base URL of the server (e.g., "http://localhost:8080")
151    pub url: String,
152    
153    /// Weight assigned to the server for weighted load balancing strategies
154    pub weight: u32,
155    
156    /// Maximum number of concurrent connections allowed to this server
157    pub max_connections: usize,
158    
159    /// Path to check for health status (e.g., "/health")
160    pub health_check_path: String,
161    
162    /// Current health status of the server (true = healthy, false = unhealthy)
163    pub is_healthy: bool,
164}
165
166impl RiBackendServer {
167    /// Creates a new backend server with the specified ID and URL.
168    /// 
169    /// # Parameters
170    /// 
171    /// - `id`: Unique identifier for the server
172    /// - `url`: Base URL of the server
173    /// 
174    /// # Returns
175    /// 
176    /// A new `RiBackendServer` instance with default values
177    pub fn new(id: String, url: String) -> Self {
178        Self {
179            id,
180            url,
181            weight: 1,
182            max_connections: 100,
183            health_check_path: "/health".to_string(),
184            is_healthy: true,
185        }
186    }
187}
188
189#[cfg_attr(feature = "pyo3", pyo3::prelude::pymethods)]
190impl RiBackendServer {
191    #[cfg(feature = "pyo3")]
192    #[new]
193    fn py_new(id: String, url: String) -> Self {
194        Self::new(id, url)
195    }
196
197    #[cfg(feature = "pyo3")]
198    fn set_weight(&mut self, weight: u32) {
199        self.weight = weight;
200    }
201
202    #[cfg(feature = "pyo3")]
203    fn set_max_connections(&mut self, max_connections: usize) {
204        self.max_connections = max_connections;
205    }
206
207    #[cfg(feature = "pyo3")]
208    fn set_health_check_path(&mut self, path: String) {
209        self.health_check_path = path;
210    }
211}
212
213/// Internal server statistics tracking.
214/// 
215/// This struct tracks real-time statistics for each backend server, including active connections,
216/// request counts, failures, and response times. It is designed to be thread-safe for use in
217/// multi-threaded environments.
218#[derive(Debug)]
219struct ServerStats {
220    /// Number of currently active connections to the server
221    active_connections: AtomicUsize,
222    
223    /// Total number of requests sent to the server since it was added
224    total_requests: AtomicUsize,
225    
226    /// Number of failed requests to the server
227    failed_requests: AtomicUsize,
228    
229    /// Most recent response time in milliseconds
230    response_time_ms: AtomicUsize,
231    
232    /// Timestamp of when the server was last used
233    last_used: StdRwLock<Instant>,
234}
235
236impl ServerStats {
237    /// Creates a new server statistics instance with default values.
238    /// 
239    /// Initializes all counters to zero and sets the last used time to now.
240    /// 
241    /// # Returns
242    /// 
243    /// A new `ServerStats` instance with default values
244    fn new() -> Self {
245        Self {
246            active_connections: AtomicUsize::new(0),
247            total_requests: AtomicUsize::new(0),
248            failed_requests: AtomicUsize::new(0),
249            response_time_ms: AtomicUsize::new(0),
250            last_used: StdRwLock::new(Instant::now()),
251        }
252    }
253
254    /// Gets the current number of active connections to the server.
255    /// 
256    /// # Returns
257    /// 
258    /// The number of active connections as a `usize`
259    fn get_active_connections(&self) -> usize {
260        self.active_connections.load(Ordering::Relaxed)
261    }
262
263    /// Increments the active connection count and updates request statistics.
264    /// 
265    /// This method should be called when a new connection is established to the server.
266    /// 
267    /// - Increments active_connections by 1
268    /// - Increments total_requests by 1
269    /// - Updates last_used to the current time
270    fn increment_connections(&self) {
271        self.active_connections.fetch_add(1, Ordering::Relaxed);
272        self.total_requests.fetch_add(1, Ordering::Relaxed);
273        if let Ok(mut last_used) = self.last_used.write() {
274            *last_used = Instant::now();
275        }
276    }
277
278    /// Decrements the active connection count.
279    /// 
280    /// This method should be called when a connection to the server is closed.
281    fn decrement_connections(&self) {
282        self.active_connections.fetch_sub(1, Ordering::Relaxed);
283    }
284
285    /// Records a failed request to the server.
286    /// 
287    /// This method should be called when a request to the server fails.
288    /// 
289    /// - Increments failed_requests by 1
290    /// - Decrements active_connections by 1 (since the connection failed)
291    fn record_failure(&self) {
292        self.failed_requests.fetch_add(1, Ordering::Relaxed);
293        self.decrement_connections();
294    }
295
296    /// Records the response time for a successful request.
297    /// 
298    /// This method should be called when a request to the server completes successfully.
299    /// 
300    /// # Parameters
301    /// 
302    /// - `response_time_ms`: Response time in milliseconds
303    fn record_response_time(&self, response_time_ms: u64) {
304        self.response_time_ms.store(response_time_ms as usize, Ordering::Relaxed);
305    }
306
307    /// Gets a snapshot of the current server statistics.
308    ///
309    /// This method converts the internal statistics into a public-facing `RiLoadBalancerServerStats` struct.
310    ///
311    /// # Returns
312    ///
313    /// A `RiLoadBalancerServerStats` struct containing the current statistics
314    fn get_stats(&self) -> RiLoadBalancerServerStats {
315        RiLoadBalancerServerStats {
316            active_connections: self.get_active_connections(),
317            total_requests: self.total_requests.load(Ordering::Relaxed),
318            failed_requests: self.failed_requests.load(Ordering::Relaxed),
319            response_time_ms: self.response_time_ms.load(Ordering::Relaxed),
320        }
321    }
322}
323
324/// Load balancer server statistics for monitoring and reporting.
325///
326/// This struct contains metrics for a backend server, providing insights into its
327/// performance, load, and reliability.
328#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
329#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
330pub struct RiLoadBalancerServerStats {
331    /// Number of currently active connections to the server
332    pub active_connections: usize,
333
334    /// Total number of requests sent to the server since it was added
335    pub total_requests: usize,
336
337    /// Number of failed requests to the server
338    pub failed_requests: usize,
339
340    /// Most recent response time in milliseconds
341    pub response_time_ms: usize,
342}
343
344/// Main load balancer implementation.
345/// 
346/// This struct provides a comprehensive load balancing solution with support for multiple
347/// strategies, health checking, connection management, and detailed statistics.
348#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
349pub struct RiLoadBalancer {
350    /// Load balancing strategy to use
351    strategy: RiLoadBalancerStrategy,
352    
353    /// List of backend servers
354    servers: RwLock<Vec<RiBackendServer>>,
355    
356    /// Statistics for each backend server
357    server_stats: RwLock<FxHashMap<String, Arc<ServerStats>>>,
358    
359    /// Counter for round robin scheduling
360    round_robin_counter: AtomicUsize,
361}
362
363impl Clone for RiLoadBalancer {
364    /// Creates a clone of the load balancer.
365    /// 
366    /// Note: The clone will have the same strategy and counter, but empty servers and stats
367    /// since we can't await in the Clone trait.
368    fn clone(&self) -> Self {
369        Self {
370            strategy: self.strategy.clone(),
371            servers: RwLock::new(Vec::new()),
372            server_stats: RwLock::new(FxHashMap::default()),
373            round_robin_counter: AtomicUsize::new(self.round_robin_counter.load(Ordering::Relaxed)),
374        }
375    }
376}
377
378impl RiLoadBalancer {
379    /// Creates a new load balancer with the specified strategy.
380    /// 
381    /// # Parameters
382    /// 
383    /// - `strategy`: The load balancing strategy to use
384    /// 
385    /// # Returns
386    /// 
387    /// A new `RiLoadBalancer` instance with the specified strategy
388    pub fn new(strategy: RiLoadBalancerStrategy) -> Self {
389        Self {
390            strategy,
391            servers: RwLock::new(Vec::new()),
392            server_stats: RwLock::new(FxHashMap::default()),
393            round_robin_counter: AtomicUsize::new(0),
394        }
395    }
396
397    /// Adds a backend server to the load balancer.
398    /// 
399    /// # Parameters
400    /// 
401    /// - `server`: The backend server to add
402    /// 
403    /// # Security
404    /// 
405    /// This method validates the server URL to prevent SSRF (Server-Side Request Forgery) attacks.
406    /// Only HTTP and HTTPS URLs are allowed. File://, gopher://, and other schemes are blocked.
407    /// Private IP addresses (127.0.0.0/8, 10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 169.254.0.0/16)
408    /// are logged as warnings but still allowed for internal deployments.
409    pub async fn add_server(&self, server: RiBackendServer) -> RiResult<()> {
410        // Security: Validate the server URL
411        self.validate_server_url(&server.url)?;
412        
413        let mut servers = self.servers.write().await;
414        let mut stats = self.server_stats.write().await;
415        
416        servers.push(server.clone());
417        stats.insert(server.id.clone(), Arc::new(ServerStats::new()));
418        
419        Ok(())
420    }
421
422    /// Validates a server URL to prevent SSRF attacks.
423    ///
424    /// # Security
425    ///
426    /// This method validates:
427    /// 1. URL scheme must be HTTP or HTTPS
428    /// 2. URL must be well-formed
429    /// 3. Warns if URL points to private IP address
430    fn validate_server_url(&self, url: &str) -> RiResult<()> {
431        // Security: Check URL length
432        if url.is_empty() || url.len() > 2048 {
433            return Err(crate::core::RiError::Other(
434                "Server URL must be 1-2048 characters".to_string()
435            ));
436        }
437
438        // Security: Parse and validate URL
439        let parsed_url = url::Url::parse(url)
440            .map_err(|e| crate::core::RiError::Other(
441                format!("Invalid server URL: {}", e)
442            ))?;
443
444        // Security: Only allow HTTP and HTTPS schemes
445        let scheme = parsed_url.scheme();
446        if scheme != "http" && scheme != "https" {
447            log::warn!(
448                "[Ri.LoadBalancer] Blocked non-HTTP(S) URL scheme: {} for URL: {}",
449                scheme, url
450            );
451            return Err(crate::core::RiError::Other(
452                format!("Invalid URL scheme '{}'. Only HTTP and HTTPS are allowed.", scheme)
453            ));
454        }
455
456        // Security: Check for private IP addresses (informational warning)
457        if let Some(host) = parsed_url.host_str() {
458            // Check for localhost
459            if host == "localhost" || host == "127.0.0.1" || host == "::1" {
460                log::warn!(
461                    "[Ri.LoadBalancer] Server URL points to localhost: {}",
462                    url
463                );
464            }
465            
466            // Check for private IP ranges
467            if let Ok(ip) = host.parse::<std::net::IpAddr>() {
468                let is_private_or_link_local = match ip {
469                    std::net::IpAddr::V4(ipv4) => ipv4.is_private() || ipv4.is_link_local(),
470                    std::net::IpAddr::V6(ipv6) => ipv6.is_unicast_link_local(),
471                };
472                if ip.is_loopback() || is_private_or_link_local {
473                    log::warn!(
474                        "[Ri.LoadBalancer] Server URL points to private IP {}: {}",
475                        ip, url
476                    );
477                }
478            }
479        }
480
481        // Security: Check for credentials in URL
482        if parsed_url.username() != "" || parsed_url.password().is_some() {
483            log::warn!(
484                "[Ri.LoadBalancer] Server URL contains credentials (not recommended): {}",
485                url.split('@').last().unwrap_or(url)
486            );
487        }
488
489        Ok(())
490    }
491
492    /// Removes a backend server from the load balancer.
493    /// 
494    /// # Parameters
495    /// 
496    /// - `server_id`: The ID of the server to remove
497    /// 
498    /// # Returns
499    /// 
500    /// `true` if the server was removed, `false` otherwise
501    pub async fn remove_server(&self, server_id: &str) -> bool {
502        let mut servers = self.servers.write().await;
503        let mut stats = self.server_stats.write().await;
504        
505        let initial_len = servers.len();
506        servers.retain(|s| s.id != server_id);
507        stats.remove(server_id);
508        
509        servers.len() < initial_len
510    }
511
512    /// Gets a list of all healthy backend servers.
513    /// 
514    /// # Returns
515    /// 
516    /// A vector of healthy `RiBackendServer` instances
517    pub async fn get_healthy_servers(&self) -> Vec<RiBackendServer> {
518        let servers = self.servers.read().await;
519        servers.iter()
520            .filter(|s| s.is_healthy)
521            .cloned()
522            .collect()
523    }
524
525    /// Selects the most appropriate backend server for a client request.
526    /// 
527    /// This method applies the configured load balancing strategy to select a server,
528    /// considering only healthy servers with available connections.
529    /// 
530    /// # Parameters
531    /// 
532    /// - `client_ip`: Optional client IP address for IP Hash strategy
533    /// 
534    /// # Returns
535    /// 
536    /// A `RiResult<RiBackendServer>` with the selected server, or an error if no servers are available
537    pub async fn select_server(&self, client_ip: Option<&str>) -> RiResult<RiBackendServer> {
538        let healthy_servers = self.get_healthy_servers().await;
539        
540        if healthy_servers.is_empty() {
541            return Err(crate::core::RiError::Other("No healthy servers available".to_string()));
542        }
543        
544        let stats = self.server_stats.read().await;
545        
546        // Filter servers that have available connections
547        let available_servers: Vec<RiBackendServer> = healthy_servers.into_iter()
548            .filter(|server| {
549                if let Some(server_stats) = stats.get(&server.id) {
550                    let connections = server_stats.get_active_connections();
551                    connections < server.max_connections
552                } else {
553                    true // If no stats, assume server is available
554                }
555            })
556            .collect();
557        
558        if available_servers.is_empty() {
559            return Err(crate::core::RiError::Other("No servers with available connections".to_string()));
560        }
561
562        let server = match self.strategy {
563            RiLoadBalancerStrategy::RoundRobin => self.select_round_robin(&available_servers).await,
564            RiLoadBalancerStrategy::WeightedRoundRobin => self.select_weighted_round_robin(&available_servers).await,
565            RiLoadBalancerStrategy::LeastConnections => self.select_least_connections(&available_servers).await,
566            RiLoadBalancerStrategy::Random => self.select_random(&available_servers),
567            RiLoadBalancerStrategy::IpHash => self.select_ip_hash(&available_servers, client_ip),
568            RiLoadBalancerStrategy::LeastResponseTime => self.select_least_response_time(&available_servers).await,
569            RiLoadBalancerStrategy::ConsistentHash => self.select_consistent_hash(&available_servers, client_ip),
570        };
571
572        if let Some(server) = server {
573            // Increment connection count
574            if let Some(stats) = self.server_stats.read().await.get(&server.id) {
575                stats.increment_connections();
576            }
577            Ok(server)
578        } else {
579            Err(crate::core::RiError::Other("Failed to select server".to_string()))
580        }
581    }
582
583    /// Selects a server using the Round Robin strategy.
584    /// 
585    /// This method sequentially selects the next available server in rotation.
586    /// 
587    /// # Parameters
588    /// 
589    /// - `servers`: List of available servers
590    /// 
591    /// # Returns
592    /// 
593    /// The selected server, or `None` if no servers are available
594    async fn select_round_robin(&self, servers: &[RiBackendServer]) -> Option<RiBackendServer> {
595        let counter = self.round_robin_counter.fetch_add(1, Ordering::Relaxed);
596        let index = counter % servers.len();
597        servers.get(index).cloned()
598    }
599
600    /// Selects a server using the Smooth Weighted Round Robin strategy.
601    /// 
602    /// This method uses a smooth weighted round robin algorithm to distribute traffic more evenly
603    /// across servers, avoiding the problem of sudden traffic spikes to high-weight servers.
604    /// 
605    /// # Parameters
606    /// 
607    /// - `servers`: List of available servers with weights
608    /// 
609    /// # Returns
610    /// 
611    /// The selected server, or `None` if no servers are available
612    async fn select_weighted_round_robin(&self, servers: &[RiBackendServer]) -> Option<RiBackendServer> {
613        if servers.is_empty() {
614            return None;
615        }
616        
617        // Simple weighted round robin implementation with improved distribution
618        let total_weight: u32 = servers.iter().map(|s| s.weight).sum();
619        let counter = self.round_robin_counter.fetch_add(1, Ordering::Relaxed);
620        let weighted_index = counter % total_weight as usize;
621        
622        let mut accumulated_weight = 0;
623        for server in servers {
624            accumulated_weight += server.weight as usize;
625            if weighted_index < accumulated_weight {
626                return Some(server.clone());
627            }
628        }
629        
630        servers.first().cloned()
631    }
632
633    /// Selects a server using the Weighted Least Connections strategy.
634    /// 
635    /// This method selects the server with the fewest active connections relative to its weight,
636    /// ensuring balanced load across servers with different capacities.
637    /// 
638    /// # Parameters
639    /// 
640    /// - `servers`: List of available servers
641    /// 
642    /// # Returns
643    /// 
644    /// The selected server, or `None` if no servers are available
645    async fn select_least_connections(&self, servers: &[RiBackendServer]) -> Option<RiBackendServer> {
646        let stats = self.server_stats.read().await;
647        
648        let mut best_server = None;
649        let mut best_score = f64::MAX; // Lower score is better (connections per weight unit)
650        
651        for server in servers {
652            if let Some(server_stats) = stats.get(&server.id) {
653                let connections = server_stats.get_active_connections();
654                
655                // Skip servers that have reached max connections
656                if connections >= server.max_connections {
657                    continue;
658                }
659                
660                // Calculate score as connections per weight unit
661                // Use a small epsilon to avoid division by zero
662                let weight = server.weight as f64 + 0.001;
663                let score = connections as f64 / weight;
664                
665                if score < best_score {
666                    best_score = score;
667                    best_server = Some(server.clone());
668                }
669            }
670        }
671        
672        best_server.or_else(|| servers.first().cloned())
673    }
674
675    /// Selects a server using the Random strategy.
676    /// 
677    /// This method randomly selects an available server, providing good distribution characteristics.
678    /// 
679    /// # Parameters
680    /// 
681    /// - `servers`: List of available servers
682    /// 
683    /// # Returns
684    /// 
685    /// The selected server, or `None` if no servers are available
686    fn select_random(&self, servers: &[RiBackendServer]) -> Option<RiBackendServer> {
687        use rand::Rng;
688        let mut rng = rand::thread_rng();
689        let index = rng.gen_range(0..servers.len());
690        servers.get(index).cloned()
691    }
692
693    /// Selects a server using the Least Response Time strategy.
694    /// 
695    /// This method selects the server with the lowest response time, optimizing for user experience.
696    /// 
697    /// # Parameters
698    /// 
699    /// - `servers`: List of available servers
700    /// 
701    /// # Returns
702    /// 
703    /// The selected server, or `None` if no servers are available
704    async fn select_least_response_time(&self, servers: &[RiBackendServer]) -> Option<RiBackendServer> {
705        let stats = self.server_stats.read().await;
706        
707        let mut best_server = None;
708        let mut min_response_time = u64::MAX;
709        
710        for server in servers {
711            if let Some(server_stats) = stats.get(&server.id) {
712                let response_time = server_stats.response_time_ms.load(Ordering::Relaxed) as u64;
713                if response_time < min_response_time {
714                    min_response_time = response_time;
715                    best_server = Some(server.clone());
716                }
717            }
718        }
719        
720        best_server.or_else(|| servers.first().cloned())
721    }
722    
723    /// Selects a server using the Consistent Hash strategy.
724    /// 
725    /// This method uses a consistent hashing algorithm to map requests to servers, minimizing
726    /// disruption when servers are added or removed.
727    /// 
728    /// # Parameters
729    /// 
730    /// - `servers`: List of available servers
731    /// - `client_ip`: Optional client IP address for hashing
732    /// 
733    /// # Returns
734    /// 
735    /// The selected server, or `None` if no servers are available
736    fn select_ip_hash(&self, servers: &[RiBackendServer], client_ip: Option<&str>) -> Option<RiBackendServer> {
737        if let Some(ip) = client_ip {
738            let hash = self.hash_ip(ip);
739            let index = hash as usize % servers.len();
740            servers.get(index).cloned()
741        } else {
742            self.select_random(servers)
743        }
744    }
745
746    /// Hashes an IP address for the IP Hash strategy.
747    /// 
748    /// # Parameters
749    /// 
750    /// - `ip`: IP address to hash
751    /// 
752    /// # Returns
753    /// 
754    /// A 64-bit hash value of the IP address
755    fn hash_ip(&self, ip: &str) -> u64 {
756        use std::collections::hash_map::DefaultHasher;
757        use std::hash::{Hash, Hasher};
758        
759        let mut hasher = DefaultHasher::new();
760        ip.hash(&mut hasher);
761        hasher.finish()
762    }
763    
764    /// Selects a server using the Consistent Hash strategy.
765    /// 
766    /// This method uses a consistent hashing algorithm to map requests to servers, minimizing
767    /// disruption when servers are added or removed.
768    /// 
769    /// # Parameters
770    /// 
771    /// - `servers`: List of available servers
772    /// - `client_ip`: Optional client IP address for hashing
773    /// 
774    /// # Returns
775    /// 
776    /// The selected server, or `None` if no servers are available
777    fn select_consistent_hash(&self, servers: &[RiBackendServer], client_ip: Option<&str>) -> Option<RiBackendServer> {
778        if servers.is_empty() {
779            return None;
780        }
781        
782        let key = client_ip.unwrap_or("127.0.0.1");
783        
784        // Create a sorted list of server hashes
785        let mut server_hashes: Vec<(u64, RiBackendServer)> = servers
786            .iter()
787            .map(|server| {
788                let hash = self.hash_ip(&server.id);
789                (hash, server.clone())
790            })
791            .collect();
792        
793        // Sort server hashes
794        server_hashes.sort_by(|a, b| a.0.cmp(&b.0));
795        
796        // Calculate hash for the key
797        let key_hash = self.hash_ip(key);
798        
799        // Find the first server with hash >= key_hash
800        for (server_hash, server) in &server_hashes {
801            if *server_hash >= key_hash {
802                return Some(server.clone());
803            }
804        }
805        
806        // If no server with hash >= key_hash, return the first server
807        server_hashes.first().map(|(_, server)| server.clone())
808    }
809
810    /// Releases a server after a request is completed.
811    /// 
812    /// This method decrements the active connection count for the specified server.
813    /// 
814    /// # Parameters
815    /// 
816    /// - `server_id`: ID of the server to release
817    pub async fn release_server(&self, server_id: &str) {
818        if let Some(stats) = self.server_stats.read().await.get(server_id) {
819            stats.decrement_connections();
820        }
821    }
822
823    /// Records a failed request to a server.
824    /// 
825    /// This method increments the failed request count and may mark the server as unhealthy
826    /// if the failure rate exceeds a threshold.
827    /// 
828    /// # Parameters
829    /// 
830    /// - `server_id`: ID of the server that failed
831    pub async fn record_server_failure(&self, server_id: &str) {
832        if let Some(stats) = self.server_stats.read().await.get(server_id) {
833            stats.record_failure();
834        }
835        
836        // Mark server as unhealthy if too many failures
837        let mut servers = self.servers.write().await;
838        if let Some(server) = servers.iter_mut().find(|s| s.id == server_id) {
839            // Simple heuristic: mark unhealthy if failure rate > 50%
840            if let Some(stats) = self.server_stats.read().await.get(server_id) {
841                let total = stats.total_requests.load(Ordering::Relaxed);
842                let failed = stats.failed_requests.load(Ordering::Relaxed);
843                
844                if total > 10 && (failed as f64 / total as f64) > 0.5 {
845                    server.is_healthy = false;
846                }
847            }
848        }
849    }
850
851    /// Records the response time for a successful request.
852    /// 
853    /// # Parameters
854    /// 
855    /// - `server_id`: ID of the server that handled the request
856    /// - `response_time_ms`: Response time in milliseconds
857    pub async fn record_response_time(&self, server_id: &str, response_time_ms: u64) {
858        if let Some(stats) = self.server_stats.read().await.get(server_id) {
859            stats.record_response_time(response_time_ms);
860        }
861    }
862
863    /// Gets statistics for a specific server.
864    /// 
865    /// # Parameters
866    /// 
867    /// - `server_id`: ID of the server to get statistics for
868    /// 
869    /// # Returns
870    /// 
871    /// An `Option<RiLoadBalancerServerStats>` with the server statistics, or `None` if the server doesn't exist
872    pub async fn get_server_stats(&self, server_id: &str) -> Option<RiLoadBalancerServerStats> {
873        self.server_stats.read().await
874            .get(server_id)
875            .map(|stats| stats.get_stats())
876    }
877
878    /// Gets statistics for all servers.
879    ///
880    /// # Returns
881    ///
882    /// A `FxHashMap<String, RiLoadBalancerServerStats>` with statistics for all servers
883    pub async fn get_all_stats(&self) -> FxHashMap<String, RiLoadBalancerServerStats> {
884        let stats = self.server_stats.read().await;
885        let mut result = FxHashMap::default();
886        
887        for (server_id, server_stats) in stats.iter() {
888            result.insert(server_id.clone(), server_stats.get_stats());
889        }
890        
891        result
892    }
893
894    /// Marks a server as healthy or unhealthy.
895    /// 
896    /// # Parameters
897    /// 
898    /// - `server_id`: ID of the server to update
899    /// - `healthy`: New health status (true = healthy, false = unhealthy)
900    pub async fn mark_server_healthy(&self, server_id: &str, healthy: bool) {
901        let mut servers = self.servers.write().await;
902        if let Some(server) = servers.iter_mut().find(|s| s.id == server_id) {
903            server.is_healthy = healthy;
904        }
905    }
906
907    /// Performs an HTTP health check on a server.
908    /// 
909    /// This method sends an HTTP GET request to the server's health check path and
910    /// considers the server healthy if it returns a 2xx status code.
911    /// 
912    /// # Parameters
913    /// 
914    /// - `server_id`: ID of the server to check
915    /// 
916    /// # Returns
917    /// 
918    /// `true` if the server is healthy, `false` otherwise
919    #[cfg(feature = "gateway")]
920    pub async fn perform_health_check(&self, server_id: &str) -> bool {
921        let servers = self.servers.read().await;
922        
923        if let Some(server) = servers.iter().find(|s| s.id == server_id) {
924            let health_check_url = format!("{}{}", server.url, server.health_check_path);
925            
926            let uri = match hyper::Uri::from_maybe_shared(health_check_url.clone()) {
927                Ok(uri) => uri,
928                Err(e) => {
929                    // Log warning for invalid health check URL
930        if let Ok(fs) = crate::fs::RiFileSystem::new_auto_root() {
931            let logger = crate::log::RiLogger::new(&crate::log::RiLogConfig::default(), fs);
932            let _ = logger.warn("load_balancer", format!("Invalid health check URL for server {server_id}: {e}"));
933        }
934                    return false;
935                }
936            };
937            
938            match hyper::Client::new().get(uri).await {
939                Ok(response) => {
940                    (200..300).contains(&response.status().as_u16())
941                },
942                Err(_) => false,
943            }
944        } else {
945            false
946        }
947    }
948    
949    #[cfg(not(feature = "gateway"))]
950    pub async fn perform_health_check(&self, _server_id: &str) -> bool {
951        // If gateway feature is not enabled, assume all servers are healthy
952        true
953    }
954    
955    /// Starts periodic health checks for all servers.
956    /// 
957    /// This method spawns a background task that performs health checks on all servers
958    /// at the specified interval.
959    /// 
960    /// # Parameters
961    /// 
962    /// - `interval_secs`: Interval between health checks in seconds
963    /// 
964    /// # Returns
965    /// 
966    /// A `tokio::task::JoinHandle` for the background health check task
967    pub async fn start_health_checks(self: Arc<Self>, interval_secs: u64) -> tokio::task::JoinHandle<()> {
968        let this = self.clone();
969        
970        tokio::spawn(async move {
971            let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(interval_secs));
972            
973            loop {
974                interval.tick().await;
975                
976                let servers = this.servers.read().await;
977                let server_ids: Vec<String> = servers.iter().map(|s| s.id.clone()).collect();
978                
979                for server_id in server_ids {
980                    let is_healthy = this.perform_health_check(&server_id).await;
981                    let _ = this.mark_server_healthy(&server_id, is_healthy).await;
982                    
983                    // If server is unhealthy, record the failure
984                    if !is_healthy {
985                        let _ = this.record_server_failure(&server_id).await;
986                    }
987                }
988            }
989        })
990    }
991
992    /// Gets the current load balancing strategy.
993    /// 
994    /// # Returns
995    /// 
996    /// A reference to the current `RiLoadBalancerStrategy`
997    pub fn get_strategy(&self) -> &RiLoadBalancerStrategy {
998        &self.strategy
999    }
1000
1001    /// Sets the load balancing strategy.
1002    /// 
1003    /// # Parameters
1004    /// 
1005    /// - `strategy`: The new load balancing strategy to use
1006    pub async fn set_strategy(&mut self, strategy: RiLoadBalancerStrategy) {
1007        self.strategy = strategy;
1008    }
1009
1010    /// Gets the total number of servers.
1011    /// 
1012    /// # Returns
1013    /// 
1014    /// The total number of servers in the load balancer
1015    pub async fn get_server_count(&self) -> usize {
1016        self.servers.read().await.len()
1017    }
1018
1019    /// Gets the number of healthy servers.
1020    /// 
1021    /// # Returns
1022    /// 
1023    /// The number of healthy servers in the load balancer
1024    pub async fn get_healthy_server_count(&self) -> usize {
1025        self.get_healthy_servers().await.len()
1026    }
1027}