Watch spawned tasks for panics and report them

Author: shepmaster

ui/src/main.rs

@@ -202,6 +202,8 @@ pub enum Error {
     RequestMissing,
     #[snafu(display("The cache has been poisoned"))]
     CachePoisoned,
+    #[snafu(display("The WebSocket worker panicked: {}", text))]
+    WebSocketTaskPanic { text: String },
 }
 
 type Result<T, E = Error> = ::std::result::Result<T, E>;

ui/src/server_axum/websocket.rs

@@ -2,7 +2,7 @@ use crate::{
     metrics::{DURATION_WS, LIVE_WS},
     parse_channel, parse_crate_type, parse_edition, parse_mode,
     sandbox::{self, Sandbox},
-    Error, ExecutionSnafu, Result, SandboxCreationSnafu,
+    Error, ExecutionSnafu, Result, SandboxCreationSnafu, WebSocketTaskPanicSnafu,
 };
 
 use axum::extract::ws::{Message, WebSocket};
@@ -11,7 +11,7 @@ use std::{
     convert::{TryFrom, TryInto},
     time::Instant,
 };
-use tokio::sync::mpsc;
+use tokio::{sync::mpsc, task::JoinSet};
 
 #[derive(serde::Deserialize)]
 #[serde(tag = "type")]
@@ -109,6 +109,9 @@ pub async fn handle(mut socket: WebSocket) {
     let start = Instant::now();
 
     let (tx, mut rx) = mpsc::channel(3);
+    let mut tasks = JoinSet::new();
+
+    // TODO: Implement some kind of timeout to shutdown running work?
 
     loop {
         tokio::select! {
@@ -118,7 +121,7 @@ pub async fn handle(mut socket: WebSocket) {
                         // browser disconnected
                         break;
                     }
-                    Some(Ok(Message::Text(txt))) => handle_msg(txt, &tx).await,
+                    Some(Ok(Message::Text(txt))) => handle_msg(txt, &tx, &mut tasks).await,
                     Some(Ok(_)) => {
                         // unknown message type
                         continue;
@@ -128,10 +131,31 @@ pub async fn handle(mut socket: WebSocket) {
             },
             resp = rx.recv() => {
                 let resp = resp.expect("The rx should never close as we have a tx");
-                let resp = resp.unwrap_or_else(|e| WSMessageResponse::Error(WSError { error: e.to_string() }));
-                const LAST_CHANCE_ERROR: &str = r#"{ "type": "WEBSOCKET_ERROR", "error": "Unable to serialize JSON" }"#;
-                let resp = serde_json::to_string(&resp).unwrap_or_else(|_| LAST_CHANCE_ERROR.into());
-                let resp = Message::Text(resp);
+                let resp = resp.unwrap_or_else(error_to_response);
+                let resp = response_to_message(resp);
+
+                if let Err(_) = socket.send(resp).await {
+                    // We can't send a response
+                    break;
+                }
+            },
+            // We don't care if there are no running tasks
+            Some(task) = tasks.join_next() => {
+                let Err(error) = task else { continue };
+                // The task was cancelled; no need to report
+                let Ok(panic) = error.try_into_panic() else { continue };
+
+                let text = match panic.downcast::<String>() {
+                    Ok(text) => *text,
+                    Err(panic) => match panic.downcast::<&str>() {
+                        Ok(text) => text.to_string(),
+                        _ => "An unknown panic occurred".into(),
+                    }
+                };
+                let error = WebSocketTaskPanicSnafu { text }.build();
+
+                let resp = error_to_response(error);
+                let resp = response_to_message(resp);
 
                 if let Err(_) = socket.send(resp).await {
                     // We can't send a response
@@ -141,22 +165,42 @@ pub async fn handle(mut socket: WebSocket) {
         }
     }
 
+    drop((tx, rx, socket));
+    tasks.shutdown().await;
+
     LIVE_WS.dec();
     let elapsed = start.elapsed();
     DURATION_WS.observe(elapsed.as_secs_f64());
 }
 
-async fn handle_msg(txt: String, tx: &mpsc::Sender<Result<WSMessageResponse>>) {
+fn error_to_response(error: Error) -> WSMessageResponse {
+    let error = error.to_string();
+    WSMessageResponse::Error(WSError { error })
+}
+
+fn response_to_message(response: WSMessageResponse) -> Message {
+    const LAST_CHANCE_ERROR: &str =
+        r#"{ "type": "WEBSOCKET_ERROR", "error": "Unable to serialize JSON" }"#;
+    let resp = serde_json::to_string(&response).unwrap_or_else(|_| LAST_CHANCE_ERROR.into());
+    Message::Text(resp)
+}
+
+async fn handle_msg(
+    txt: String,
+    tx: &mpsc::Sender<Result<WSMessageResponse>>,
+    tasks: &mut JoinSet<Result<()>>,
+) {
     use WSMessageRequest::*;
 
     let msg = serde_json::from_str(&txt).context(crate::DeserializationSnafu);
 
     match msg {
         Ok(WSExecuteRequest(req)) => {
             let tx = tx.clone();
-            tokio::spawn(async move {
+            tasks.spawn(async move {
                 let resp = handle_execute(req).await;
                 tx.send(resp).await.ok(/* We don't care if the channel is closed */);
+                Ok(())
             });
         }
         Err(e) => {