feat: graceful shutdown

src/lib.rs

@@ -23,7 +23,8 @@
 //! async fn main() {
 //!     let config = Config::new(Path::new("config.toml"));
 //!     let lb = Layer4::new(Arc::new(Mutex::new(config)));
-//!     let _ = lb.start().await;
+//!     let (_, shutdown_rx) = tokio::sync::watch::channel(false);
+//!     let _ = lb.start(shutdown_rx).await;
 //! }
 //! ```
 

src/load_balancer/layer4.rs

@@ -24,7 +24,10 @@ impl load_balancer::LoadBalancer for Layer4 {
     //will listen to incoming requests at given address
     //calls pick_server to pick a server when user sends a request
     //calls Server::transfer_data to transfer data between server and client
-    async fn start(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    async fn start(
+        &self,
+        mut shutdown_rx: tokio::sync::watch::Receiver<bool>,
+    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         //load balancer address from config
         let lb_address = {
             let config = self.config.lock().unwrap();
@@ -36,28 +39,45 @@ impl load_balancer::LoadBalancer for Layer4 {
         //create a TcpListener and binds it to load balancer address
         let listener = TcpListener::bind((host, port)).await?;
 
+        println!("Layer 4 Load Balancer listening on {}:{}", host, port);
+
         //loop to continuously accept incoming connections
         loop {
-            //accept incoming connections
-            let (stream, addr) = listener.accept().await?;
-
-            //clone the server list to safely share across multiple threads
-            let config_clone = self.config.clone();
-
-            //spawn a tokio task to server multiple connections concurrently
-            tokio::task::spawn(async move {
-                //pick a server
-                let server = Self::pick_server(config_clone, addr)
-                    .await
-                    .expect("No server");
-                //call Server::transfer_data to transfer data between server and client
-                if let Err(err) = Server::transfer_data(server, stream).await {
-                    eprintln!("Error transferring data {:?}", err);
+            tokio::select! {
+                // Check if shutdown signal is received
+                _ = shutdown_rx.changed() => {
+                    if *shutdown_rx.borrow() {
+                        println!("Shutdown signal received, stopping Layer 4 Load Balancer...");
+                        break;
+                    }
+                }
+                // Accept incoming connections
+                result = listener.accept() => {
+                    match result {
+                        Ok((stream, addr)) => {
+                            //clone the server list to safely share across multiple threads
+                            let config_clone = self.config.clone();
+
+                            //spawn a tokio task to server multiple connections concurrently
+                            tokio::task::spawn(async move {
+                                //pick a server
+                                let server = Self::pick_server(config_clone, addr)
+                                    .await
+                                    .expect("No server");
+                                //call Server::transfer_data to transfer data between server and client
+                                if let Err(err) = Server::transfer_data(server, stream).await {
+                                    eprintln!("Error transferring data {:?}", err);
+                                }
+                            });
+                        }
+                        Err(e) => {
+                            eprintln!("Error accepting connection: {:?}", e);
+                        }
+                    }
                 }
-            });
+            }
         }
-    }
 
-    //stops layer 4 load balancer
-    fn stop(&self) {}
+        Ok(())
+    }
 }

src/load_balancer/layer7.rs

@@ -28,7 +28,10 @@ impl LoadBalancer for Layer7 {
     //will listen to incoming requests at given address
     //calls pick_server to pick a server when user sends a request
     //calls Server::handle_request to forward request to the server
-    async fn start(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    async fn start(
+        &self,
+        mut shutdown_rx: tokio::sync::watch::Receiver<bool>,
+    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         //load balancer address from config
         let lb_address = {
             let config = self.config.lock().unwrap();
@@ -40,45 +43,63 @@ impl LoadBalancer for Layer7 {
         //create a TcpListener and binds it to load balancer address
         let listener = TcpListener::bind((host, port)).await?;
 
+        println!("Layer 7 Load Balancer listening on {}:{}", host, port);
+
         //loop to continuously accept incoming connections
         loop {
-            //accept incoming connections
-            let (stream, addr) = listener.accept().await?;
-            let io = TokioIo::new(stream);
-
-            //clone the server list to safely share across multiple threads
-            let config_clone = self.config.clone();
+            tokio::select! {
+                // Check if shutdown signal is received
+                _ = shutdown_rx.changed() => {
+                    if *shutdown_rx.borrow() {
+                        println!("Shutdown signal received, stopping Layer 7 Load Balancer...");
+                        break;
+                    }
+                }
+                // Accept incoming connections
+                result = listener.accept() => {
+                    match result {
+                        Ok((stream, addr)) => {
+                            let io = TokioIo::new(stream);
 
-            //spawn a tokio task to server multiple connections concurrently
-            tokio::task::spawn(async move {
-                if let Err(err) = http1::Builder::new()
-                    .preserve_header_case(true)
-                    .title_case_headers(true)
-                    //bind the incoming connection to handle_request
-                    .serve_connection(
-                        io,
-                        service_fn(move |req| {
                             //clone the server list to safely share across multiple threads
-                            let config_clone = config_clone.clone();
-                            async move {
-                                //pick a server
-                                let config_clone = config_clone.clone();
-                                let server = Self::pick_server(config_clone, addr)
+                            let config_clone = self.config.clone();
+
+                            //spawn a tokio task to server multiple connections concurrently
+                            tokio::task::spawn(async move {
+                                if let Err(err) = http1::Builder::new()
+                                    .preserve_header_case(true)
+                                    .title_case_headers(true)
+                                    //bind the incoming connection to handle_request
+                                    .serve_connection(
+                                        io,
+                                        service_fn(move |req| {
+                                            //clone the server list to safely share across multiple threads
+                                            let config_clone = config_clone.clone();
+                                            async move {
+                                                //pick a server
+                                                let config_clone = config_clone.clone();
+                                                let server = Self::pick_server(config_clone, addr)
+                                                    .await
+                                                    .expect("No server");
+                                                //call Server::handle_request to forward the request to server
+                                                Server::handle_request(server, req, addr).await
+                                            }
+                                        }),
+                                    )
                                     .await
-                                    .expect("No server");
-                                //call Server::handle_request to forward the request to server
-                                Server::handle_request(server, req, addr).await
-                            }
-                        }),
-                    )
-                    .await
-                {
-                    eprintln!("Error serving connection: {:?}", err);
+                                {
+                                    eprintln!("Error serving connection: {:?}", err);
+                                }
+                            });
+                        }
+                        Err(e) => {
+                            eprintln!("Error accepting connection: {:?}", e);
+                        }
+                    }
                 }
-            });
+            }
         }
-    }
 
-    //stops layer 7 load balancer
-    fn stop(&self) {}
+        Ok(())
+    }
 }

src/load_balancer/load_balancer.rs

@@ -15,7 +15,13 @@ pub trait LoadBalancer {
     fn new(config: SyncConfig) -> Self;
 
     /// Starts the load balancer and begins accepting connections
-    async fn start(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
+    ///
+    /// # Arguments
+    /// * `shutdown_rx` - A watch receiver that signals when to initiate shutdown
+    async fn start(
+        &self,
+        shutdown_rx: tokio::sync::watch::Receiver<bool>,
+    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
 
     /// Picks a server based on the configured algorithm to handle an incoming request
     ///
@@ -35,8 +41,4 @@ pub trait LoadBalancer {
         //return picked server
         Some(server)
     }
-
-    //stops the load balancer
-    #[allow(dead_code)]
-    fn stop(&self);
 }

src/main.rs

@@ -8,27 +8,46 @@ use load_balancer::layer7::Layer7;
 use std::path::Path;
 
 use std::sync::{Arc, Mutex};
+use tokio::signal;
 
 #[tokio::main]
 async fn main() {
     let config = Config::new(Path::new("config.toml"));
     let layer_mode = config.layer_mode.clone();
     let config_arc = Arc::new(Mutex::new(config));
 
+    // Create a channel for graceful shutdown
+    let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false);
+
+    // Spawn a task to handle shutdown signals
+    tokio::spawn(async move {
+        match signal::ctrl_c().await {
+            Ok(()) => {
+                println!("\nReceived shutdown signal, initiating graceful shutdown...");
+                let _ = shutdown_tx.send(true);
+            }
+            Err(err) => {
+                eprintln!("Unable to listen for shutdown signal: {}", err);
+            }
+        }
+    });
+
     match layer_mode {
         LayerMode::L4 => {
             println!("Starting Layer 4 (TCP) Load Balancer...");
             let lb = Layer4::new(config_arc);
-            if let Err(e) = lb.start().await {
+            if let Err(e) = lb.start(shutdown_rx).await {
                 eprintln!("Error starting Layer 4 load balancer: {:?}", e);
             }
         }
         LayerMode::L7 => {
             println!("Starting Layer 7 (HTTP) Load Balancer...");
             let lb = Layer7::new(config_arc);
-            if let Err(e) = lb.start().await {
+            if let Err(e) = lb.start(shutdown_rx).await {
                 eprintln!("Error starting Layer 7 load balancer: {:?}", e);
             }
         }
     }
+
+    println!("Load balancer shut down gracefully.");
 }