pub mod cookie_hash;
mod ip_extractor;
pub mod ip_hash;
pub mod round_robin;
pub mod rps_autoscale;
use super::worker::WorkerConfig;
use crate::client::Client;
use crate::error::FaucetResult;
use crate::leak;
use cookie_hash::CookieHash;
use hyper::Request;
pub use ip_extractor::IpExtractor;
use std::net::IpAddr;
use std::str::FromStr;
use uuid::Uuid;
use self::ip_hash::IpHash;
use self::round_robin::RoundRobin;
use self::rps_autoscale::RpsAutoscale;
const DEFAULT_MAX_RPS: f64 = 10.0;
trait LoadBalancingStrategy {
type Input;
async fn entry(&self, ip: Self::Input) -> Client;
}
#[derive(Debug, Clone, Copy, clap::ValueEnum, Eq, PartialEq, serde::Deserialize)]
#[serde(rename = "snake_case")]
pub enum Strategy {
#[serde(alias = "round_robin", alias = "RoundRobin", alias = "round-robin")]
RoundRobin,
#[serde(alias = "ip_hash", alias = "IpHash", alias = "ip-hash")]
IpHash,
#[serde(alias = "cookie_hash", alias = "CookieHash", alias = "cookie-hash")]
CookieHash,
#[serde(alias = "rps", alias = "Rps", alias = "rps")]
Rps,
}
impl FromStr for Strategy {
type Err = &'static str;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"round_robin" => Ok(Self::RoundRobin),
"ip_hash" => Ok(Self::IpHash),
"cookie_hash" => Ok(Self::CookieHash),
"rps" => Ok(Self::Rps),
_ => Err("invalid strategy"),
}
}
}
#[derive(Debug, Clone, Copy)]
enum LBIdent {
Ip(IpAddr),
Uuid(Uuid),
}
#[derive(Copy, Clone)]
enum DynLoadBalancer {
IpHash(&'static ip_hash::IpHash),
RoundRobin(&'static round_robin::RoundRobin),
CookieHash(&'static cookie_hash::CookieHash),
Rps(&'static rps_autoscale::RpsAutoscale),
}
impl LoadBalancingStrategy for DynLoadBalancer {
type Input = LBIdent;
async fn entry(&self, ip: LBIdent) -> Client {
match ip {
LBIdent::Ip(ip) => match self {
DynLoadBalancer::RoundRobin(rr) => rr.entry(ip).await,
DynLoadBalancer::IpHash(ih) => ih.entry(ip).await,
DynLoadBalancer::Rps(rr) => rr.entry(ip).await,
_ => unreachable!(
"This should never happen, ip should never be passed to cookie hash"
),
},
LBIdent::Uuid(uuid) => match self {
DynLoadBalancer::CookieHash(ch) => ch.entry(uuid).await,
_ => unreachable!(
"This should never happen, uuid should never be passed to round robin or ip hash"
),
},
}
}
}
pub(crate) struct LoadBalancer {
strategy: DynLoadBalancer,
extractor: IpExtractor,
}
impl LoadBalancer {
pub async fn new(
strategy: Strategy,
extractor: IpExtractor,
workers: &[&'static WorkerConfig],
max_rps_config: Option<f64>, ) -> FaucetResult<Self> {
let strategy: DynLoadBalancer = match strategy {
Strategy::RoundRobin => {
DynLoadBalancer::RoundRobin(leak!(RoundRobin::new(workers).await))
}
Strategy::IpHash => DynLoadBalancer::IpHash(leak!(IpHash::new(workers).await)),
Strategy::CookieHash => {
DynLoadBalancer::CookieHash(leak!(CookieHash::new(workers).await))
}
Strategy::Rps => {
let rps_value = max_rps_config.unwrap_or(DEFAULT_MAX_RPS);
DynLoadBalancer::Rps(leak!(RpsAutoscale::new(workers, rps_value).await))
}
};
Ok(Self {
strategy,
extractor,
})
}
pub fn get_strategy(&self) -> Strategy {
match self.strategy {
DynLoadBalancer::RoundRobin(_) => Strategy::RoundRobin,
DynLoadBalancer::IpHash(_) => Strategy::IpHash,
DynLoadBalancer::CookieHash(_) => Strategy::CookieHash,
DynLoadBalancer::Rps(_) => Strategy::Rps,
}
}
async fn get_client_ip(&self, ip: IpAddr) -> FaucetResult<Client> {
Ok(self.strategy.entry(LBIdent::Ip(ip)).await)
}
async fn get_client_uuid(&self, uuid: Uuid) -> FaucetResult<Client> {
Ok(self.strategy.entry(LBIdent::Uuid(uuid)).await)
}
pub async fn get_client(&self, ip: IpAddr, uuid: Option<Uuid>) -> FaucetResult<Client> {
if let Some(uuid) = uuid {
self.get_client_uuid(uuid).await
} else {
self.get_client_ip(ip).await
}
}
pub fn extract_ip<B>(
&self,
request: &Request<B>,
socket: Option<IpAddr>,
) -> FaucetResult<IpAddr> {
self.extractor.extract(request, socket)
}
}
impl Clone for LoadBalancer {
fn clone(&self) -> Self {
Self {
strategy: self.strategy,
extractor: self.extractor,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_strategy_from_str() {
assert_eq!(
Strategy::from_str("round_robin").unwrap(),
Strategy::RoundRobin
);
assert_eq!(Strategy::from_str("ip_hash").unwrap(), Strategy::IpHash);
assert!(Strategy::from_str("invalid").is_err());
}
#[tokio::test]
async fn test_load_balancer_new_round_robin() {
let configs = Vec::new();
let _ = LoadBalancer::new(
Strategy::RoundRobin,
IpExtractor::XForwardedFor,
&configs,
None,
)
.await
.expect("failed to create load balancer");
}
#[tokio::test]
async fn test_load_balancer_new_ip_hash() {
let configs = Vec::new();
let _ = LoadBalancer::new(Strategy::IpHash, IpExtractor::XForwardedFor, &configs, None)
.await
.expect("failed to create load balancer");
}
#[tokio::test]
async fn test_load_balancer_extract_ip() {
let configs = Vec::new();
let load_balancer = LoadBalancer::new(
Strategy::RoundRobin,
IpExtractor::XForwardedFor,
&configs,
None,
)
.await
.expect("failed to create load balancer");
let request = Request::builder()
.header("x-forwarded-for", "192.168.0.1")
.body(())
.unwrap();
let ip = load_balancer
.extract_ip(&request, Some("127.0.0.1".parse().unwrap()))
.expect("failed to extract ip");
assert_eq!(ip, "192.168.0.1".parse::<IpAddr>().unwrap());
}
#[tokio::test]
async fn test_load_balancer_get_client() {
use crate::client::ExtractSocketAddr;
let configs: [&'static WorkerConfig; 2] = [
&*Box::leak(Box::new(WorkerConfig::dummy(
"test",
"127.0.0.1:9999",
true,
))),
&*Box::leak(Box::new(WorkerConfig::dummy(
"test",
"127.0.0.1:9998",
true,
))),
];
let load_balancer = LoadBalancer::new(
Strategy::RoundRobin,
IpExtractor::XForwardedFor,
&configs,
None,
)
.await
.expect("failed to create load balancer");
let ip = "192.168.0.1".parse().unwrap();
let client = load_balancer
.get_client_ip(ip)
.await
.expect("failed to get client");
assert_eq!(client.socket_addr(), "127.0.0.1:9999".parse().unwrap());
let client = load_balancer
.get_client_ip(ip)
.await
.expect("failed to get client");
assert_eq!(client.socket_addr(), "127.0.0.1:9998".parse().unwrap());
for config in configs.iter() {
config.wait_until_done().await;
}
}
#[tokio::test]
async fn test_clone_load_balancer() {
let configs = Vec::new();
let load_balancer = LoadBalancer::new(
Strategy::RoundRobin,
IpExtractor::XForwardedFor,
&configs,
None,
)
.await
.expect("failed to create load balancer");
let _ = load_balancer.clone();
}
}