graceful shutdown on task error or Ctrl-C

.cargo/config.toml

@@ -0,0 +1,4 @@
+
+
+[build]
+rustflags = ["--cfg", "tokio_unstable"]

.gitignore

@@ -1 +1,2 @@
 /target
+.vscode

src/main.rs

@@ -1,13 +1,14 @@
 use anyhow::Result;
 use bytes::{Buf, BytesMut};
-use futures_util::future::try_join_all;
 use futures_util::{SinkExt, StreamExt};
 use serde_json::{json, Value};
 use std::path::PathBuf;
 use tokio::io::AsyncReadExt;
+use tokio::io::AsyncWriteExt;
 use tokio::net::{TcpListener, TcpStream};
+use tokio::signal;
 use tokio::sync::mpsc::{Receiver, Sender};
-use tokio::{io::AsyncWriteExt, spawn};
+use tokio::task::JoinSet;
 use tokio_tungstenite::accept_async;
 use wokwi_server::{GdbInstruction, SimulationPacket};
 
@@ -46,7 +47,7 @@ struct Args {
 }
 
 #[tokio::main]
-async fn main() -> Result<()> {
+async fn main() {
     let opts = Args::parse();
     if  !matches!(opts.chip, Chip::Esp32 | Chip::Esp32c3 | Chip::Esp32s2) {
         anyhow::bail!("Chip not supported in Wokwi. See available chips and features at https://docs.wokwi.com/guides/esp32#simulation-features");
@@ -55,12 +56,29 @@ async fn main() -> Result<()> {
     let (wsend, wrecv) = tokio::sync::mpsc::channel(1);
     let (gsend, grecv) = tokio::sync::mpsc::channel(1);
 
-    let main_wss = spawn(wokwi_task(opts, gsend, wrecv));
-    let gdb = spawn(gdb_task(wsend, grecv));
+    let mut set = JoinSet::new();
+    set.spawn(wokwi_task(opts, gsend, wrecv));
+    set.spawn(gdb_task(wsend, grecv));
 
-    try_join_all([main_wss, gdb]).await?;
-
-    Ok(())
+    loop {
+        tokio::select! {
+            _ = signal::ctrl_c() => {
+                set.shutdown().await;
+                break;
+            },
+            task = set.join_one() => {
+                match task.map_err(Into::into) {
+                    Ok(Some(Err(e))) | Err(e) => {
+                        println!("Task failed: {:?}", e);
+                        set.shutdown().await;
+                        break;
+                    }
+                    Ok(None) => break, /* All tasks completed */
+                    Ok(Some(Ok(_))) => {}, // graceful shutdown
+                }
+            }
+        }
+    }
 }
 
 async fn wokwi_task(
@@ -96,13 +114,10 @@ async fn wokwi_task(
     );
     opener::open_browser(url).ok(); // we don't care if this fails
 
-    while let Ok((stream, _)) = server.accept().await {
-        if let Err(e) = process(opts.clone(), stream, (&mut send, &mut recv)).await {
-            // only one connection at a time
-            println!("Woki websocket closed, error: {:?}", e);
-        }
+    loop {
+        let (stream, _) = server.accept().await?;
+        process(opts.clone(), stream, (&mut send, &mut recv)).await?;
     }
-    Ok(())
 }
 
 async fn process(
@@ -214,15 +229,14 @@ async fn process(
 
 async fn gdb_task(mut send: Sender<GdbInstruction>, mut recv: Receiver<String>) -> Result<()> {
     let server = TcpListener::bind(("127.0.0.1", GDB_PORT)).await?;
+    loop {
+        let (stream, _) = server.accept().await?;
 
-    while let Ok((stream, _)) = server.accept().await {
-        println!("GDB client connected");
         match handle_gdb_client(stream, &mut send, &mut recv).await {
             Ok(_) => println!("GDB Session ended cleanly."),
             Err(e) => println!("GDB Session ended with error: {:?}", e),
         }
     }
-    Ok(())
 }
 
 async fn handle_gdb_client(