67 lines
2.2 KiB
Rust
67 lines
2.2 KiB
Rust
//! Load balancing algorithms for upstream selection.
|
|
//!
|
|
//! Uses Pingora's built-in load balancing primitives where possible,
|
|
//! with custom extensions for session-affinity consistent hashing.
|
|
|
|
use crate::config::Endpoint;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
|
|
/// Load balancer that selects an upstream endpoint from a pool.
|
|
#[derive(Debug)]
|
|
pub struct LoadBalancer {
|
|
endpoints: Vec<Endpoint>,
|
|
counter: AtomicUsize,
|
|
}
|
|
|
|
impl LoadBalancer {
|
|
/// Create a new load balancer with the given endpoints.
|
|
pub fn new(endpoints: Vec<Endpoint>) -> Self {
|
|
Self {
|
|
endpoints,
|
|
counter: AtomicUsize::new(0),
|
|
}
|
|
}
|
|
|
|
/// Update the endpoint pool (e.g., on hot-reload).
|
|
pub fn update_endpoints(&mut self, endpoints: Vec<Endpoint>) {
|
|
self.endpoints = endpoints;
|
|
}
|
|
|
|
/// Select an endpoint using round-robin.
|
|
pub fn round_robin(&self) -> Option<&Endpoint> {
|
|
let healthy: Vec<&Endpoint> = self.endpoints.iter().filter(|e| e.ready).collect();
|
|
if healthy.is_empty() {
|
|
return None;
|
|
}
|
|
let idx = self.counter.fetch_add(1, Ordering::Relaxed) % healthy.len();
|
|
Some(healthy[idx])
|
|
}
|
|
|
|
/// Select an endpoint using the least-connections strategy (placeholder).
|
|
///
|
|
/// Full implementation would track active connections per endpoint.
|
|
pub fn least_connections(&self) -> Option<&Endpoint> {
|
|
self.endpoints.iter().filter(|e| e.ready).min_by_key(|_| {
|
|
// Placeholder: return 0 for now. Real impl would track connection counts.
|
|
0usize
|
|
})
|
|
}
|
|
|
|
/// Consistent hash selection based on a key (for session affinity).
|
|
pub fn consistent_hash(&self, key: &str) -> Option<&Endpoint> {
|
|
use std::collections::hash_map::DefaultHasher;
|
|
use std::hash::{Hash, Hasher};
|
|
|
|
let healthy: Vec<&Endpoint> = self.endpoints.iter().filter(|e| e.ready).collect();
|
|
if healthy.is_empty() {
|
|
return None;
|
|
}
|
|
|
|
let mut hasher = DefaultHasher::new();
|
|
key.hash(&mut hasher);
|
|
let hash = hasher.finish();
|
|
let idx = hash as usize % healthy.len();
|
|
Some(healthy[idx])
|
|
}
|
|
}
|