4 moves server state handling to redis

src/algorithms/least_connection.rs

@@ -1,10 +1,12 @@
-use crate::{error::Error, middleware::Server};
+use std::collections::HashMap;
 
-pub async fn least_connection(available_servers: &[Server]) -> Result<Server, Error> {
-    let server = available_servers
-        .iter()
-        .min_by_key(|server| server.load())
+use crate::error::Error;
+
+pub async fn least_connection(server_loads: HashMap<String, u32>) -> Result<String, Error> {
+    let (url, _) = server_loads
+        .into_iter()
+        .min_by_key(|(_, load)| *load)
         .ok_or_else(|| Error::NoServerAvailable)?;
 
-    Ok(server.clone())
+    Ok(url)
 }

src/algorithms/mod.rs

@@ -1,4 +1,6 @@
-use crate::{error::Error, middleware::Server};
+use reqwest::Url;
+
+use crate::{db::RedisClient, error::Error, middleware::ServerClient};
 
 mod least_connection;
 mod resource_based;
@@ -27,18 +29,34 @@ impl From<String> for Algorithm {
 }
 
 impl Algorithm {
-    pub async fn select_server(&self, available_servers: &[Server]) -> Result<Server, Error> {
-        match self {
-            Algorithm::LeastConnection => {
-                least_connection::least_connection(available_servers).await
-            }
-            Algorithm::ResourceBased => resource_based::resource_based(available_servers).await,
+    pub async fn select_server(
+        &self,
+        mut redis_client: RedisClient,
+    ) -> 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::ResourceBased => unimplemented!(),
             Algorithm::WeightedLeastConnection => {
-                weighted_least_connection::weighted_least_connection(available_servers).await
+                weighted_least_connection::weighted_least_connection(server_loads, weights).await
             }
             Algorithm::WeightedResponseTime => {
-                weighted_response_time::weighted_response_time(available_servers).await
+                weighted_response_time::weighted_response_time(
+                    redis_client.get_all_server_mean_latency().await?,
+                    weights,
+                )
+                .await
             }
-        }
+        }?;
+
+        redis_client.update_server_load(&url, 1).await?;
+
+        let url = url.parse::<Url>().map_err(|e| Error::Other(e.into()))?;
+
+        Ok(ServerClient {
+            url,
+            client: reqwest::Client::new(),
+        })
     }
 }

src/algorithms/resource_based.rs

@@ -1,6 +1,8 @@
-use crate::{error::Error, middleware::Server};
+use crate::{error::Error, middleware::StaticServerData};
 
 // TODO: Implement resource-based load balancing algorithm
