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}