change everything to rwlocks for thread safety (#5416)

* make everything in engine a rwlock and cleanup repetitive code

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* updates

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* updates

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* updates

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* updates

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* updates

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* docs

Signed-off-by: Jess Frazelle <github@jessfraz.com>

---------

Signed-off-by: Jess Frazelle <github@jessfraz.com>

src/wasm-lib/kcl/src/engine/conn.rs

@@ -1,13 +1,9 @@
 //! Functions for setting up our WebSocket and WebRTC connections for communications with the
 //! engine.
 
-use std::{
-    collections::HashMap,
-    sync::{Arc, Mutex},
-};
+use std::{collections::HashMap, sync::Arc};
 
 use anyhow::{anyhow, Result};
-use dashmap::DashMap;
 use futures::{SinkExt, StreamExt};
 use indexmap::IndexMap;
 use kcmc::{
@@ -17,9 +13,7 @@ use kcmc::{
     },
     ModelingCmd,
 };
-use kittycad_modeling_cmds::{
-    self as kcmc, id::ModelingCmdId, ok_response::OkModelingCmdResponse, websocket::ModelingBatch,
-};
+use kittycad_modeling_cmds::{self as kcmc, ok_response::OkModelingCmdResponse, websocket::ModelingBatch};
 use tokio::sync::{mpsc, oneshot, RwLock};
 use tokio_tungstenite::tungstenite::Message as WsMsg;
 use uuid::Uuid;
@@ -43,21 +37,21 @@ type WebSocketTcpWrite = futures::stream::SplitSink<tokio_tungstenite::WebSocket
 pub struct EngineConnection {
     engine_req_tx: mpsc::Sender<ToEngineReq>,
     shutdown_tx: mpsc::Sender<()>,
-    responses: Arc<DashMap<uuid::Uuid, WebSocketResponse>>,
-    pending_errors: Arc<Mutex<Vec<String>>>,
+    responses: Arc<RwLock<IndexMap<uuid::Uuid, WebSocketResponse>>>,
+    pending_errors: Arc<RwLock<Vec<String>>>,
     #[allow(dead_code)]
     tcp_read_handle: Arc<TcpReadHandle>,
-    socket_health: Arc<Mutex<SocketHealth>>,
-    batch: Arc<Mutex<Vec<(WebSocketRequest, SourceRange)>>>,
-    batch_end: Arc<Mutex<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>>,
-    artifact_commands: Arc<Mutex<Vec<ArtifactCommand>>>,
+    socket_health: Arc<RwLock<SocketHealth>>,
+    batch: Arc<RwLock<Vec<(WebSocketRequest, SourceRange)>>>,
+    batch_end: Arc<RwLock<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>>,
+    artifact_commands: Arc<RwLock<Vec<ArtifactCommand>>>,
 
     /// The default planes for the scene.
     default_planes: Arc<RwLock<Option<DefaultPlanes>>>,
     /// If the server sends session data, it'll be copied to here.
-    session_data: Arc<Mutex<Option<ModelingSessionData>>>,
+    session_data: Arc<RwLock<Option<ModelingSessionData>>>,
 
-    execution_kind: Arc<Mutex<ExecutionKind>>,
+    execution_kind: Arc<RwLock<ExecutionKind>>,
 }
 
 pub struct TcpRead {
@@ -230,12 +224,12 @@ impl EngineConnection {
 
         let mut tcp_read = TcpRead { stream: tcp_read };
 
-        let session_data: Arc<Mutex<Option<ModelingSessionData>>> = Arc::new(Mutex::new(None));
+        let session_data: Arc<RwLock<Option<ModelingSessionData>>> = Arc::new(RwLock::new(None));
         let session_data2 = session_data.clone();
-        let responses: Arc<DashMap<uuid::Uuid, WebSocketResponse>> = Arc::new(DashMap::new());
+        let responses: Arc<RwLock<IndexMap<uuid::Uuid, WebSocketResponse>>> = Arc::new(RwLock::new(IndexMap::new()));
         let responses_clone = responses.clone();
-        let socket_health = Arc::new(Mutex::new(SocketHealth::Active));
-        let pending_errors = Arc::new(Mutex::new(Vec::new()));
+        let socket_health = Arc::new(RwLock::new(SocketHealth::Active));
+        let pending_errors = Arc::new(RwLock::new(Vec::new()));
         let pending_errors_clone = pending_errors.clone();
 
         let socket_health_tcp_read = socket_health.clone();
@@ -260,7 +254,7 @@ impl EngineConnection {
                                     let id: uuid::Uuid = (*resp_id).into();
                                     match batch_response {
                                         BatchResponse::Success { response } => {
-                                            responses_clone.insert(
+                                            responses_clone.write().await.insert(
                                                 id,
                                                 WebSocketResponse::Success(SuccessWebSocketResponse {
                                                     success: true,
@@ -272,7 +266,7 @@ impl EngineConnection {
                                             );
                                         }
                                         BatchResponse::Failure { errors } => {
-                                            responses_clone.insert(
+                                            responses_clone.write().await.insert(
                                                 id,
                                                 WebSocketResponse::Failure(FailureWebSocketResponse {
                                                     success: false,
@@ -288,7 +282,7 @@ impl EngineConnection {
                                 resp: OkWebSocketResponseData::ModelingSessionData { session },
                                 ..
                             }) => {
-                                let mut sd = session_data2.lock().unwrap();
+                                let mut sd = session_data2.write().await;
                                 sd.replace(session.clone());
                             }
                             WebSocketResponse::Failure(FailureWebSocketResponse {
@@ -297,7 +291,7 @@ impl EngineConnection {
                                 errors,
                             }) => {
                                 if let Some(id) = request_id {
-                                    responses_clone.insert(
+                                    responses_clone.write().await.insert(
                                         *id,
                                         WebSocketResponse::Failure(FailureWebSocketResponse {
                                             success: false,
@@ -307,19 +301,20 @@ impl EngineConnection {
                                     );
                                 } else {
                                     // Add it to our pending errors.
-                                    let mut pe = pending_errors_clone.lock().unwrap();
+                                    let mut pe = pending_errors_clone.write().await;
                                     for error in errors {
                                         if !pe.contains(&error.message) {
                                             pe.push(error.message.clone());
                                         }
                                     }
+                                    drop(pe);
                                 }
                             }
                             _ => {}
                         }
 
                         if let Some(id) = id {
-                            responses_clone.insert(id, ws_resp.clone());
+                            responses_clone.write().await.insert(id, ws_resp.clone());
                         }
                     }
                     Err(e) => {
@@ -327,7 +322,7 @@ impl EngineConnection {
                             WebSocketReadError::Read(e) => crate::logln!("could not read from WS: {:?}", e),
                             WebSocketReadError::Deser(e) => crate::logln!("could not deserialize msg from WS: {:?}", e),
                         }
-                        *socket_health_tcp_read.lock().unwrap() = SocketHealth::Inactive;
+                        *socket_health_tcp_read.write().await = SocketHealth::Inactive;
                         return Err(e);
                     }
                 }
@@ -343,70 +338,41 @@ impl EngineConnection {
             responses,
             pending_errors,
             socket_health,
-            batch: Arc::new(Mutex::new(Vec::new())),
-            batch_end: Arc::new(Mutex::new(IndexMap::new())),
-            artifact_commands: Arc::new(Mutex::new(Vec::new())),
+            batch: Arc::new(RwLock::new(Vec::new())),
+            batch_end: Arc::new(RwLock::new(IndexMap::new())),
+            artifact_commands: Arc::new(RwLock::new(Vec::new())),
             default_planes: Default::default(),
             session_data,
             execution_kind: Default::default(),
         })
     }
-
-    fn handle_command(
-        &self,
-        cmd: &ModelingCmd,
-        cmd_id: ModelingCmdId,
-        id_to_source_range: &HashMap<Uuid, SourceRange>,
-    ) -> Result<(), KclError> {
-        let cmd_id = *cmd_id.as_ref();
-        let range = id_to_source_range
-            .get(&cmd_id)
-            .copied()
-            .ok_or_else(|| KclError::internal(format!("Failed to get source range for command ID: {:?}", cmd_id)))?;
-
-        // Add artifact command.
-        let mut artifact_commands = self.artifact_commands.lock().unwrap();
-        artifact_commands.push(ArtifactCommand {
-            cmd_id,
-            range,
-            command: cmd.clone(),
-        });
-        Ok(())
-    }
 }
 
 #[async_trait::async_trait]
 impl EngineManager for EngineConnection {
-    fn batch(&self) -> Arc<Mutex<Vec<(WebSocketRequest, SourceRange)>>> {
+    fn batch(&self) -> Arc<RwLock<Vec<(WebSocketRequest, SourceRange)>>> {
         self.batch.clone()
     }
 
-    fn batch_end(&self) -> Arc<Mutex<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>> {
+    fn batch_end(&self) -> Arc<RwLock<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>> {
         self.batch_end.clone()
     }
 
-    fn responses(&self) -> IndexMap<Uuid, WebSocketResponse> {
-        self.responses
-            .iter()
-            .map(|entry| {
-                let (k, v) = entry.pair();
-                (*k, v.clone())
-            })
-            .collect()
+    fn responses(&self) -> Arc<RwLock<IndexMap<Uuid, WebSocketResponse>>> {
+        self.responses.clone()
     }
 
-    fn take_artifact_commands(&self) -> Vec<ArtifactCommand> {
-        let mut artifact_commands = self.artifact_commands.lock().unwrap();
-        std::mem::take(&mut *artifact_commands)
+    fn artifact_commands(&self) -> Arc<RwLock<Vec<ArtifactCommand>>> {
+        self.artifact_commands.clone()
     }
 
-    fn execution_kind(&self) -> ExecutionKind {
-        let guard = self.execution_kind.lock().unwrap();
+    async fn execution_kind(&self) -> ExecutionKind {
+        let guard = self.execution_kind.read().await;
         *guard
     }
 
-    fn replace_execution_kind(&self, execution_kind: ExecutionKind) -> ExecutionKind {
-        let mut guard = self.execution_kind.lock().unwrap();
+    async fn replace_execution_kind(&self, execution_kind: ExecutionKind) -> ExecutionKind {
+        let mut guard = self.execution_kind.write().await;
         let original = *guard;
         *guard = execution_kind;
         original
@@ -447,22 +413,10 @@ impl EngineManager for EngineConnection {
         id: uuid::Uuid,
         source_range: SourceRange,
         cmd: WebSocketRequest,
-        id_to_source_range: HashMap<Uuid, SourceRange>,
+        _id_to_source_range: HashMap<Uuid, SourceRange>,
     ) -> Result<WebSocketResponse, KclError> {
-        match &cmd {
-            WebSocketRequest::ModelingCmdBatchReq(ModelingBatch { requests, .. }) => {
-                for request in requests {
-                    self.handle_command(&request.cmd, request.cmd_id, &id_to_source_range)?;
-                }
-            }
-            WebSocketRequest::ModelingCmdReq(request) => {
-                self.handle_command(&request.cmd, request.cmd_id, &id_to_source_range)?;
-            }
-            _ => {}
-        }
-
         // In isolated mode, we don't send the command to the engine.
-        if self.execution_kind().is_isolated() {
+        if self.execution_kind().await.is_isolated() {
             return match &cmd {
                 WebSocketRequest::ModelingCmdBatchReq(ModelingBatch { requests, .. }) => {
                     let mut responses = HashMap::with_capacity(requests.len());
@@ -524,25 +478,24 @@ impl EngineManager for EngineConnection {
         // Wait for the response.
         let current_time = std::time::Instant::now();
         while current_time.elapsed().as_secs() < 60 {
-            if let Ok(guard) = self.socket_health.lock() {
-                if *guard == SocketHealth::Inactive {
-                    // Check if we have any pending errors.
-                    let pe = self.pending_errors.lock().unwrap();
-                    if !pe.is_empty() {
-                        return Err(KclError::Engine(KclErrorDetails {
-                            message: pe.join(", ").to_string(),
-                            source_ranges: vec![source_range],
-                        }));
-                    } else {
-                        return Err(KclError::Engine(KclErrorDetails {
-                            message: "Modeling command failed: websocket closed early".to_string(),
-                            source_ranges: vec![source_range],
-                        }));
-                    }
+            let guard = self.socket_health.read().await;
+            if *guard == SocketHealth::Inactive {
+                // Check if we have any pending errors.
+                let pe = self.pending_errors.read().await;
+                if !pe.is_empty() {
+                    return Err(KclError::Engine(KclErrorDetails {
+                        message: pe.join(", ").to_string(),
+                        source_ranges: vec![source_range],
+                    }));
+                } else {
+                    return Err(KclError::Engine(KclErrorDetails {
+                        message: "Modeling command failed: websocket closed early".to_string(),
+                        source_ranges: vec![source_range],
+                    }));
                 }
             }
             // We pop off the responses to cleanup our mappings.
-            if let Some((_, resp)) = self.responses.remove(&id) {
+            if let Some(resp) = self.responses.write().await.shift_remove(&id) {
                 return Ok(resp);
             }
         }
@@ -553,17 +506,16 @@ impl EngineManager for EngineConnection {
         }))
     }
 
-    fn get_session_data(&self) -> Option<ModelingSessionData> {
-        self.session_data.lock().unwrap().clone()
+    async fn get_session_data(&self) -> Option<ModelingSessionData> {
+        self.session_data.read().await.clone()
     }
 
     async fn close(&self) {
         let _ = self.shutdown_tx.send(()).await;
         loop {
-            if let Ok(guard) = self.socket_health.lock() {
-                if *guard == SocketHealth::Inactive {
-                    return;
-                }
+            let guard = self.socket_health.read().await;
+            if *guard == SocketHealth::Inactive {
+                return;
             }
         }
     }

src/wasm-lib/kcl/src/engine/conn_mock.rs

@@ -1,10 +1,7 @@
 //! Functions for setting up our WebSocket and WebRTC connections for communications with the
 //! engine.
 
-use std::{
-    collections::HashMap,
-    sync::{Arc, Mutex},
-};
+use std::{collections::HashMap, sync::Arc};
 
 use anyhow::Result;
 use indexmap::IndexMap;
@@ -15,7 +12,8 @@ use kcmc::{
         WebSocketResponse,
     },
 };
-use kittycad_modeling_cmds::{self as kcmc, id::ModelingCmdId, ModelingCmd};
+use kittycad_modeling_cmds::{self as kcmc};
+use tokio::sync::RwLock;
 use uuid::Uuid;
 
 use super::ExecutionKind;
@@ -28,71 +26,48 @@ use crate::{
 
 #[derive(Debug, Clone)]
 pub struct EngineConnection {
-    batch: Arc<Mutex<Vec<(WebSocketRequest, SourceRange)>>>,
-    batch_end: Arc<Mutex<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>>,
-    artifact_commands: Arc<Mutex<Vec<ArtifactCommand>>>,
-    execution_kind: Arc<Mutex<ExecutionKind>>,
+    batch: Arc<RwLock<Vec<(WebSocketRequest, SourceRange)>>>,
+    batch_end: Arc<RwLock<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>>,
+    artifact_commands: Arc<RwLock<Vec<ArtifactCommand>>>,
+    execution_kind: Arc<RwLock<ExecutionKind>>,
 }
 
 impl EngineConnection {
     pub async fn new() -> Result<EngineConnection> {
         Ok(EngineConnection {
-            batch: Arc::new(Mutex::new(Vec::new())),
-            batch_end: Arc::new(Mutex::new(IndexMap::new())),
-            artifact_commands: Arc::new(Mutex::new(Vec::new())),
+            batch: Arc::new(RwLock::new(Vec::new())),
+            batch_end: Arc::new(RwLock::new(IndexMap::new())),
+            artifact_commands: Arc::new(RwLock::new(Vec::new())),
             execution_kind: Default::default(),
         })
     }
-
-    fn handle_command(
-        &self,
-        cmd: &ModelingCmd,
-        cmd_id: ModelingCmdId,
-        id_to_source_range: &HashMap<Uuid, SourceRange>,
-    ) -> Result<(), KclError> {
-        let cmd_id = *cmd_id.as_ref();
-        let range = id_to_source_range
-            .get(&cmd_id)
-            .copied()
-            .ok_or_else(|| KclError::internal(format!("Failed to get source range for command ID: {:?}", cmd_id)))?;
-
-        // Add artifact command.
-        let mut artifact_commands = self.artifact_commands.lock().unwrap();
-        artifact_commands.push(ArtifactCommand {
-            cmd_id,
-            range,
-            command: cmd.clone(),
-        });
-        Ok(())
-    }
 }
 
 #[async_trait::async_trait]
 impl crate::engine::EngineManager for EngineConnection {
-    fn batch(&self) -> Arc<Mutex<Vec<(WebSocketRequest, SourceRange)>>> {
+    fn batch(&self) -> Arc<RwLock<Vec<(WebSocketRequest, SourceRange)>>> {
         self.batch.clone()
     }
 
-    fn batch_end(&self) -> Arc<Mutex<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>> {
+    fn batch_end(&self) -> Arc<RwLock<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>> {
         self.batch_end.clone()
     }
 
-    fn responses(&self) -> IndexMap<Uuid, WebSocketResponse> {
-        IndexMap::new()
+    fn responses(&self) -> Arc<RwLock<IndexMap<Uuid, WebSocketResponse>>> {
+        Arc::new(RwLock::new(IndexMap::new()))
     }
 
-    fn take_artifact_commands(&self) -> Vec<ArtifactCommand> {
-        let mut artifact_commands = self.artifact_commands.lock().unwrap();
-        std::mem::take(&mut *artifact_commands)
+    fn artifact_commands(&self) -> Arc<RwLock<Vec<ArtifactCommand>>> {
+        self.artifact_commands.clone()
     }
 
-    fn execution_kind(&self) -> ExecutionKind {
-        let guard = self.execution_kind.lock().unwrap();
+    async fn execution_kind(&self) -> ExecutionKind {
+        let guard = self.execution_kind.read().await;
         *guard
     }
 
-    fn replace_execution_kind(&self, execution_kind: ExecutionKind) -> ExecutionKind {
-        let mut guard = self.execution_kind.lock().unwrap();
+    async fn replace_execution_kind(&self, execution_kind: ExecutionKind) -> ExecutionKind {
+        let mut guard = self.execution_kind.write().await;
         let original = *guard;
         *guard = execution_kind;
         original
@@ -119,7 +94,7 @@ impl crate::engine::EngineManager for EngineConnection {
         id: uuid::Uuid,
         _source_range: SourceRange,
         cmd: WebSocketRequest,
-        id_to_source_range: HashMap<Uuid, SourceRange>,
+        _id_to_source_range: HashMap<Uuid, SourceRange>,
     ) -> Result<WebSocketResponse, KclError> {
         match cmd {
             WebSocketRequest::ModelingCmdBatchReq(ModelingBatch {
@@ -130,7 +105,6 @@ impl crate::engine::EngineManager for EngineConnection {
                 // Create the empty responses.
                 let mut responses = HashMap::with_capacity(requests.len());
                 for request in requests {
-                    self.handle_command(&request.cmd, request.cmd_id, &id_to_source_range)?;
                     responses.insert(
                         request.cmd_id,
                         BatchResponse::Success {
@@ -144,17 +118,13 @@ impl crate::engine::EngineManager for EngineConnection {
                     success: true,
                 }))
             }
-            WebSocketRequest::ModelingCmdReq(request) => {
-                self.handle_command(&request.cmd, request.cmd_id, &id_to_source_range)?;
-
-                Ok(WebSocketResponse::Success(SuccessWebSocketResponse {
-                    request_id: Some(id),
-                    resp: OkWebSocketResponseData::Modeling {
-                        modeling_response: OkModelingCmdResponse::Empty {},
-                    },
-                    success: true,
-                }))
-            }
+            WebSocketRequest::ModelingCmdReq(_) => Ok(WebSocketResponse::Success(SuccessWebSocketResponse {
+                request_id: Some(id),
+                resp: OkWebSocketResponseData::Modeling {
+                    modeling_response: OkModelingCmdResponse::Empty {},
+                },
+                success: true,
+            })),
             _ => Ok(WebSocketResponse::Success(SuccessWebSocketResponse {
                 request_id: Some(id),
                 resp: OkWebSocketResponseData::Modeling {

src/wasm-lib/kcl/src/engine/conn_wasm.rs

@@ -1,25 +1,22 @@
 //! Functions for setting up our WebSocket and WebRTC connections for communications with the
 //! engine.
-use std::{
-    collections::HashMap,
-    sync::{Arc, Mutex},
-};
+use std::{collections::HashMap, sync::Arc};
 
 use anyhow::Result;
 use indexmap::IndexMap;
 use kcmc::{
-    id::ModelingCmdId,
     ok_response::OkModelingCmdResponse,
     websocket::{
         BatchResponse, ModelingBatch, OkWebSocketResponseData, SuccessWebSocketResponse, WebSocketRequest,
         WebSocketResponse,
     },
-    ModelingCmd,
 };
 use kittycad_modeling_cmds as kcmc;
+use tokio::sync::RwLock;
 use uuid::Uuid;
 use wasm_bindgen::prelude::*;
 
+use crate::engine::EngineManager;
 use crate::{
     engine::ExecutionKind,
     errors::{KclError, KclErrorDetails},
@@ -54,11 +51,11 @@ extern "C" {
 #[derive(Debug, Clone)]
 pub struct EngineConnection {
     manager: Arc<EngineCommandManager>,
-    batch: Arc<Mutex<Vec<(WebSocketRequest, SourceRange)>>>,
-    batch_end: Arc<Mutex<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>>,
-    responses: Arc<Mutex<IndexMap<Uuid, WebSocketResponse>>>,
-    artifact_commands: Arc<Mutex<Vec<ArtifactCommand>>>,
-    execution_kind: Arc<Mutex<ExecutionKind>>,
+    batch: Arc<RwLock<Vec<(WebSocketRequest, SourceRange)>>>,
+    batch_end: Arc<RwLock<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>>,
+    responses: Arc<RwLock<IndexMap<Uuid, WebSocketResponse>>>,
+    artifact_commands: Arc<RwLock<Vec<ArtifactCommand>>>,
+    execution_kind: Arc<RwLock<ExecutionKind>>,
 }
 
 // Safety: WebAssembly will only ever run in a single-threaded context.
@@ -70,66 +67,130 @@ impl EngineConnection {
         #[allow(clippy::arc_with_non_send_sync)]
         Ok(EngineConnection {
             manager: Arc::new(manager),
-            batch: Arc::new(Mutex::new(Vec::new())),
-            batch_end: Arc::new(Mutex::new(IndexMap::new())),
-            responses: Arc::new(Mutex::new(IndexMap::new())),
-            artifact_commands: Arc::new(Mutex::new(Vec::new())),
+            batch: Arc::new(RwLock::new(Vec::new())),
+            batch_end: Arc::new(RwLock::new(IndexMap::new())),
+            responses: Arc::new(RwLock::new(IndexMap::new())),
+            artifact_commands: Arc::new(RwLock::new(Vec::new())),
             execution_kind: Default::default(),
         })
     }
-}
 
-impl EngineConnection {
-    fn handle_command(
+    async fn do_send_modeling_cmd(
         &self,
-        cmd: &ModelingCmd,
-        cmd_id: ModelingCmdId,
-        id_to_source_range: &HashMap<Uuid, SourceRange>,
-    ) -> Result<(), KclError> {
-        let cmd_id = *cmd_id.as_ref();
-        let range = id_to_source_range
-            .get(&cmd_id)
-            .copied()
-            .ok_or_else(|| KclError::internal(format!("Failed to get source range for command ID: {:?}", cmd_id)))?;
+        id: uuid::Uuid,
+        source_range: SourceRange,
+        cmd: WebSocketRequest,
+        id_to_source_range: HashMap<uuid::Uuid, SourceRange>,
+    ) -> Result<WebSocketResponse, KclError> {
+        // In isolated mode, we don't send the command to the engine.
+        if self.execution_kind().await.is_isolated() {
+            return match &cmd {
+                WebSocketRequest::ModelingCmdBatchReq(ModelingBatch { requests, .. }) => {
+                    let mut responses = HashMap::with_capacity(requests.len());
+                    for request in requests {
+                        responses.insert(
+                            request.cmd_id,
+                            BatchResponse::Success {
+                                response: OkModelingCmdResponse::Empty {},
+                            },
+                        );
+                    }
+                    Ok(WebSocketResponse::Success(SuccessWebSocketResponse {
+                        request_id: Some(id),
+                        resp: OkWebSocketResponseData::ModelingBatch { responses },
+                        success: true,
+                    }))
+                }
+                _ => Ok(WebSocketResponse::Success(SuccessWebSocketResponse {
+                    request_id: Some(id),
+                    resp: OkWebSocketResponseData::Modeling {
+                        modeling_response: OkModelingCmdResponse::Empty {},
+                    },
+                    success: true,
+                })),
+            };
+        }
 
-        // Add artifact command.
-        let mut artifact_commands = self.artifact_commands.lock().unwrap();
-        artifact_commands.push(ArtifactCommand {
-            cmd_id,
-            range,
-            command: cmd.clone(),
-        });
-        Ok(())
+        let source_range_str = serde_json::to_string(&source_range).map_err(|e| {
+            KclError::Engine(KclErrorDetails {
+                message: format!("Failed to serialize source range: {:?}", e),
+                source_ranges: vec![source_range],
+            })
+        })?;
+        let cmd_str = serde_json::to_string(&cmd).map_err(|e| {
+            KclError::Engine(KclErrorDetails {
+                message: format!("Failed to serialize modeling command: {:?}", e),
+                source_ranges: vec![source_range],
+            })
+        })?;
+        let id_to_source_range_str = serde_json::to_string(&id_to_source_range).map_err(|e| {
+            KclError::Engine(KclErrorDetails {
+                message: format!("Failed to serialize id to source range: {:?}", e),
+                source_ranges: vec![source_range],
+            })
+        })?;
+
+        let promise = self
+            .manager
+            .send_modeling_cmd_from_wasm(id.to_string(), source_range_str, cmd_str, id_to_source_range_str)
+            .map_err(|e| {
+                KclError::Engine(KclErrorDetails {
+                    message: e.to_string().into(),
+                    source_ranges: vec![source_range],
+                })
+            })?;
+
+        let value = crate::wasm::JsFuture::from(promise).await.map_err(|e| {
+            KclError::Engine(KclErrorDetails {
+                message: format!("Failed to wait for promise from engine: {:?}", e),
+                source_ranges: vec![source_range],
+            })
+        })?;
+
+        // Parse the value as a string.
+        let s = value.as_string().ok_or_else(|| {
+            KclError::Engine(KclErrorDetails {
+                message: format!("Failed to get string from response from engine: `{:?}`", value),
+                source_ranges: vec![source_range],
+            })
+        })?;
+
+        let ws_result: WebSocketResponse = serde_json::from_str(&s).map_err(|e| {
+            KclError::Engine(KclErrorDetails {
+                message: format!("Failed to deserialize response from engine: {:?}", e),
+                source_ranges: vec![source_range],
+            })
+        })?;
+
+        Ok(ws_result)
     }
 }
 
 #[async_trait::async_trait]
 impl crate::engine::EngineManager for EngineConnection {
-    fn batch(&self) -> Arc<Mutex<Vec<(WebSocketRequest, SourceRange)>>> {
+    fn batch(&self) -> Arc<RwLock<Vec<(WebSocketRequest, SourceRange)>>> {
         self.batch.clone()
     }
 
-    fn batch_end(&self) -> Arc<Mutex<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>> {
+    fn batch_end(&self) -> Arc<RwLock<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>> {
         self.batch_end.clone()
     }
 
-    fn responses(&self) -> IndexMap<Uuid, WebSocketResponse> {
-        let responses = self.responses.lock().unwrap();
-        responses.clone()
+    fn responses(&self) -> Arc<RwLock<IndexMap<Uuid, WebSocketResponse>>> {
+        self.responses.clone()
     }
 
-    fn take_artifact_commands(&self) -> Vec<ArtifactCommand> {
-        let mut artifact_commands = self.artifact_commands.lock().unwrap();
-        std::mem::take(&mut *artifact_commands)
+    fn artifact_commands(&self) -> Arc<RwLock<Vec<ArtifactCommand>>> {
+        self.artifact_commands.clone()
     }
 
-    fn execution_kind(&self) -> ExecutionKind {
-        let guard = self.execution_kind.lock().unwrap();
+    async fn execution_kind(&self) -> ExecutionKind {
+        let guard = self.execution_kind.read().await;
         *guard
     }
 
-    fn replace_execution_kind(&self, execution_kind: ExecutionKind) -> ExecutionKind {
-        let mut guard = self.execution_kind.lock().unwrap();
+    async fn replace_execution_kind(&self, execution_kind: ExecutionKind) -> ExecutionKind {
+        let mut guard = self.execution_kind.write().await;
         let original = *guard;
         *guard = execution_kind;
         original
@@ -214,100 +275,18 @@ impl crate::engine::EngineManager for EngineConnection {
         cmd: WebSocketRequest,
         id_to_source_range: HashMap<uuid::Uuid, SourceRange>,
     ) -> Result<WebSocketResponse, KclError> {
-        match &cmd {
-            WebSocketRequest::ModelingCmdBatchReq(ModelingBatch { requests, .. }) => {
-                for request in requests {
-                    self.handle_command(&request.cmd, request.cmd_id, &id_to_source_range)?;
-                }
-            }
-            WebSocketRequest::ModelingCmdReq(request) => {
-                self.handle_command(&request.cmd, request.cmd_id, &id_to_source_range)?;
-            }
-            _ => {}
+        let ws_result = self
+            .do_send_modeling_cmd(id, source_range, cmd, id_to_source_range)
+            .await?;
+
+        // In isolated mode, we don't save the response.
+        if self.execution_kind().await.is_isolated() {
+            return Ok(ws_result);
         }
 
-        // In isolated mode, we don't send the command to the engine.
-        if self.execution_kind().is_isolated() {
-            return match &cmd {
-                WebSocketRequest::ModelingCmdBatchReq(ModelingBatch { requests, .. }) => {
-                    let mut responses = HashMap::with_capacity(requests.len());
-                    for request in requests {
-                        responses.insert(
-                            request.cmd_id,
-                            BatchResponse::Success {
-                                response: OkModelingCmdResponse::Empty {},
-                            },
-                        );
-                    }
-                    Ok(WebSocketResponse::Success(SuccessWebSocketResponse {
-                        request_id: Some(id),
-                        resp: OkWebSocketResponseData::ModelingBatch { responses },
-                        success: true,
-                    }))
-                }
-                _ => Ok(WebSocketResponse::Success(SuccessWebSocketResponse {
-                    request_id: Some(id),
-                    resp: OkWebSocketResponseData::Modeling {
-                        modeling_response: OkModelingCmdResponse::Empty {},
-                    },
-                    success: true,
-                })),
-            };
-        }
-
-        let source_range_str = serde_json::to_string(&source_range).map_err(|e| {
-            KclError::Engine(KclErrorDetails {
-                message: format!("Failed to serialize source range: {:?}", e),
-                source_ranges: vec![source_range],
-            })
-        })?;
-        let cmd_str = serde_json::to_string(&cmd).map_err(|e| {
-            KclError::Engine(KclErrorDetails {
-                message: format!("Failed to serialize modeling command: {:?}", e),
-                source_ranges: vec![source_range],
-            })
-        })?;
-        let id_to_source_range_str = serde_json::to_string(&id_to_source_range).map_err(|e| {
-            KclError::Engine(KclErrorDetails {
-                message: format!("Failed to serialize id to source range: {:?}", e),
-                source_ranges: vec![source_range],
-            })
-        })?;
-
-        let promise = self
-            .manager
-            .send_modeling_cmd_from_wasm(id.to_string(), source_range_str, cmd_str, id_to_source_range_str)
-            .map_err(|e| {
-                KclError::Engine(KclErrorDetails {
-                    message: e.to_string().into(),
-                    source_ranges: vec![source_range],
-                })
-            })?;
-
-        let value = crate::wasm::JsFuture::from(promise).await.map_err(|e| {
-            KclError::Engine(KclErrorDetails {
-                message: format!("Failed to wait for promise from engine: {:?}", e),
-                source_ranges: vec![source_range],
-            })
-        })?;
-
-        // Parse the value as a string.
-        let s = value.as_string().ok_or_else(|| {
-            KclError::Engine(KclErrorDetails {
-                message: format!("Failed to get string from response from engine: `{:?}`", value),
-                source_ranges: vec![source_range],
-            })
-        })?;
-
-        let ws_result: WebSocketResponse = serde_json::from_str(&s).map_err(|e| {
-            KclError::Engine(KclErrorDetails {
-                message: format!("Failed to deserialize response from engine: {:?}", e),
-                source_ranges: vec![source_range],
-            })
-        })?;
-
-        let mut responses = self.responses.lock().unwrap();
+        let mut responses = self.responses.write().await;
         responses.insert(id, ws_result.clone());
+        drop(responses);
 
         Ok(ws_result)
     }

src/wasm-lib/kcl/src/engine/mod.rs

@@ -8,14 +8,12 @@ pub mod conn_mock;
 #[cfg(feature = "engine")]
 pub mod conn_wasm;
 
-use std::{
-    collections::HashMap,
-    sync::{Arc, Mutex},
-};
+use std::{collections::HashMap, sync::Arc};
 
 use indexmap::IndexMap;
 use kcmc::{
     each_cmd as mcmd,
+    id::ModelingCmdId,
     length_unit::LengthUnit,
     ok_response::OkModelingCmdResponse,
     shared::Color,
@@ -28,6 +26,7 @@ use kcmc::{
 use kittycad_modeling_cmds as kcmc;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
+use tokio::sync::RwLock;
 use uuid::Uuid;
 
 use crate::{
@@ -62,28 +61,38 @@ impl ExecutionKind {
 #[async_trait::async_trait]
 pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
     /// Get the batch of commands to be sent to the engine.
-    fn batch(&self) -> Arc<Mutex<Vec<(WebSocketRequest, SourceRange)>>>;
+    fn batch(&self) -> Arc<RwLock<Vec<(WebSocketRequest, SourceRange)>>>;
 
     /// Get the batch of end commands to be sent to the engine.
-    fn batch_end(&self) -> Arc<Mutex<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>>;
+    fn batch_end(&self) -> Arc<RwLock<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>>;
 
     /// Get the command responses from the engine.
-    fn responses(&self) -> IndexMap<Uuid, WebSocketResponse>;
+    fn responses(&self) -> Arc<RwLock<IndexMap<Uuid, WebSocketResponse>>>;
 
-    /// Take the artifact commands generated up to this point and clear them.
-    fn take_artifact_commands(&self) -> Vec<ArtifactCommand>;
+    /// Get the artifact commands that have accumulated so far.
+    fn artifact_commands(&self) -> Arc<RwLock<Vec<ArtifactCommand>>>;
 
     /// Clear all artifact commands that have accumulated so far.
-    fn clear_artifact_commands(&self) {
-        self.take_artifact_commands();
+    async fn clear_artifact_commands(&self) {
+        self.artifact_commands().write().await.clear();
+    }
+
+    /// Take the artifact commands that have accumulated so far and clear them.
+    async fn take_artifact_commands(&self) -> Vec<ArtifactCommand> {
+        std::mem::take(&mut *self.artifact_commands().write().await)
+    }
+
+    /// Take the responses that have accumulated so far and clear them.
+    async fn take_responses(&self) -> IndexMap<Uuid, WebSocketResponse> {
+        std::mem::take(&mut *self.responses().write().await)
     }
 
     /// Get the current execution kind.
-    fn execution_kind(&self) -> ExecutionKind;
+    async fn execution_kind(&self) -> ExecutionKind;
 
     /// Replace the current execution kind with a new value and return the
     /// existing value.
-    fn replace_execution_kind(&self, execution_kind: ExecutionKind) -> ExecutionKind;
+    async fn replace_execution_kind(&self, execution_kind: ExecutionKind) -> ExecutionKind;
 
     /// Get the default planes.
     async fn default_planes(
@@ -127,7 +136,7 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
 
         // Ensure artifact commands are cleared so that we don't accumulate them
         // across runs.
-        self.clear_artifact_commands();
+        self.clear_artifact_commands().await;
 
         // Do the after clear scene hook.
         self.clear_scene_post_hook(id_generator, source_range).await?;
@@ -151,6 +160,27 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
         Ok(())
     }
 
+    async fn handle_artifact_command(
+        &self,
+        cmd: &ModelingCmd,
+        cmd_id: ModelingCmdId,
+        id_to_source_range: &HashMap<Uuid, SourceRange>,
+    ) -> Result<(), KclError> {
+        let cmd_id = *cmd_id.as_ref();
+        let range = id_to_source_range
+            .get(&cmd_id)
+            .copied()
+            .ok_or_else(|| KclError::internal(format!("Failed to get source range for command ID: {:?}", cmd_id)))?;
+
+        // Add artifact command.
+        self.artifact_commands().write().await.push(ArtifactCommand {
+            cmd_id,
+            range,
+            command: cmd.clone(),
+        });
+        Ok(())
+    }
+
     async fn set_units(
         &self,
         units: crate::UnitLength,
@@ -203,7 +233,7 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
         });
 
         // Add cmd to the batch.
-        self.batch().lock().unwrap().push((req, source_range));
+        self.batch().write().await.push((req, source_range));
 
         Ok(())
     }
@@ -223,7 +253,7 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
         });
 
         // Add cmd to the batch end.
-        self.batch_end().lock().unwrap().insert(id, (req, source_range));
+        self.batch_end().write().await.insert(id, (req, source_range));
         Ok(())
     }
 
@@ -249,11 +279,11 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
         source_range: SourceRange,
     ) -> Result<OkWebSocketResponseData, crate::errors::KclError> {
         let all_requests = if batch_end {
-            let mut requests = self.batch().lock().unwrap().clone();
-            requests.extend(self.batch_end().lock().unwrap().values().cloned());
+            let mut requests = self.batch().read().await.clone();
+            requests.extend(self.batch_end().read().await.values().cloned());
             requests
         } else {
-            self.batch().lock().unwrap().clone()
+            self.batch().read().await.clone()
         };
 
         // Return early if we have no commands to send.
@@ -304,10 +334,27 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
             }
         }
 
+        // Do the artifact commands.
+        for (req, _) in all_requests.iter() {
+            match &req {
+                WebSocketRequest::ModelingCmdBatchReq(ModelingBatch { requests, .. }) => {
+                    for request in requests {
+                        self.handle_artifact_command(&request.cmd, request.cmd_id, &id_to_source_range)
+                            .await?;
+                    }
+                }
+                WebSocketRequest::ModelingCmdReq(request) => {
+                    self.handle_artifact_command(&request.cmd, request.cmd_id, &id_to_source_range)
+                        .await?;
+                }
+                _ => {}
+            }
+        }
+
         // Throw away the old batch queue.
-        self.batch().lock().unwrap().clear();
+        self.batch().write().await.clear();
         if batch_end {
-            self.batch_end().lock().unwrap().clear();
+            self.batch_end().write().await.clear();
         }
 
         // We pop off the responses to cleanup our mappings.
@@ -596,7 +643,7 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
 
     /// Get session data, if it has been received.
     /// Returns None if the server never sent it.
-    fn get_session_data(&self) -> Option<ModelingSessionData> {
+    async fn get_session_data(&self) -> Option<ModelingSessionData> {
         None
     }
 

src/wasm-lib/kcl/src/execution/exec_ast.rs

@@ -48,7 +48,7 @@ impl ExecutorContext {
                     let old_units = exec_state.length_unit();
                     exec_state.mod_local.settings.update_from_annotation(annotation)?;
                     let new_units = exec_state.length_unit();
-                    if !self.engine.execution_kind().is_isolated() && old_units != new_units {
+                    if !self.engine.execution_kind().await.is_isolated() && old_units != new_units {
                         self.engine
                             .set_units(new_units.into(), annotation.as_source_range())
                             .await?;
@@ -393,7 +393,7 @@ impl ExecutorContext {
         exec_state.global.mod_loader.enter_module(path);
         std::mem::swap(&mut exec_state.mod_local, &mut local_state);
         exec_state.mut_memory().push_new_root_env();
-        let original_execution = self.engine.replace_execution_kind(exec_kind);
+        let original_execution = self.engine.replace_execution_kind(exec_kind).await;
 
         let result = self
             .exec_program(program, exec_state, crate::execution::BodyType::Root)
@@ -406,7 +406,7 @@ impl ExecutorContext {
         if !exec_kind.is_isolated() && new_units != old_units {
             self.engine.set_units(old_units.into(), Default::default()).await?;
         }
-        self.engine.replace_execution_kind(original_execution);
+        self.engine.replace_execution_kind(original_execution).await;
 
         result
             .map_err(|err| {

src/wasm-lib/kcl/src/execution/import.rs

@@ -289,7 +289,7 @@ pub struct PreImportedGeometry {
 }
 
 pub async fn send_to_engine(pre: PreImportedGeometry, ctxt: &ExecutorContext) -> Result<ImportedGeometry, KclError> {
-    if ctxt.no_engine_commands() {
+    if ctxt.no_engine_commands().await {
         return Ok(ImportedGeometry {
             id: pre.id,
             value: pre.command.files.iter().map(|f| f.path.to_string()).collect(),

src/wasm-lib/kcl/src/execution/mod.rs

@@ -484,8 +484,8 @@ impl ExecutorContext {
     }
 
     /// Returns true if we should not send engine commands for any reason.
-    pub fn no_engine_commands(&self) -> bool {
-        self.is_mock() || self.engine.execution_kind().is_isolated()
+    pub async fn no_engine_commands(&self) -> bool {
+        self.is_mock() || self.engine.execution_kind().await.is_isolated()
     }
 
     pub async fn send_clear_scene(
@@ -713,7 +713,7 @@ impl ExecutorContext {
             "Post interpretation KCL memory stats: {:#?}",
             exec_state.memory().stats
         ));
-        let session_data = self.engine.get_session_data();
+        let session_data = self.engine.get_session_data().await;
         Ok(session_data)
     }
 
@@ -734,8 +734,11 @@ impl ExecutorContext {
         exec_state
             .global
             .artifact_commands
-            .extend(self.engine.take_artifact_commands());
-        exec_state.global.artifact_responses.extend(self.engine.responses());
+            .extend(self.engine.take_artifact_commands().await);
+        exec_state
+            .global
+            .artifact_responses
+            .extend(self.engine.take_responses().await);
         // Build the artifact graph.
         match build_artifact_graph(
             &exec_state.global.artifact_commands,

src/wasm-lib/kcl/src/std/args.rs

@@ -277,12 +277,12 @@ impl Args {
         // before what ever we call next.
         for id in ids {
             // Pop it off the batch_end and add it to the batch.
-            let Some(item) = self.ctx.engine.batch_end().lock().unwrap().shift_remove(&id) else {
+            let Some(item) = self.ctx.engine.batch_end().write().await.shift_remove(&id) else {
                 // It might be in the batch already.
                 continue;
             };
             // Add it to the batch.
-            self.ctx.engine.batch().lock().unwrap().push(item);
+            self.ctx.engine.batch().write().await.push(item);
         }
 
         // Run flush.

src/wasm-lib/kcl/src/std/extrude.rs

@@ -232,8 +232,9 @@ pub(crate) async fn do_post_extrude(
         sides: face_id_map,
         start_cap_id,
         end_cap_id,
-    } = analyze_faces(exec_state, &args, face_infos);
+    } = analyze_faces(exec_state, &args, face_infos).await;
     // Iterate over the sketch.value array and add face_id to GeoMeta
+    let no_engine_commands = args.ctx.no_engine_commands().await;
     let new_value = sketch
         .paths
         .iter()
@@ -267,7 +268,7 @@ pub(crate) async fn do_post_extrude(
                         Some(extrude_surface)
                     }
                 }
-            } else if args.ctx.no_engine_commands() {
+            } else if no_engine_commands {
                 // Only pre-populate the extrude surface if we are in mock mode.
 
                 let extrude_surface = ExtrudeSurface::ExtrudePlane(crate::execution::ExtrudePlane {
@@ -313,12 +314,12 @@ struct Faces {
     start_cap_id: Option<Uuid>,
 }
 
-fn analyze_faces(exec_state: &mut ExecState, args: &Args, face_infos: Vec<ExtrusionFaceInfo>) -> Faces {
+async fn analyze_faces(exec_state: &mut ExecState, args: &Args, face_infos: Vec<ExtrusionFaceInfo>) -> Faces {
     let mut faces = Faces {
         sides: HashMap::with_capacity(face_infos.len()),
         ..Default::default()
     };
-    if args.ctx.no_engine_commands() {
+    if args.ctx.no_engine_commands().await {
         // Create fake IDs for start and end caps, to make extrudes mock-execute safe
         faces.start_cap_id = Some(exec_state.next_uuid());
         faces.end_cap_id = Some(exec_state.next_uuid());

src/wasm-lib/kcl/src/std/fillet.rs

@@ -222,7 +222,7 @@ pub async fn get_opposite_edge(exec_state: &mut ExecState, args: Args) -> Result
     name = "getOppositeEdge",
 }]
 async fn inner_get_opposite_edge(tag: TagIdentifier, exec_state: &mut ExecState, args: Args) -> Result<Uuid, KclError> {
-    if args.ctx.no_engine_commands() {
+    if args.ctx.no_engine_commands().await {
         return Ok(exec_state.next_uuid());
     }
     let face_id = args.get_adjacent_face_to_tag(exec_state, &tag, false).await?;
@@ -299,7 +299,7 @@ async fn inner_get_next_adjacent_edge(
     exec_state: &mut ExecState,
     args: Args,
 ) -> Result<Uuid, KclError> {
-    if args.ctx.no_engine_commands() {
+    if args.ctx.no_engine_commands().await {
         return Ok(exec_state.next_uuid());
     }
     let face_id = args.get_adjacent_face_to_tag(exec_state, &tag, false).await?;
@@ -384,7 +384,7 @@ async fn inner_get_previous_adjacent_edge(
     exec_state: &mut ExecState,
     args: Args,
 ) -> Result<Uuid, KclError> {
-    if args.ctx.no_engine_commands() {
+    if args.ctx.no_engine_commands().await {
         return Ok(exec_state.next_uuid());
     }
     let face_id = args.get_adjacent_face_to_tag(exec_state, &tag, false).await?;

src/wasm-lib/kcl/src/std/helix.rs

@@ -125,7 +125,7 @@ async fn inner_helix(
         meta: vec![args.source_range.into()],
     });
 
-    if args.ctx.no_engine_commands() {
+    if args.ctx.no_engine_commands().await {
         return Ok(helix_result);
     }
 

src/wasm-lib/kcl/src/std/mirror.rs

@@ -112,7 +112,7 @@ async fn inner_mirror_2d(
         SketchSet::Sketches(sketches) => sketches,
     };
 
-    if args.ctx.no_engine_commands() {
+    if args.ctx.no_engine_commands().await {
         return Ok(starting_sketches);
     }