Improve logging and health checks

Author: HueCodes

README.md

@@ -5,7 +5,7 @@ A lightweight, high-performance HTTP reverse proxy written in Rust, prioritizing
 ## Features
 
 - **High Performance**: Built with Tokio and Hyper for maximum throughput
-- **Load Balancing**: Round-robin load balancing across multiple backend servers
+- **Load Balancing**: Round-robin and weighted round-robin across healthy backends
 - **Health Checks**: Automatic health monitoring of backend servers
 - **Configuration-Driven**: YAML-based configuration for easy management
 - **Structured Logging**: Comprehensive request/response logging with configurable levels
@@ -58,7 +58,7 @@ backends:
     health_check_path: "/"
 
 load_balancing:
-  strategy: "round_robin"  # Load balancing strategy
+  strategy: "round_robin"  # round_robin or weighted_round_robin
 
 health_checks:
   enabled: true           # Enable/disable health checks
@@ -86,10 +86,10 @@ The proxy consists of several key components:
 
 ## Load Balancing
 
-Currently supports round-robin load balancing:
-- Requests are distributed evenly across healthy backends
-- Unhealthy backends are automatically excluded
-- Backends are marked unhealthy after consecutive failures
+Supports both round-robin and weighted round-robin strategies:
+- Round-robin distributes requests evenly across healthy backends
+- Weighted round-robin honors configured weights for skewed traffic
+- Unhealthy backends are automatically excluded until they recover
 
 ## Health Checks
 

config.example.yaml

@@ -12,7 +12,7 @@ backends:
     health_check_path: "/posts/1"
 
 load_balancing:
-  strategy: "round_robin"
+  strategy: "round_robin"  # round_robin or weighted_round_robin
 
 health_checks:
   enabled: true

config.yaml

@@ -11,7 +11,7 @@ backends:
     health_check_path: "/"
 
 load_balancing:
-  strategy: "round_robin"  # round_robin, weighted_round_robin, least_connections
+  strategy: "round_robin"  # round_robin or weighted_round_robin
 
 health_checks:
   enabled: true

src/main.rs

@@ -10,6 +10,7 @@ use hyper_util::rt::TokioIo;
 use log::{debug, error, info, warn};
 use serde::Deserialize;
 use std::collections::HashMap;
+use std::io::Write;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::{Duration, Instant};
@@ -107,6 +108,12 @@ struct LoadBalancer {
     http_client: HttpClient,
 }
 
+#[derive(Debug)]
+enum HealthProbeError {
+    Timeout,
+    Transport(String),
+}
+
 impl LoadBalancer {
     fn new(
         backends: Vec<Backend>,
@@ -215,7 +222,26 @@ impl LoadBalancer {
                 }
             };
 
-            match self.http_client.get(uri).await {
+            let request_future = self.http_client.get(uri);
+            let response_result = if config.timeout_seconds > 0 {
+                match tokio::time::timeout(
+                    Duration::from_secs(config.timeout_seconds),
+                    request_future,
+                )
+                .await
+                {
+                    Ok(inner_result) => {
+                        inner_result.map_err(|err| HealthProbeError::Transport(err.to_string()))
+                    }
+                    Err(_) => Err(HealthProbeError::Timeout),
+                }
+            } else {
+                request_future
+                    .await
+                    .map_err(|err| HealthProbeError::Transport(err.to_string()))
+            };
+
+            match response_result {
                 Ok(response) => {
                     let status = response.status();
                     let mut health = self.health_status.write().await;
@@ -240,8 +266,19 @@ impl LoadBalancer {
                         }
                     }
                 }
-                Err(e) => {
-                    debug!("Health check failed for {}: {}", backend.url, e);
+                Err(error) => {
+                    match error {
+                        HealthProbeError::Timeout => {
+                            debug!(
+                                "Health check timed out for {} after {}s",
+                                backend.url, config.timeout_seconds
+                            );
+                        }
+                        HealthProbeError::Transport(e) => {
+                            debug!("Health check failed for {}: {}", backend.url, e);
+                        }
+                    }
+
                     let mut health = self.health_status.write().await;
 
                     if let Some(backend_health) = health.get_mut(&backend.url) {
@@ -465,17 +502,32 @@ fn load_config() -> Result<Config> {
 }
 
 fn setup_logging(config: &LoggingConfig) {
-    let log_level = match config.level.as_str() {
+    let level_key = config.level.to_lowercase();
+    let log_level = match level_key.as_str() {
         "debug" => log::LevelFilter::Debug,
         "info" => log::LevelFilter::Info,
         "warn" => log::LevelFilter::Warn,
         "error" => log::LevelFilter::Error,
         _ => log::LevelFilter::Info,
     };
 
-    env_logger::Builder::from_default_env()
-        .filter_level(log_level)
-        .init();
+    let mut builder = env_logger::Builder::from_default_env();
+    builder.filter_level(log_level);
+
+    if config.format.eq_ignore_ascii_case("json") {
+        builder.format(|buf, record| {
+            let timestamp = Utc::now().to_rfc3339();
+            let payload = serde_json::json!({
+                "timestamp": timestamp,
+                "level": record.level().to_string(),
+                "target": record.target(),
+                "message": record.args().to_string(),
+            });
+            writeln!(buf, "{}", payload)
+        });
+    }
+
+    builder.init();
 }
 
 #[tokio::main]