diff --git a/Readme.md b/Readme.md index 79db28e..ddbd8ab 100644 --- a/Readme.md +++ b/Readme.md @@ -34,7 +34,7 @@ The server binds to the configured PORT. This project reads configuration from environment variables. The important variables are: -`AVAILABLE_SERVERS` — comma-separated list of backend base URLs and their weights (e.g. http://host:port$weight). +`AVAILABLE_SERVERS` — comma-separated list of backend base URLs and their weights (e.g. http://host:port|weight). `PORT` — port to bind the load balancer to. diff --git a/src/algorithms/location_based.rs b/src/algorithms/location_based.rs new file mode 100644 index 0000000..8e58e56 --- /dev/null +++ b/src/algorithms/location_based.rs @@ -0,0 +1,81 @@ +use std::{collections::HashMap, sync::LazyLock}; + +use crate::error::Error; + +pub async fn location_based(location: &str) -> Result { + SERVERS_LOCATION_MAPPINGS // Probably not the best solution, but it works for now + .get(location) + .cloned() + .ok_or(Error::InternalServerError) +} + +static SERVERS_LOCATION_MAPPINGS: LazyLock> = LazyLock::new(|| { + HashMap::from([ + // ====================== + // North America + // ====================== + ("us".into(), "https://us-east.example.com".into()), + ("us-east".into(), "https://us-east.example.com".into()), + ("us-central".into(), "https://us-east.example.com".into()), + ("ca".into(), "https://us-east.example.com".into()), + ("ca-east".into(), "https://us-east.example.com".into()), + ("mx".into(), "https://us-east.example.com".into()), + ("us-west".into(), "https://us-west.example.com".into()), + ("ca-west".into(), "https://us-west.example.com".into()), + // ====================== + // Europe + // ====================== + ("ie".into(), "https://eu-west.example.com".into()), + ("uk".into(), "https://eu-west.example.com".into()), + ("fr".into(), "https://eu-west.example.com".into()), + ("de".into(), "https://eu-west.example.com".into()), + ("nl".into(), "https://eu-west.example.com".into()), + ("be".into(), "https://eu-west.example.com".into()), + ("es".into(), "https://eu-west.example.com".into()), + ("pt".into(), "https://eu-west.example.com".into()), + ("pl".into(), "https://eu-central.example.com".into()), + ("cz".into(), "https://eu-central.example.com".into()), + ("at".into(), "https://eu-central.example.com".into()), + ("ch".into(), "https://eu-central.example.com".into()), + ("hu".into(), "https://eu-central.example.com".into()), + // ====================== + // Africa + // ====================== + ("ng".into(), "https://africa.example.com".into()), + ("gh".into(), "https://africa.example.com".into()), + ("ke".into(), "https://africa.example.com".into()), + ("za".into(), "https://africa.example.com".into()), + ("eg".into(), "https://africa.example.com".into()), + // ====================== + // Middle East + // ====================== + ("ae".into(), "https://middle-east.example.com".into()), + ("sa".into(), "https://middle-east.example.com".into()), + ("qa".into(), "https://middle-east.example.com".into()), + ("il".into(), "https://middle-east.example.com".into()), + // ====================== + // Asia + // ====================== + ("in".into(), "https://asia-south.example.com".into()), + ("pk".into(), "https://asia-south.example.com".into()), + ("bd".into(), "https://asia-south.example.com".into()), + ("lk".into(), "https://asia-south.example.com".into()), + ("jp".into(), "https://asia-east.example.com".into()), + ("kr".into(), "https://asia-east.example.com".into()), + ("tw".into(), "https://asia-east.example.com".into()), + ("sg".into(), "https://asia-southeast.example.com".into()), + ("id".into(), "https://asia-southeast.example.com".into()), + ("th".into(), "https://asia-southeast.example.com".into()), + ("vn".into(), "https://asia-southeast.example.com".into()), + ("ph".into(), "https://asia-southeast.example.com".into()), + // ====================== + // Oceania + // ====================== + ("au".into(), "https://australia.example.com".into()), + ("nz".into(), "https://australia.example.com".into()), + // ====================== + // Fallback / Generic + // ====================== + ("global".into(), "https://us-east.example.com".into()), + ]) +}); diff --git a/src/algorithms/mod.rs b/src/algorithms/mod.rs index dee1207..a64d722 100644 --- a/src/algorithms/mod.rs +++ b/src/algorithms/mod.rs @@ -3,6 +3,7 @@ use reqwest::Url; use crate::{db::RedisClient, error::Error, middleware::ServerClient}; mod least_connection; +mod location_based; mod resource_based; mod weighted_least_connection; mod weighted_response_time; @@ -11,6 +12,7 @@ mod weighted_response_time; pub enum Algorithm { #[default] LeastConnection, + LocationBased, ResourceBased, WeightedLeastConnection, WeightedResponseTime, @@ -20,6 +22,7 @@ impl From for Algorithm { fn from(algorithm: String) -> Self { match algorithm.as_str() { "least_connection" => Algorithm::LeastConnection, + "location_based" | "location" => Algorithm::LocationBased, "resource_based" => Algorithm::ResourceBased, "weighted_least_connection" => Algorithm::WeightedLeastConnection, "weighted_response_time" => Algorithm::WeightedResponseTime, @@ -32,11 +35,14 @@ impl Algorithm { pub async fn select_server( &self, mut redis_client: RedisClient, + location: &str, ) -> Result { let server_loads = redis_client.get_all_server_load().await?; let weights = redis_client.get_all_server_weights().await?; + let url = match self { Algorithm::LeastConnection => least_connection::least_connection(server_loads).await, + Algorithm::LocationBased => location_based::location_based(location).await, Algorithm::ResourceBased => unimplemented!(), Algorithm::WeightedLeastConnection => { weighted_least_connection::weighted_least_connection(server_loads, weights).await diff --git a/src/config.rs b/src/config.rs index 62683ac..16b9560 100644 --- a/src/config.rs +++ b/src/config.rs @@ -13,6 +13,7 @@ pub struct SystemConfig { pub redis_url: String, pub algorithm: String, pub trace_level: String, + pub default_location: String, } impl SystemConfig { @@ -27,6 +28,7 @@ impl SystemConfig { pub struct State { pub redis_conn: RedisClient, pub algorithm: Algorithm, + pub default_location: String, } impl State { @@ -45,6 +47,7 @@ impl State { Ok(State { redis_conn, algorithm: config.algorithm.clone().into(), + default_location: config.default_location.clone(), }) } } diff --git a/src/db/redis.rs b/src/db/redis.rs index 943b605..3635c2e 100644 --- a/src/db/redis.rs +++ b/src/db/redis.rs @@ -31,7 +31,7 @@ impl RedisClient { client .update_server_weight(server.url.as_str(), server.weight) - .await? + .await?; } Ok(client) @@ -58,13 +58,13 @@ impl RedisClient { /// Update the data of a server in Redis. pub async fn update_server_url(&mut self, value: &str) -> Result<(), Error> { - Ok(self.0.rpush("server_url", value).await.map(|_| ())?) + Ok(self.0.rpush("server_urls", value).await.map(|_| ())?) } /// Get all server data from Redis. pub async fn get_all_server_url(&mut self) -> Result, Error> { self.0 - .lrange("server_url", 0, -1) + .lrange("server_urls", 0, -1) .await? .into_iter() .map(|v| Ok(StaticServerData::from_json(v)?.into())) diff --git a/src/middleware/mod.rs b/src/middleware/mod.rs index 2b50826..3e3f358 100644 --- a/src/middleware/mod.rs +++ b/src/middleware/mod.rs @@ -36,9 +36,15 @@ pub async fn request_route( let route = parts.uri.to_string(); + let location = parts + .headers + .get("X-Location") + .and_then(|l| l.to_str().ok()) + .unwrap_or(&state.default_location); + let server_client = state .algorithm - .select_server(state.redis_conn.clone()) + .select_server(state.redis_conn.clone(), location) .await?; let start_time = std::time::Instant::now();