introduces location based algorithms

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.
 

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<String, Error> {
+    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<HashMap<String, String>> = 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()),
+    ])
+});

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<String> 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<ServerClient, Error> {
         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

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(),
         })
     }
 }

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<Vec<ServerClient>, Error> {
         self.0
-            .lrange("server_url", 0, -1)
+            .lrange("server_urls", 0, -1)
             .await?
             .into_iter()
             .map(|v| Ok(StaticServerData::from_json(v)?.into()))

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();