From c4c0a96a85aeb4f53d3b3d061dd6ca1e8d31e1cd Mon Sep 17 00:00:00 2001
From: pie-314 <aadarsh.official.xz@gmail.com>
Date: Sun, 17 May 2026 16:54:43 +0530
Subject: [PATCH] add timeouts and improve networking resilience

---
 src/config/types.rs |  2 ++
 src/proxy/tcp.rs    | 26 ++++++++++++++++++++------
 src/state/app.rs    | 10 +++++++++-
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/src/config/types.rs b/src/config/types.rs
index 70300f0..54e6ca8 100644
--- a/src/config/types.rs
+++ b/src/config/types.rs
@@ -28,6 +28,8 @@ pub struct LoadBalancerConfig {
     pub sticky_sessions: bool,
 
     pub health_check_interval_secs: u64,
+    pub connect_timeout_secs: u64,
+    pub idle_timeout_secs: u64,
 }
 
 // Static backend server definition loaded from configuration.
diff --git a/src/proxy/tcp.rs b/src/proxy/tcp.rs
index 933a455..11074f6 100644
--- a/src/proxy/tcp.rs
+++ b/src/proxy/tcp.rs
@@ -1,6 +1,7 @@
 use tokio::{
     io::copy_bidirectional,
     net::{TcpListener, TcpStream},
+    time::timeout,
 };
 use tracing::{error, info};
 
@@ -27,9 +28,9 @@ pub async fn start_tcp_proxy(address: &str, state: SharedAppState) -> anyhow::Re
 }
 
 async fn handle_connection(mut stream: TcpStream, state: SharedAppState) -> anyhow::Result<()> {
-    let retry_attempt = {
+    let (retry_attempt, connect_timeout, idle_timeout) = {
         let state = state.read().await;
-        state.retry_attempts
+        (state.retry_attempts, state.connect_timeout, state.idle_timeout)
     };
 
     // retry attempts mean how many times we should try
@@ -58,16 +59,29 @@ async fn handle_connection(mut stream: TcpStream, state: SharedAppState) -> anyh
 
         info!("forwarding traffic to {}", backend_address);
 
-        match TcpStream::connect(&backend_address).await {
-            Ok(mut backend_stream) => {
-                copy_bidirectional(&mut stream, &mut backend_stream).await?;
+        match timeout(connect_timeout, TcpStream::connect(&backend_address)).await {
+            Ok(Ok(mut backend_stream)) => {
+                if timeout(idle_timeout, copy_bidirectional(&mut stream, &mut backend_stream))
+                    .await
+                    .is_err()
+                {
+                    error!("connection with {} timed out (idle)", backend_address);
+                }
                 return Ok(());
             }
-            Err(error) => {
+            Ok(Err(error)) => {
                 // this is very important ..
                 guard.mark_backend_unhealthy();
                 error!("failed to connect to backend {}: {:?}", backend_address, error);
 
+                // retry another
+                continue;
+            }
+            Err(_) => {
+                // connection attempt timed out
+                guard.mark_backend_unhealthy();
+                error!("connection attempt to {} timed out", backend_address);
+
                 // retry another
                 continue;
             }
diff --git a/src/state/app.rs b/src/state/app.rs
index 7412eac..62c3a6c 100644
--- a/src/state/app.rs
+++ b/src/state/app.rs
@@ -3,6 +3,7 @@ use crate::config::LoadBalancingAlgorithm;
 use crate::{config::types::Config, state::backend::BackendState};
 use std::sync::Arc;
 use std::sync::atomic::AtomicUsize;
+use std::time::Duration;
 use tokio::sync::RwLock;
 // Contains all backend servers belonging to a single logical service.
 #[derive(Debug)]
@@ -40,6 +41,8 @@ impl UpstreamPool {
 pub struct AppState {
     pub retry_attempts: usize,
     pub upstreams: Vec<UpstreamPool>,
+    pub connect_timeout: Duration,
+    pub idle_timeout: Duration,
 }
 
 pub type SharedAppState = Arc<RwLock<AppState>>;
@@ -70,7 +73,12 @@ impl AppState {
             })
             .collect();
 
-        Self { upstreams, retry_attempts: config.load_balancer.retry_attempts }
+        Self {
+            upstreams,
+            retry_attempts: config.load_balancer.retry_attempts,
+            connect_timeout: Duration::from_secs(config.load_balancer.connect_timeout_secs),
+            idle_timeout: Duration::from_secs(config.load_balancer.idle_timeout_secs),
+        }
     }
 }