websocket improvements

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"] }

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");
     }
 }
 

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<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
 
 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<bool>) -> 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 {