Parallelize the artifact graph only time suck (#6482)

* parallelize the artifact only time suck

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

updates

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

make wasm safe

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

updates

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

* artifact graph things

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>

---------

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

rust/kcl-lib/src/engine/async_tasks/mod.rs

@@ -0,0 +1,9 @@
+#[cfg(not(target_arch = "wasm32"))]
+pub mod tasks;
+#[cfg(target_arch = "wasm32")]
+pub mod tasks_wasm;
+
+#[cfg(not(target_arch = "wasm32"))]
+pub use tasks::AsyncTasks;
+#[cfg(target_arch = "wasm32")]
+pub use tasks_wasm::AsyncTasks;

rust/kcl-lib/src/engine/async_tasks/tasks.rs

@@ -0,0 +1,51 @@
+//! This module contains the `AsyncTasks` struct, which is used to manage a set of asynchronous
+//! tasks.
+
+use std::sync::Arc;
+
+use tokio::sync::RwLock;
+
+use crate::errors::KclError;
+
+#[derive(Debug, Clone)]
+pub struct AsyncTasks {
+    pub tasks: Arc<RwLock<tokio::task::JoinSet<anyhow::Result<(), KclError>>>>,
+}
+
+impl AsyncTasks {
+    pub fn new() -> Self {
+        Self {
+            tasks: Arc::new(RwLock::new(tokio::task::JoinSet::new())),
+        }
+    }
+
+    pub async fn spawn<F>(&mut self, task: F)
+    where
+        F: std::future::Future<Output = anyhow::Result<(), KclError>>,
+        F: Send + 'static,
+    {
+        self.tasks.write().await.spawn(task);
+    }
+
+    // Wait for all tasks to finish.
+    // Return an error if any of them failed.
+    pub async fn join_all(&mut self) -> anyhow::Result<(), KclError> {
+        let tasks = std::mem::take(&mut *self.tasks.write().await);
+        let results = tasks.join_all().await;
+        for result in results {
+            result?;
+        }
+
+        Ok(())
+    }
+
+    pub async fn clear(&mut self) {
+        *self.tasks.write().await = tokio::task::JoinSet::new();
+    }
+}
+
+impl Default for AsyncTasks {
+    fn default() -> Self {
+        Self::new()
+    }
+}

rust/kcl-lib/src/engine/async_tasks/tasks_wasm.rs

@@ -0,0 +1,89 @@
+//! This module contains the `AsyncTasks` struct, which is used to manage a set of asynchronous
+//! tasks.
+
+use std::{ops::AddAssign, sync::Arc};
+
+use tokio::sync::RwLock;
+
+use crate::errors::KclError;
+
+#[derive(Debug, Clone)]
+pub struct AsyncTasks {
+    pub sender: Arc<RwLock<tokio::sync::mpsc::Sender<Result<(), KclError>>>>,
+    pub receiver: Arc<RwLock<tokio::sync::mpsc::Receiver<Result<(), KclError>>>>,
+    pub sent: Arc<RwLock<usize>>,
+}
+
+impl AsyncTasks {
+    pub fn new() -> Self {
+        let (results_tx, results_rx) = tokio::sync::mpsc::channel(1);
+        Self {
+            sender: Arc::new(RwLock::new(results_tx)),
+            receiver: Arc::new(RwLock::new(results_rx)),
+            sent: Arc::new(RwLock::new(0)),
+        }
+    }
+
+    pub async fn spawn<F>(&mut self, task: F)
+    where
+        F: std::future::Future<Output = anyhow::Result<(), KclError>>,
+        F: Send + 'static,
+    {
+        // Add one to the sent counter.
+        self.sent.write().await.add_assign(1);
+
+        // Spawn the task and send the result to the channel.
+        let sender_clone = self.sender.clone();
+        wasm_bindgen_futures::spawn_local(async move {
+            let result = task.await;
+            let sender = sender_clone.read().await;
+            if let Err(_) = sender.send(result).await {
+                web_sys::console::error_1(&"Failed to send result".into());
+            }
+        });
+    }
+
+    // Wait for all tasks to finish.
+    // Return an error if any of them failed.
+    pub async fn join_all(&mut self) -> anyhow::Result<(), KclError> {
+        if *self.sent.read().await == 0 {
+            return Ok(());
+        }
+
+        let mut results = Vec::new();
+        let mut receiver = self.receiver.write().await;
+
+        // Wait for all tasks to finish.
+        while let Some(result) = receiver.recv().await {
+            results.push(result);
+
+            // Check if all tasks have finished.
+            if results.len() == *self.sent.read().await {
+                break;
+            }
+        }
+
+        // Check if any of the tasks failed.
+        for result in results {
+            result?;
+        }
+
+        Ok(())
+    }
+
+    pub async fn clear(&mut self) {
+        // Clear the sent counter.
+        *self.sent.write().await = 0;
+
+        // Clear the channel.
+        let (results_tx, results_rx) = tokio::sync::mpsc::channel(1);
+        *self.sender.write().await = results_tx;
+        *self.receiver.write().await = results_rx;
+    }
+}
+
+impl Default for AsyncTasks {
+    fn default() -> Self {
+        Self::new()
+    }
+}

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

@@ -18,11 +18,12 @@ use tokio::sync::{mpsc, oneshot, RwLock};
 use tokio_tungstenite::tungstenite::Message as WsMsg;
 use uuid::Uuid;
 
-use super::EngineStats;
+#[cfg(feature = "artifact-graph")]
+use crate::execution::ArtifactCommand;
 use crate::{
-    engine::EngineManager,
+    engine::{AsyncTasks, EngineManager, EngineStats},
     errors::{KclError, KclErrorDetails},
-    execution::{ArtifactCommand, DefaultPlanes, IdGenerator},
+    execution::{DefaultPlanes, IdGenerator},
     SourceRange,
 };
 
@@ -37,13 +38,14 @@ type WebSocketTcpWrite = futures::stream::SplitSink<tokio_tungstenite::WebSocket
 pub struct EngineConnection {
     engine_req_tx: mpsc::Sender<ToEngineReq>,
     shutdown_tx: mpsc::Sender<()>,
-    responses: Arc<RwLock<IndexMap<uuid::Uuid, WebSocketResponse>>>,
+    responses: ResponseInformation,
     pending_errors: Arc<RwLock<Vec<String>>>,
     #[allow(dead_code)]
     tcp_read_handle: Arc<TcpReadHandle>,
     socket_health: Arc<RwLock<SocketHealth>>,
     batch: Arc<RwLock<Vec<(WebSocketRequest, SourceRange)>>>,
     batch_end: Arc<RwLock<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>>,
+    #[cfg(feature = "artifact-graph")]
     artifact_commands: Arc<RwLock<Vec<ArtifactCommand>>>,
     ids_of_async_commands: Arc<RwLock<IndexMap<Uuid, SourceRange>>>,
 
@@ -53,6 +55,8 @@ pub struct EngineConnection {
     session_data: Arc<RwLock<Option<ModelingSessionData>>>,
 
     stats: EngineStats,
+
+    async_tasks: AsyncTasks,
 }
 
 pub struct TcpRead {
@@ -116,12 +120,14 @@ impl Drop for TcpReadHandle {
     }
 }
 
-struct ResponsesInformation {
+/// Information about the responses from the engine.
+#[derive(Clone, Debug)]
+struct ResponseInformation {
     /// The responses from the engine.
     responses: Arc<RwLock<IndexMap<uuid::Uuid, WebSocketResponse>>>,
 }
 
-impl ResponsesInformation {
+impl ResponseInformation {
     pub async fn add(&self, id: Uuid, response: WebSocketResponse) {
         self.responses.write().await.insert(id, response);
     }
@@ -238,14 +244,14 @@ impl EngineConnection {
 
         let session_data: Arc<RwLock<Option<ModelingSessionData>>> = Arc::new(RwLock::new(None));
         let session_data2 = session_data.clone();
-        let responses: Arc<RwLock<IndexMap<uuid::Uuid, WebSocketResponse>>> = Arc::new(RwLock::new(IndexMap::new()));
         let ids_of_async_commands: Arc<RwLock<IndexMap<Uuid, SourceRange>>> = Arc::new(RwLock::new(IndexMap::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 responses_information = ResponsesInformation {
-            responses: responses.clone(),
+        let response_information = ResponseInformation {
+            responses: Arc::new(RwLock::new(IndexMap::new())),
         };
+        let response_information_cloned = response_information.clone();
 
         let socket_health_tcp_read = socket_health.clone();
         let tcp_read_handle = tokio::spawn(async move {
@@ -270,7 +276,7 @@ impl EngineConnection {
                                         BatchResponse::Success { response } => {
                                             // If the id is in our ids of async commands, remove
                                             // it.
-                                            responses_information
+                                            response_information_cloned
                                                 .add(
                                                     id,
                                                     WebSocketResponse::Success(SuccessWebSocketResponse {
@@ -284,7 +290,7 @@ impl EngineConnection {
                                                 .await;
                                         }
                                         BatchResponse::Failure { errors } => {
-                                            responses_information
+                                            response_information_cloned
                                                 .add(
                                                     id,
                                                     WebSocketResponse::Failure(FailureWebSocketResponse {
@@ -311,7 +317,7 @@ impl EngineConnection {
                                 errors,
                             }) => {
                                 if let Some(id) = request_id {
-                                    responses_information
+                                    response_information_cloned
                                         .add(
                                             *id,
                                             WebSocketResponse::Failure(FailureWebSocketResponse {
@@ -336,7 +342,7 @@ impl EngineConnection {
                         }
 
                         if let Some(id) = id {
-                            responses_information.add(id, ws_resp.clone()).await;
+                            response_information_cloned.add(id, ws_resp.clone()).await;
                         }
                     }
                     Err(e) => {
@@ -357,16 +363,18 @@ impl EngineConnection {
             tcp_read_handle: Arc::new(TcpReadHandle {
                 handle: Arc::new(tcp_read_handle),
             }),
-            responses,
+            responses: response_information,
             pending_errors,
             socket_health,
             batch: Arc::new(RwLock::new(Vec::new())),
             batch_end: Arc::new(RwLock::new(IndexMap::new())),
+            #[cfg(feature = "artifact-graph")]
             artifact_commands: Arc::new(RwLock::new(Vec::new())),
             ids_of_async_commands,
             default_planes: Default::default(),
             session_data,
             stats: Default::default(),
+            async_tasks: AsyncTasks::new(),
         })
     }
 }
@@ -382,9 +390,10 @@ impl EngineManager for EngineConnection {
     }
 
     fn responses(&self) -> Arc<RwLock<IndexMap<Uuid, WebSocketResponse>>> {
-        self.responses.clone()
+        self.responses.responses.clone()
     }
 
+    #[cfg(feature = "artifact-graph")]
     fn artifact_commands(&self) -> Arc<RwLock<Vec<ArtifactCommand>>> {
         self.artifact_commands.clone()
     }
@@ -393,6 +402,10 @@ impl EngineManager for EngineConnection {
         self.ids_of_async_commands.clone()
     }
 
+    fn async_tasks(&self) -> AsyncTasks {
+        self.async_tasks.clone()
+    }
+
     fn stats(&self) -> &EngineStats {
         &self.stats
     }
@@ -483,9 +496,19 @@ impl EngineManager for EngineConnection {
                     }));
                 }
             }
-            // We pop off the responses to cleanup our mappings.
-            if let Some(resp) = self.responses.write().await.shift_remove(&id) {
-                return Ok(resp);
+
+            #[cfg(feature = "artifact-graph")]
+            {
+                // We cannot pop here or it will break the artifact graph.
+                if let Some(resp) = self.responses.responses.read().await.get(&id) {
+                    return Ok(resp.clone());
+                }
+            }
+            #[cfg(not(feature = "artifact-graph"))]
+            {
+                if let Some(resp) = self.responses.responses.write().await.shift_remove(&id) {
+                    return Ok(resp);
+                }
             }
         }
 

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

@@ -16,11 +16,13 @@ use kittycad_modeling_cmds::{self as kcmc, websocket::ModelingCmdReq, ImportFile
 use tokio::sync::RwLock;
 use uuid::Uuid;
 
-use super::EngineStats;
+#[cfg(feature = "artifact-graph")]
+use crate::execution::ArtifactCommand;
 use crate::{
+    engine::{AsyncTasks, EngineStats},
     errors::KclError,
     exec::DefaultPlanes,
-    execution::{ArtifactCommand, IdGenerator},
+    execution::IdGenerator,
     SourceRange,
 };
 
@@ -28,12 +30,14 @@ use crate::{
 pub struct EngineConnection {
     batch: Arc<RwLock<Vec<(WebSocketRequest, SourceRange)>>>,
     batch_end: Arc<RwLock<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>>,
+    #[cfg(feature = "artifact-graph")]
     artifact_commands: Arc<RwLock<Vec<ArtifactCommand>>>,
     ids_of_async_commands: Arc<RwLock<IndexMap<Uuid, SourceRange>>>,
     responses: Arc<RwLock<IndexMap<Uuid, WebSocketResponse>>>,
     /// The default planes for the scene.
     default_planes: Arc<RwLock<Option<DefaultPlanes>>>,
     stats: EngineStats,
+    async_tasks: AsyncTasks,
 }
 
 impl EngineConnection {
@@ -41,11 +45,13 @@ impl EngineConnection {
         Ok(EngineConnection {
             batch: Arc::new(RwLock::new(Vec::new())),
             batch_end: Arc::new(RwLock::new(IndexMap::new())),
+            #[cfg(feature = "artifact-graph")]
             artifact_commands: Arc::new(RwLock::new(Vec::new())),
             ids_of_async_commands: Arc::new(RwLock::new(IndexMap::new())),
             responses: Arc::new(RwLock::new(IndexMap::new())),
             default_planes: Default::default(),
             stats: Default::default(),
+            async_tasks: AsyncTasks::new(),
         })
     }
 }
@@ -68,6 +74,7 @@ impl crate::engine::EngineManager for EngineConnection {
         &self.stats
     }
 
+    #[cfg(feature = "artifact-graph")]
     fn artifact_commands(&self) -> Arc<RwLock<Vec<ArtifactCommand>>> {
         self.artifact_commands.clone()
     }
@@ -76,6 +83,10 @@ impl crate::engine::EngineManager for EngineConnection {
         self.ids_of_async_commands.clone()
     }
 
+    fn async_tasks(&self) -> AsyncTasks {
+        self.async_tasks.clone()
+    }
+
     fn get_default_planes(&self) -> Arc<RwLock<Option<DefaultPlanes>>> {
         self.default_planes.clone()
     }

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

@@ -11,7 +11,7 @@ use uuid::Uuid;
 use wasm_bindgen::prelude::*;
 
 use crate::{
-    engine::EngineStats,
+    engine::{AsyncTasks, EngineStats},
     errors::{KclError, KclErrorDetails},
     execution::{ArtifactCommand, DefaultPlanes, IdGenerator},
     SourceRange,
@@ -55,6 +55,7 @@ pub struct EngineConnection {
     /// The default planes for the scene.
     default_planes: Arc<RwLock<Option<DefaultPlanes>>>,
     stats: EngineStats,
+    async_tasks: AsyncTasks,
 }
 
 #[wasm_bindgen]
@@ -128,6 +129,7 @@ impl EngineConnection {
             ids_of_async_commands: Arc::new(RwLock::new(IndexMap::new())),
             default_planes: Default::default(),
             stats: Default::default(),
+            async_tasks: AsyncTasks::new(),
         })
     }
 
@@ -270,6 +272,10 @@ impl crate::engine::EngineManager for EngineConnection {
         self.ids_of_async_commands.clone()
     }
 
+    fn async_tasks(&self) -> AsyncTasks {
+        self.async_tasks.clone()
+    }
+
     fn get_default_planes(&self) -> Arc<RwLock<Option<DefaultPlanes>>> {
         self.default_planes.clone()
     }

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

@@ -1,5 +1,6 @@
 //! Functions for managing engine communications.
 
+pub mod async_tasks;
 #[cfg(not(target_arch = "wasm32"))]
 #[cfg(feature = "engine")]
 pub mod conn;
@@ -16,10 +17,12 @@ use std::{
     },
 };
 
+pub use async_tasks::AsyncTasks;
 use indexmap::IndexMap;
+#[cfg(feature = "artifact-graph")]
+use kcmc::id::ModelingCmdId;
 use kcmc::{
     each_cmd as mcmd,
-    id::ModelingCmdId,
     length_unit::LengthUnit,
     ok_response::OkModelingCmdResponse,
     shared::Color,
@@ -35,9 +38,11 @@ use serde::{Deserialize, Serialize};
 use tokio::sync::RwLock;
 use uuid::Uuid;
 
+#[cfg(feature = "artifact-graph")]
+use crate::execution::ArtifactCommand;
 use crate::{
     errors::{KclError, KclErrorDetails},
-    execution::{types::UnitLen, ArtifactCommand, DefaultPlanes, IdGenerator, Point3d},
+    execution::{types::UnitLen, DefaultPlanes, IdGenerator, Point3d},
     SourceRange,
 };
 
@@ -74,11 +79,15 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
     fn responses(&self) -> Arc<RwLock<IndexMap<Uuid, WebSocketResponse>>>;
 
     /// Get the artifact commands that have accumulated so far.
+    #[cfg(feature = "artifact-graph")]
     fn artifact_commands(&self) -> Arc<RwLock<Vec<ArtifactCommand>>>;
 
     /// Get the ids of the async commands we are waiting for.
     fn ids_of_async_commands(&self) -> Arc<RwLock<IndexMap<Uuid, SourceRange>>>;
 
+    /// Get the async tasks we are waiting for.
+    fn async_tasks(&self) -> AsyncTasks;
+
     /// Take the batch of commands that have accumulated so far and clear them.
     async fn take_batch(&self) -> Vec<(WebSocketRequest, SourceRange)> {
         std::mem::take(&mut *self.batch().write().await)
@@ -90,11 +99,13 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
     }
 
     /// Clear all artifact commands that have accumulated so far.
+    #[cfg(feature = "artifact-graph")]
     async fn clear_artifact_commands(&self) {
         self.artifact_commands().write().await.clear();
     }
 
     /// Take the artifact commands that have accumulated so far and clear them.
+    #[cfg(feature = "artifact-graph")]
     async fn take_artifact_commands(&self) -> Vec<ArtifactCommand> {
         std::mem::take(&mut *self.artifact_commands().write().await)
     }
@@ -145,6 +156,7 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
         self.batch().write().await.clear();
         self.batch_end().write().await.clear();
         self.ids_of_async_commands().write().await.clear();
+        self.async_tasks().clear().await;
     }
 
     /// Send a modeling command and do not wait for the response message.
@@ -186,6 +198,7 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
 
         // Ensure artifact commands are cleared so that we don't accumulate them
         // across runs.
+        #[cfg(feature = "artifact-graph")]
         self.clear_artifact_commands().await;
 
         // Do the after clear scene hook.
@@ -253,6 +266,18 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
             self.ensure_async_command_completed(id, Some(source_range)).await?;
         }
 
+        // Make sure we check for all async tasks as well.
+        // The reason why we ignore the error here is that, if a model fillets an edge
+        // we previously called something on, it might no longer exist. In which case,
+        // the artifact graph won't care either if its gone since you can't select it
+        // anymore anyways.
+        if let Err(err) = self.async_tasks().join_all().await {
+            crate::log::logln!("Error waiting for async tasks (this is typically fine and just means that an edge became something else): {:?}", err);
+        }
+
+        // Flush the batch to make sure nothing remains.
+        self.flush_batch(true, SourceRange::default()).await?;
+
         Ok(())
     }
 
@@ -273,6 +298,7 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
         Ok(())
     }
 
+    #[cfg(feature = "artifact-graph")]
     async fn handle_artifact_command(
         &self,
         cmd: &ModelingCmd,
@@ -413,6 +439,7 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
         self.ids_of_async_commands().write().await.insert(id, source_range);
 
         // Add to artifact commands.
+        #[cfg(feature = "artifact-graph")]
         self.handle_artifact_command(cmd, id.into(), &HashMap::from([(id, source_range)]))
             .await?;
 
@@ -486,6 +513,7 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
         }
 
         // Do the artifact commands.
+        #[cfg(feature = "artifact-graph")]
         for (req, _) in orig_requests.iter() {
             match &req {
                 WebSocketRequest::ModelingCmdBatchReq(ModelingBatch { requests, .. }) => {