-pub async fn resource_based(available_servers: &[Server]) -> Result<Server, Error> {
+pub async fn _resource_based(
+    available_servers: &[StaticServerData],
+) -> Result<StaticServerData, Error> {
     Ok(available_servers[0].clone())
 }

src/algorithms/weighted_least_connection.rs

@@ -1,10 +1,18 @@
-use crate::{error::Error, middleware::Server};
+use std::collections::HashMap;
 
-pub async fn weighted_least_connection(available_servers: &[Server]) -> Result<Server, Error> {
-    let server = available_servers
-        .iter()
-        .min_by_key(|server| server.load() / server.weight())
+use crate::error::Error;
+
+pub async fn weighted_least_connection(
+    server_loads: HashMap<String, u32>,
+    weights: HashMap<String, u32>,
+) -> Result<String, Error> {
+    let (url, _) = server_loads
+        .into_iter()
+        .min_by_key(|(key, load)| {
+            let weight = weights.get(key).unwrap_or(&1);
+            *load / weight
+        })
         .ok_or_else(|| Error::NoServerAvailable)?;
 
-    Ok(server.clone())
+    Ok(url)
 }

src/algorithms/weighted_response_time.rs

@@ -1,10 +1,18 @@
-use crate::{error::Error, middleware::Server};
+use std::collections::HashMap;
 
-pub async fn weighted_response_time(available_servers: &[Server]) -> Result<Server, Error> {
-    let server = available_servers
-        .iter()
-        .min_by_key(|server| server.mean_latency())
+use crate::error::Error;
+
+pub async fn weighted_response_time(
+    latencies: HashMap<String, u32>,
+    weights: HashMap<String, u32>,
+) -> Result<String, Error> {
+    let (url, _) = latencies
+        .into_iter()
+        .min_by_key(|(key, latency)| {
+            let weight = weights.get(key).unwrap_or(&1);
+            *latency / weight
+        })
         .ok_or_else(|| Error::NoServerAvailable)?;
 
-    Ok(server.clone())
+    Ok(url)
 }

src/app.rs

@@ -0,0 +1,75 @@
+use axum::{Router, routing::get};
+use tokio::{net::TcpListener, task::JoinHandle};
+use tower_http::{
+    cors::{AllowHeaders, AllowMethods, AllowOrigin, CorsLayer},
+    trace::TraceLayer,
+};
+
+use crate::config::State;
+use crate::middleware::request_route;
+use crate::route::health::status;
+use crate::services::{latency_tracker_worker, server_status_worker};
+
+type JoinHandleWrapper = JoinHandle<Result<(), std::io::Error>>;
+
+/// Application struct to hold main and background worker tasks
+pub struct App {
+    main: JoinHandleWrapper,
+    server_status_background_worker: JoinHandleWrapper,
+    latency_tracker_background_worker: JoinHandleWrapper,
+}
+
+impl App {
+    pub async fn setup(
+        state: State,
+        listener: TcpListener,
+    ) -> Result<App, Box<dyn std::error::Error>> {
+        let server = Router::new()
+            .route("/status", get(status))
+            .layer(
+                CorsLayer::new()
+                    .allow_headers(AllowHeaders::any())
+                    .allow_origin(AllowOrigin::any())
+                    .allow_methods(AllowMethods::any()),
+            )
+            .layer(TraceLayer::new_for_http())
+            // add rate limitter middleware mechanism
+            .layer(axum::middleware::from_fn_with_state(
+                state.clone(),
+                request_route,
+            ))
+            .with_state(state.clone());
+
+        let main = tokio::spawn(async move { axum::serve(listener, server).await });
+
+        let redis_conn_1 = state.redis_conn.clone();
+        let redis_conn_2 = state.redis_conn.clone();
+
+        let server_status_background_worker = tokio::spawn(async move {
+            let _: () = server_status_worker(redis_conn_1).await;
+            Ok(())
+        });
+
+        let latency_tracker_background_worker = tokio::spawn(async move {
+            let _: () = latency_tracker_worker(redis_conn_2).await;
+            Ok(())
+        });
+
+        Ok(Self {
+            main,
+            server_status_background_worker,
+            latency_tracker_background_worker,
+        })
+    }
+
+    pub async fn start(self) -> Result<(), Box<dyn std::error::Error>> {
+        match tokio::try_join!(
+            self.main,
+            self.server_status_background_worker,
+            self.latency_tracker_background_worker
+        ) {
+            Ok(_) => Ok(()),
+            Err(err) => Err(err)?,
+        }
+    }
+}

src/config.rs

@@ -3,12 +3,12 @@ use serde::Deserialize;
 use crate::{
     algorithms::Algorithm,
     db::{self, RedisClient},
-    middleware::{Server, ServerClients},
+    middleware::StaticServerData,
 };
 
 #[derive(Deserialize)]
 pub struct SystemConfig {
-    pub available_servers: String, // TODO: This should be hosted in redis
+    pub available_servers: String,
     pub port: u16,
     pub redis_url: String,
     pub algorithm: String,
@@ -20,13 +20,12 @@ impl SystemConfig {
         dotenvy::dotenv_override().ok();
 
         envy::from_env::<Self>()
-            .map_err(|e| anyhow::anyhow!("Failed to load environment variables: {}", e))
+            .map_err(|e| anyhow::anyhow!("Failed to load environment variable(s): {}", e))
     }
 }
 
 #[derive(Clone)]
 pub struct State {
-    pub available_servers: ServerClients,
     pub redis_conn: RedisClient,
     pub algorithm: Algorithm,
 }
@@ -35,17 +34,16 @@ impl State {
     pub async fn new(config: &SystemConfig) -> Result<Self, Box<dyn std::error::Error>> {
         let servers = config.available_servers.split(',').collect::<Vec<&str>>();
 
-        let available_servers: Vec<Server> = servers
+        let available_servers: Vec<StaticServerData> = servers
             .clone()
             .into_iter()
-            .map(Server::new)
-            .collect::<Result<Vec<Server>, _>>()?;
+            .map(StaticServerData::new)
+            .collect::<Result<Vec<StaticServerData>, _>>()?;
 
         let redis_conn =
             db::RedisClient::init_redis(&config.redis_url, available_servers.clone()).await?;
 
         Ok(State {
-            available_servers: ServerClients::new(available_servers),
             redis_conn,
             algorithm: config.algorithm.clone().into(),
         })