diff --git a/src/lib.rs b/src/lib.rs index 92edc63..1b22490 100644 --- a/src/lib.rs +++ b/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; //! } //! ``` diff --git a/src/load_balancer/layer4.rs b/src/load_balancer/layer4.rs index 22af0f5..7f553c6 100644 --- a/src/load_balancer/layer4.rs +++ b/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> { + async fn start( + &self, + mut shutdown_rx: tokio::sync::watch::Receiver, + ) -> Result<(), Box> { //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(()) + } } diff --git a/src/load_balancer/layer7.rs b/src/load_balancer/layer7.rs index 781bc2e..6e209b9 100644 --- a/src/load_balancer/layer7.rs +++ b/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> { + async fn start( + &self, + mut shutdown_rx: tokio::sync::watch::Receiver, + ) -> Result<(), Box> { //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(()) + } } diff --git a/src/load_balancer/load_balancer.rs b/src/load_balancer/load_balancer.rs index f5dd4a8..cadebef 100644 --- a/src/load_balancer/load_balancer.rs +++ b/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>; + /// + /// # Arguments + /// * `shutdown_rx` - A watch receiver that signals when to initiate shutdown + async fn start( + &self, + shutdown_rx: tokio::sync::watch::Receiver, + ) -> Result<(), Box>; /// 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); } diff --git a/src/main.rs b/src/main.rs index 05fcf8a..73870f4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,6 +8,7 @@ use load_balancer::layer7::Layer7; use std::path::Path; use std::sync::{Arc, Mutex}; +use tokio::signal; #[tokio::main] async fn main() { @@ -15,20 +16,38 @@ async fn main() { 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."); } diff --git a/tests/shutdown_tests.rs b/tests/shutdown_tests.rs new file mode 100644 index 0000000..1779c9f --- /dev/null +++ b/tests/shutdown_tests.rs @@ -0,0 +1,167 @@ +use deston::config::config::Config; +use deston::load_balancer::layer4::Layer4; +use deston::load_balancer::layer7::Layer7; +use deston::load_balancer::load_balancer::LoadBalancer; +use std::fs; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio::time::timeout; + +#[tokio::test] +async fn test_layer4_graceful_shutdown() { + // Create a test config + let config_content = r#" +[load_balancer] +address = "127.0.0.1" +port = 18080 +algorithm = "round_robin" +layer = "L4" + +[[server]] +address = "127.0.0.1" +port = 13000 +max_connections = 1000 +weight = 1 +"#; + + let mut config_path = std::env::temp_dir(); + config_path.push("test_shutdown_l4.toml"); + fs::write(&config_path, config_content).unwrap(); + + let config = Config::new(&config_path); + let config_arc = Arc::new(Mutex::new(config)); + + // Create shutdown channel + let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false); + + let lb = Layer4::new(config_arc); + + // Start the load balancer in a separate task + let lb_handle = tokio::spawn(async move { lb.start(shutdown_rx).await }); + + // Give it a moment to start + tokio::time::sleep(Duration::from_millis(100)).await; + + // Send shutdown signal + let _ = shutdown_tx.send(true); + + // The load balancer should shut down gracefully + let result = timeout(Duration::from_secs(5), lb_handle).await; + + // Clean up + fs::remove_file(config_path).ok(); + + // Assert that the load balancer shut down successfully + assert!( + result.is_ok(), + "Load balancer should shut down within timeout" + ); + assert!( + result.unwrap().is_ok(), + "Load balancer should return Ok on shutdown" + ); +} + +#[tokio::test] +async fn test_layer7_graceful_shutdown() { + // Create a test config + let config_content = r#" +[load_balancer] +address = "127.0.0.1" +port = 18081 +algorithm = "round_robin" +layer = "L7" + +[[server]] +address = "127.0.0.1" +port = 13001 +max_connections = 1000 +weight = 1 +"#; + + let mut config_path = std::env::temp_dir(); + config_path.push("test_shutdown_l7.toml"); + fs::write(&config_path, config_content).unwrap(); + + let config = Config::new(&config_path); + let config_arc = Arc::new(Mutex::new(config)); + + // Create shutdown channel + let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false); + + let lb = Layer7::new(config_arc); + + // Start the load balancer in a separate task + let lb_handle = tokio::spawn(async move { lb.start(shutdown_rx).await }); + + // Give it a moment to start + tokio::time::sleep(Duration::from_millis(100)).await; + + // Send shutdown signal + let _ = shutdown_tx.send(true); + + // The load balancer should shut down gracefully + let result = timeout(Duration::from_secs(5), lb_handle).await; + + // Clean up + fs::remove_file(config_path).ok(); + + // Assert that the load balancer shut down successfully + assert!( + result.is_ok(), + "Load balancer should shut down within timeout" + ); + assert!( + result.unwrap().is_ok(), + "Load balancer should return Ok on shutdown" + ); +} + +#[tokio::test] +async fn test_shutdown_without_signal() { + // Test that shutdown channel works correctly when not triggered + let config_content = r#" +[load_balancer] +address = "127.0.0.1" +port = 18082 +algorithm = "round_robin" +layer = "L7" + +[[server]] +address = "127.0.0.1" +port = 13002 +max_connections = 1000 +weight = 1 +"#; + + let mut config_path = std::env::temp_dir(); + config_path.push("test_shutdown_no_signal.toml"); + fs::write(&config_path, config_content).unwrap(); + + let config = Config::new(&config_path); + let config_arc = Arc::new(Mutex::new(config)); + + // Create shutdown channel + let (_shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false); + + let lb = Layer7::new(config_arc); + + // Start the load balancer in a separate task + let lb_handle = tokio::spawn(async move { lb.start(shutdown_rx).await }); + + // Give it a moment to start + tokio::time::sleep(Duration::from_millis(100)).await; + + // Without sending a shutdown signal, the task should still be running + // We use a very short timeout to verify it's still running + let result = timeout(Duration::from_millis(200), lb_handle).await; + + // Clean up + fs::remove_file(config_path).ok(); + + // Assert that the load balancer is still running (timeout occurs) + assert!( + result.is_err(), + "Load balancer should still be running without shutdown signal" + ); +}