feat: improve error handling

src/computer.rs

@@ -23,6 +23,7 @@ pub fn start_sending(socket: &mut WebSocket<MaybeTlsStream<TcpStream>>) {
     let key_counter_clone = Arc::clone(&key_counter);
     let click_counter_clone = Arc::clone(&click_counter);
 
+    // Keys and clicks handler
     thread::spawn(move || {
         KeybdKey::bind_all(move |_| {
             key_counter_clone.fetch_add(1, Ordering::SeqCst);
@@ -38,6 +39,7 @@ pub fn start_sending(socket: &mut WebSocket<MaybeTlsStream<TcpStream>>) {
         inputbot::handle_input_events();
     });
 
+    // Send to WebSocket every 60 seconds
     loop {
         sys.refresh_cpu_usage();
         sys.refresh_memory();
@@ -51,7 +53,18 @@ pub fn start_sending(socket: &mut WebSocket<MaybeTlsStream<TcpStream>>) {
         let keys = key_counter.load(Ordering::SeqCst);
         let clicks = click_counter.load(Ordering::SeqCst);
 
-        websocket::send(socket, cpu_usage, memory_usage, keys, clicks);
+        match websocket::send(socket, cpu_usage, memory_usage, keys, clicks) {
+            Ok(_) => {
+                // Reset counters after sending
+                key_counter.store(0, Ordering::SeqCst);
+                click_counter.store(0, Ordering::SeqCst);
+            }
+            Err(e) => {
+                eprintln!("Failed to send WebSocket message: {}", e);
+                // Avoid resetting counters because we'll try to resend them after reconnection
+                break; // triggers reconnection in main.rs
+            }
+        }
 
         // Reset counters after sending
         key_counter.store(0, Ordering::SeqCst);

src/main.rs

@@ -1,4 +1,4 @@
-use std::error::Error;
+use std::{error::Error, thread, time::Duration};
 
 mod computer;
 mod notifications;
@@ -7,8 +7,23 @@ mod websocket;
 fn main() -> Result<(), Box<dyn Error>> {
     dotenvy::dotenv()?;
 
-    let mut socket = websocket::connect().expect("Could not connect to WebSocket");
-    computer::start_sending(&mut socket);
+    loop {
+        match websocket::connect() {
+            Ok(mut socket) => {
+                println!("WebSocket connected successfully");
 
-    Ok(())
+                // This will block until connection fails
+                computer::start_sending(&mut socket);
+
+                // The connection failed if code has reached here
+                println!("WebSocket connection lost, attempting to reconnect in 10 seconds...");
+            }
+            Err(_) => {
+                println!("Retrying connection in 10 seconds...");
+            }
+        }
+
+        // Wait 5 seconds before attempting to reconnect
+        thread::sleep(Duration::from_secs(5));
+    }
 }

src/websocket.rs

@@ -49,7 +49,7 @@ pub fn send(
     ram: u8,
     keys: u16,
     clicks: u16,
-) {
+) -> Result<(), tungstenite::Error> {
     let message = format!(
         "{{ \"cpu\": {}, \"ram\": {}, \"keys\": {}, \"clicks\": {} }}",
         cpu, ram, keys, clicks
@@ -57,5 +57,5 @@ pub fn send(
 
     println!("Sending to WebSocket: {}", message);
 
-    socket.send(Message::Text(message.into())).unwrap();
+    socket.send(Message::Text(message.into()))
 }