diff --git a/Cargo.toml b/Cargo.toml index fce9542..b63d003 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,7 +5,7 @@ edition = "2021" [dependencies] snmp = "0.2" -tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "time", "net", "process"] } +tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "time", "net", "process", "signal"] } tokio-tungstenite = { version = "0.21", features = ["rustls-tls-webpki-roots"] } futures = "0.3" serde = { version = "1.0", features = ["derive"] } diff --git a/src/main.rs b/src/main.rs index 69f5c5c..d9ef450 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,6 +10,7 @@ use std::env; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use tokio::sync::watch; use tokio::time::sleep; use websocket_client::AgentClient; @@ -201,6 +202,16 @@ async fn main() { let connected = Arc::new(AtomicBool::new(false)); let connected_for_health = Arc::clone(&connected); + // Create shutdown signal channel + let (shutdown_tx, shutdown_rx) = watch::channel(false); + + // Spawn signal handler for graceful shutdown + tokio::spawn(async move { + wait_for_shutdown_signal().await; + log_info!("Shutdown signal received, initiating graceful shutdown..."); + let _ = shutdown_tx.send(true); + }); + // Start simple health endpoint with connection state tokio::spawn(async move { if let Err(e) = health::start_health_server(8080, connected_for_health).await { @@ -214,6 +225,12 @@ async fn main() { let mut attempt = 0; loop { + // Check if shutdown was requested + if *shutdown_rx.borrow() { + log_info!("Shutdown requested, exiting main loop"); + break; + } + attempt += 1; if attempt > 1 { @@ -247,13 +264,56 @@ async fn main() { } }; - // Run the agent event loop - if let Err(e) = client.run().await { - log_error!("Agent disconnected: {}", e); - // Mark as disconnected for health check - connected.store(false, Ordering::Relaxed); - // Loop will retry with backoff + // Run the agent event loop with shutdown signal + match client.run(shutdown_rx.clone()).await { + Ok(()) => { + // Clean shutdown requested + if *shutdown_rx.borrow() { + log_info!("Agent shutdown complete"); + break; + } + } + Err(e) => { + log_error!("Agent disconnected: {}", e); + } } + + // Mark as disconnected for health check + connected.store(false, Ordering::Relaxed); + // Loop will retry with backoff (unless shutdown was requested) + } + + log_info!("Towerops agent stopped"); +} + +/// Wait for SIGTERM or SIGINT shutdown signal. +async fn wait_for_shutdown_signal() { + #[cfg(unix)] + { + use tokio::signal::unix::{signal, SignalKind}; + + let mut sigterm = + signal(SignalKind::terminate()).expect("Failed to register SIGTERM handler"); + let mut sigint = + signal(SignalKind::interrupt()).expect("Failed to register SIGINT handler"); + + tokio::select! { + _ = sigterm.recv() => { + log_info!("Received SIGTERM"); + } + _ = sigint.recv() => { + log_info!("Received SIGINT"); + } + } + } + + #[cfg(not(unix))] + { + // On non-Unix platforms, just wait for Ctrl+C + tokio::signal::ctrl_c() + .await + .expect("Failed to register Ctrl+C handler"); + log_info!("Received Ctrl+C"); } } diff --git a/src/websocket_client.rs b/src/websocket_client.rs index f867c30..dc0a010 100644 --- a/src/websocket_client.rs +++ b/src/websocket_client.rs @@ -10,12 +10,15 @@ use futures::{SinkExt, StreamExt}; use prost::Message; use std::collections::HashMap; use tokio::net::TcpStream; -use tokio::sync::mpsc; -use tokio::time::{interval, Duration}; +use tokio::sync::{mpsc, watch}; +use tokio::time::{interval, timeout, Duration}; use tokio_tungstenite::{ connect_async, tungstenite::protocol::Message as WsMessage, MaybeTlsStream, WebSocketStream, }; +/// Connection timeout for WebSocket establishment (30 seconds) +const CONNECTION_TIMEOUT: Duration = Duration::from_secs(30); + type Result = std::result::Result>; use crate::ping::ping; @@ -59,15 +62,31 @@ impl AgentClient { // Strip trailing slash from base URL to avoid double slashes let base_url = url.trim_end_matches('/'); let ws_url = format!("{}/socket/agent/websocket", base_url); - crate::log_info!("Connecting to WebSocket: {}", ws_url); + crate::log_info!( + "Connecting to WebSocket: {} (timeout: {}s)", + ws_url, + CONNECTION_TIMEOUT.as_secs() + ); - let (mut ws_stream, _) = connect_async(&ws_url) - .await - .map_err(|e| { + // Wrap connection in timeout to avoid hanging indefinitely on bad network + let (mut ws_stream, _) = match timeout(CONNECTION_TIMEOUT, connect_async(&ws_url)).await { + Ok(Ok(result)) => result, + Ok(Err(e)) => { crate::log_error!("WebSocket connection failed: {}", e); - e - }) - .map_err(|e| format!("Failed to connect to WebSocket: {}", e))?; + return Err(format!("Failed to connect to WebSocket: {}", e).into()); + } + Err(_) => { + crate::log_error!( + "WebSocket connection timed out after {}s", + CONNECTION_TIMEOUT.as_secs() + ); + return Err(format!( + "Connection timed out after {}s", + CONNECTION_TIMEOUT.as_secs() + ) + .into()); + } + }; crate::log_info!("Connected to Towerops server at {}", url); @@ -107,11 +126,24 @@ impl AgentClient { /// - Executing SNMP queries /// - Sending results back /// - Periodic heartbeats - pub async fn run(&mut self) -> Result<()> { + /// - Graceful shutdown on SIGTERM + pub async fn run(&mut self, mut shutdown_rx: watch::Receiver) -> Result<()> { let mut heartbeat_interval = interval(Duration::from_secs(60)); loop { tokio::select! { + // Check for shutdown signal (highest priority) + _ = shutdown_rx.changed() => { + if *shutdown_rx.borrow() { + crate::log_info!("Shutdown signal received, closing WebSocket connection gracefully"); + // Send close frame to server + if let Err(e) = self.ws_stream.close(None).await { + crate::log_warn!("Error sending WebSocket close frame: {}", e); + } + return Ok(()); + } + } + // Receive messages from server msg = self.ws_stream.next() => { match msg {