added error handling for coordinator

Author: adwinwhite

ui/frontend/websocketMiddleware.ts

@@ -88,14 +88,6 @@ export const websocketMiddleware =
 
         socket.addEventListener('open', () => {
           store.dispatch(websocketConnected());
-          // const params = new URLSearchParams(window.location.search);
-          if (socket != null) {
-            // if (params.has("pool")) {
-              socket.send("use container pool");
-            // } else {
-              // socket.send("use no pool");
-            // }
-          }
 
           wasConnected = true;
         });

ui/src/coordinator.rs

@@ -1,6 +1,6 @@
 use std::{collections::VecDeque, sync::{Arc, Mutex}};
 
-use crate::worker::{setup_worker, Container};
+use crate::worker::{spawn_worker, Container, self};
 
 use lazy_static::lazy_static;
 
@@ -20,8 +20,12 @@ impl Coordinator {
         }
     }
 
-    pub fn allocate(&mut self) -> Container {
-        self.free_containers.pop_front().unwrap_or_else(setup_worker)
+    pub fn allocate(&mut self) -> Result<Container, worker::Error> {
+        if let Some(container) = self.free_containers.pop_front() {
+            Ok(container)
+        } else {
+            Ok(spawn_worker()?)
+        }
     }
 
     pub fn recycle(&mut self, container: Container) -> bool {

ui/src/server_axum/websocket.rs

@@ -1,265 +1,65 @@
 use crate::{
+    coordinator::COORDINATOR,
     metrics::{DURATION_WS, LIVE_WS},
-    parse_channel, parse_crate_type, parse_edition, parse_mode,
-    sandbox::{self, Sandbox},
-    Error, ExecutionSnafu, Result, SandboxCreationSnafu, WebSocketTaskPanicSnafu, coordinator::{COORDINATOR, Coordinator},
 };
 
 use axum::extract::ws::{Message, WebSocket};
-use snafu::prelude::*;
-use std::{
-    convert::{TryFrom, TryInto},
-    time::Instant,
-};
-use tokio::{sync::mpsc, task::JoinSet};
-
-#[derive(serde::Deserialize)]
-#[serde(tag = "type")]
-enum WSMessageRequest {
-    #[serde(rename = "WS_EXECUTE_REQUEST")]
-    WSExecuteRequest(WSExecuteRequest),
-}
-
-#[derive(serde::Deserialize)]
-#[serde(rename_all = "camelCase")]
-struct WSExecuteRequest {
-    channel: String,
-    mode: String,
-    edition: String,
-    crate_type: String,
-    tests: bool,
-    code: String,
-    backtrace: bool,
-    extra: serde_json::Value,
-}
-
-impl TryFrom<WSExecuteRequest> for (sandbox::ExecuteRequest, serde_json::Value) {
-    type Error = Error;
-
-    fn try_from(value: WSExecuteRequest) -> Result<Self, Self::Error> {
-        let WSExecuteRequest {
-            channel,
-            mode,
-            edition,
-            crate_type,
-            tests,
-            code,
-            backtrace,
-            extra,
-        } = value;
-
-        let req = sandbox::ExecuteRequest {
-            channel: parse_channel(&channel)?,
-            mode: parse_mode(&mode)?,
-            edition: parse_edition(&edition)?,
-            crate_type: parse_crate_type(&crate_type)?,
-            tests,
-            backtrace,
-            code,
-        };
 
-        Ok((req, extra))
-    }
-}
-
-#[derive(Debug, serde::Serialize)]
-#[serde(tag = "type")]
-enum WSMessageResponse {
-    #[serde(rename = "WEBSOCKET_ERROR")]
-    Error(WSError),
-    #[serde(rename = "WS_EXECUTE_RESPONSE")]
-    WSExecuteResponse(WSExecuteResponse),
-}
-
-#[derive(Debug, serde::Serialize)]
-#[serde(rename_all = "camelCase")]
-struct WSError {
-    error: String,
-}
-
-#[derive(Debug, serde::Serialize)]
-#[serde(rename_all = "camelCase")]
-struct WSExecuteResponse {
-    success: bool,
-    stdout: String,
-    stderr: String,
-    extra: serde_json::Value,
-}
-
-impl From<(sandbox::ExecuteResponse, serde_json::Value)> for WSExecuteResponse {
-    fn from(value: (sandbox::ExecuteResponse, serde_json::Value)) -> Self {
-        let sandbox::ExecuteResponse {
-            success,
-            stdout,
-            stderr,
-        } = value.0;
-        let extra = value.1;
-
-        WSExecuteResponse {
-            success,
-            stdout,
-            stderr,
-            extra,
-        }
-    }
-}
+use std::time::Instant;
 
 pub async fn handle(mut socket: WebSocket) {
     LIVE_WS.inc();
     let start = Instant::now();
 
-    let mut use_pool = false;
-    if let Some(Ok(Message::Text(txt))) = socket.recv().await {
-        dbg!(&txt);
-        if txt == *"use container pool" {
-            use_pool = true;
-        }
-    };
-    if use_pool {
-        // let mut coordinator = Coordinator::new();
-        // Avoid task gets cancelled and make sure recycling happen.
-        tokio::spawn(async move {
-            let (ws_sender, mut ws_receiver) = COORDINATOR.lock().unwrap().allocate();
-            // let (ws_sender, mut ws_receiver) = coordinator.allocate();
-            loop {
-                tokio::select! {
-                    client_msg = socket.recv() => {
-                        match client_msg {
-                            None => {
-                                // browser disconnected
-                                break;
-                            }
-                            Some(Ok(Message::Text(txt))) => ws_sender.send(txt).await.unwrap(),
-                            Some(Ok(_)) => {
-                                // unknown message type
-                                continue;
-                            }
-                            Some(Err(e)) => super::record_websocket_error(e.to_string()),
-                        }
-                    },
-                    worker_msg = ws_receiver.recv() => {
-                        if let Some(txt) = worker_msg {
-                            socket.send(Message::Text(txt)).await.unwrap();
-                        } else {
-                            // Worker message channel closed.
-                            // Maybe it panicked or got killed.
-                            // TODO: send a notification to user and clean the worker process.
-                            break;
-                        }
-                    }
-                }
-            }
-            COORDINATOR.lock().unwrap().recycle((ws_sender, ws_receiver));
-            // coordinator.recycle((ws_sender, ws_receiver));
-        });
-    } else {
-        let (tx, mut rx) = mpsc::channel(3);
-        let mut tasks = JoinSet::new();
-
-        // TODO: Implement some kind of timeout to shutdown running work?
-
+    let allocated = COORDINATOR
+        .lock()
+        .unwrap_or_else(|e| e.into_inner())
+        .allocate();
+    // TODO: notify the user that container allocation failed.
+    if let Ok((ws_sender, mut ws_receiver)) = allocated {
         loop {
             tokio::select! {
-                request = socket.recv() => {
-                    match request {
+                client_msg = socket.recv() => {
+                    match client_msg {
                         None => {
                             // browser disconnected
                             break;
                         }
-                        Some(Ok(Message::Text(txt))) => handle_msg(txt, &tx, &mut tasks).await,
+                        Some(Ok(Message::Text(txt))) => {
+                            if ws_sender.send(txt).await.is_err() {
+                                // Worker exited.
+                                // TODO: reallocate a worker to this WebSocket.
+                                break;
+                            }
+                        }
                         Some(Ok(_)) => {
                             // unknown message type
                             continue;
                         }
                         Some(Err(e)) => super::record_websocket_error(e.to_string()),
                     }
                 },
-                resp = rx.recv() => {
-                    let resp = resp.expect("The rx should never close as we have a tx");
-                    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(),
+                worker_msg = ws_receiver.recv() => {
+                    if let Some(txt) = worker_msg {
+                        if (socket.send(Message::Text(txt)).await).is_err() {
+                            // We can't send a message to user.
+                            break;
                         }
-                    };
-                    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
+                    } else {
+                        // Worker exited.
+                        // TODO: reallocate a worker to this WebSocket.
                         break;
                     }
-                },
+                }
             }
         }
-
-        drop((tx, rx, socket));
-        tasks.shutdown().await;
+        COORDINATOR
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .recycle((ws_sender, ws_receiver));
     }
 
     LIVE_WS.dec();
     let elapsed = start.elapsed();
     DURATION_WS.observe(elapsed.as_secs_f64());
 }
-
-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();
-            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) => {
-            let resp = Err(e);
-            tx.send(resp).await.ok(/* We don't care if the channel is closed */);
-        }
-    }
-}
-
-async fn handle_execute(req: WSExecuteRequest) -> Result<WSMessageResponse> {
-    let sb = Sandbox::new().await.context(SandboxCreationSnafu)?;
-
-    let (req, extra) = req.try_into()?;
-    let resp = sb.execute(&req).await.context(ExecutionSnafu)?;
-    Ok(WSMessageResponse::WSExecuteResponse((resp, extra).into()))
-}