From 74002a78b0e9e80a88fd893059c45d853db8a0b4 Mon Sep 17 00:00:00 2001
From: Rohit <rohitnarnolia88@gmail.com>
Date: Wed, 20 May 2026 10:19:49 +0530
Subject: [PATCH] refactor: clean up code formatting and improve health check
 logic

---
 ROADMAP.md              |  2 +-
 src/config/validator.rs |  2 +-
 src/health/tcp.rs       | 18 +++++++++++-------
 src/proxy/tcp.rs        | 10 ++++++++--
 4 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/ROADMAP.md b/ROADMAP.md
index 407273d..3e9f3bb 100644
--- a/ROADMAP.md
+++ b/ROADMAP.md
@@ -129,7 +129,7 @@
 - [ ] Refactor duplicated runtime logic
 - [ ] Separate balancing module
 - [ ] Separate health module
-- [*] Improve state ownership model
+- [ ] Improve state ownership model
 
 ---
 
diff --git a/src/config/validator.rs b/src/config/validator.rs
index 6b765d0..252f3b1 100644
--- a/src/config/validator.rs
+++ b/src/config/validator.rs
@@ -3,7 +3,7 @@ use anyhow::{Result, bail};
 use std::collections::HashSet;
 
 // The load balancer follows a fail-fast startup philosophy.
-// Invalid topology or malformed configuration should prevent startup
+// Invalid configuration should prevent startup
 // rather than causing runtime instability.
 
 pub fn validate_config(config: &Config) -> Result<()> {
diff --git a/src/health/tcp.rs b/src/health/tcp.rs
index 52f1bb7..59564c3 100644
--- a/src/health/tcp.rs
+++ b/src/health/tcp.rs
@@ -31,15 +31,19 @@ pub async fn check_backend_status(backend: &BackendState) -> Result<()> {
 
 pub async fn start_health_checker(state: SharedAppState, interval_secs: u64) {
     loop {
-        let state = state.read().await;
-        for upstream in &state.upstreams {
-            for backend in &upstream.backends {
-                let _ = check_backend_status(backend).await;
-            }
+        let backends = {
+            let state = state.read().await;
+
+            state
+                .upstreams
+                .iter()
+                .flat_map(|upstream| upstream.backends.clone())
+                .collect::<Vec<_>>()
+        };
+        for backend in backends {
+            let _ = check_backend_status(&backend).await;
         }
 
-        // releasing the lock before going to sleep ..
-        drop(state);
         sleep(Duration::from_secs(interval_secs)).await;
     }
 }
diff --git a/src/proxy/tcp.rs b/src/proxy/tcp.rs
index b5bea43..71c89db 100644
--- a/src/proxy/tcp.rs
+++ b/src/proxy/tcp.rs
@@ -1,6 +1,6 @@
 use crate::state::app::SharedAppState;
 use crate::state::backend::ConnectionGuard;
-use std::time::Duration;
+use std::{collections::HashSet, time::Duration};
 use tokio::{
     io::copy_bidirectional,
     net::{TcpListener, TcpStream},
@@ -32,6 +32,8 @@ async fn handle_connection(mut stream: TcpStream, state: SharedAppState) -> anyh
         (state.retry_attempts, state.connect_timeout, state.idle_timeout)
     };
 
+    let mut attempted_backends = HashSet::new();
+
     // retry attempts mean how many times we should try
     // connecting the client to a suitable backend server
     //
@@ -55,6 +57,9 @@ async fn handle_connection(mut stream: TcpStream, state: SharedAppState) -> anyh
             // format!("{}:{}", backend.config.host, backend.config.port)
         };
         let backend_address = guard.address();
+        if attempted_backends.contains(guard.backend_id()) {
+            continue;
+        }
 
         info!("forwarding traffic to {}", backend_address);
 
@@ -64,7 +69,8 @@ async fn handle_connection(mut stream: TcpStream, state: SharedAppState) -> anyh
             }
             Err(error) => {
                 guard.mark_backend_unhealthy();
-                error!("backend {} failed: {:?}", backend_address, error);
+                error!(backend_id = %guard.backend_id(),backend = %backend_address,attempt = attempted_backends.len() + 1,"backend request failed: {:?}",error);
+                attempted_backends.insert(guard.backend_id().to_string());
                 continue;
             }
         }