race w batch

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

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

@@ -67,9 +67,17 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
     /// Get the batch of commands to be sent to the engine.
     fn batch(&self) -> Arc<RwLock<Vec<(WebSocketRequest, SourceRange)>>>;
 
+    async fn take_batch(&self) -> Vec<(WebSocketRequest, SourceRange)> {
+        std::mem::take(&mut *self.batch().write().await)
+    }
+
     /// Get the batch of end commands to be sent to the engine.
     fn batch_end(&self) -> Arc<RwLock<IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)>>>;
 
+    async fn take_batch_end(&self) -> IndexMap<uuid::Uuid, (WebSocketRequest, SourceRange)> {
+        std::mem::take(&mut *self.batch_end().write().await)
+    }
+
     /// Get the command responses from the engine.
     fn responses(&self) -> Arc<RwLock<IndexMap<Uuid, WebSocketResponse>>>;
 
@@ -330,12 +338,13 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
         batch_end: bool,
         source_range: SourceRange,
     ) -> Result<OkWebSocketResponseData, crate::errors::KclError> {
+        println!("Flushing batch queue");
         let all_requests = if batch_end {
-            let mut requests = self.batch().read().await.clone();
+            let mut requests = self.take_batch().await.clone();
-            requests.extend(self.batch_end().read().await.values().cloned());
+            requests.extend(self.take_batch_end().await.values().cloned());
             requests
         } else {
-            self.batch().read().await.clone()
+            self.take_batch().await.clone()
         };
 
         // Return early if we have no commands to send.
@@ -404,10 +413,6 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
         }
 
         // Throw away the old batch queue.
-        self.batch().write().await.clear();
-        if batch_end {
-            self.batch_end().write().await.clear();
-        }
         self.stats().batches_sent.fetch_add(1, Ordering::Relaxed);
 
         // We pop off the responses to cleanup our mappings.

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

@@ -175,7 +175,6 @@ impl ExecutorContext {
 
                     match &import_stmt.selector {
                         ImportSelector::List { items } => {
-                            println!("Importing items from module {}", import_stmt.path,);
                             let (env_ref, module_exports) =
                                 self.exec_module_for_items(module_id, exec_state, source_range).await?;
                             for import_item in items {
@@ -218,7 +217,6 @@ impl ExecutorContext {
                             }
                         }
                         ImportSelector::Glob(_) => {
-                            println!("Importing all items from module {}", import_stmt.path);
                             let (env_ref, module_exports) =
                                 self.exec_module_for_items(module_id, exec_state, source_range).await?;
                             for name in module_exports.iter() {
@@ -558,7 +556,6 @@ impl ExecutorContext {
         exec_state: &mut ExecState,
         source_range: SourceRange,
     ) -> Result<(Option<KclValue>, EnvironmentRef, Vec<String>), KclError> {
-        println!("exec_module_from_ast {path}");
         exec_state.global.mod_loader.enter_module(path);
         let result = self.exec_module_body(program, exec_state, false, module_id, path).await;
         exec_state.global.mod_loader.leave_module(path);

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

@@ -766,7 +766,6 @@ impl ExecutorContext {
                 {
                     wasm_bindgen_futures::spawn_local(async move {
                         //set.spawn(async move {
-                        println!("Running module {module} from run_concurrent");
                         let mut exec_state = exec_state;
                         let exec_ctxt = exec_ctxt;
 
@@ -789,7 +788,6 @@ impl ExecutorContext {
                 #[cfg(not(target_arch = "wasm32"))]
                 {
                     set.spawn(async move {
-                        println!("Running module {module} from run_concurrent");
                         let mut exec_state = exec_state;
                         let exec_ctxt = exec_ctxt;
 
@@ -813,7 +811,7 @@ impl ExecutorContext {
 
             drop(results_tx);
 
-            while let Some((module_id, _, result)) = results_rx.recv().await {
+            while let Some((module_id, module_path, result)) = results_rx.recv().await {
                 match result {
                     Ok((val, session_data, variables)) => {
                         let mut repr = exec_state.global.module_infos[&module_id].take_repr();
@@ -826,7 +824,24 @@ impl ExecutorContext {
                         exec_state.global.module_infos[&module_id].restore_repr(repr);
                     }
                     Err(e) => {
-                        return Err(KclErrorWithOutputs::no_outputs(e));
+                        println!("Error in module {:?}: {e}", module_path);
+                        let module_id_to_module_path: IndexMap<ModuleId, ModulePath> = exec_state
+                            .global
+                            .path_to_source_id
+                            .iter()
+                            .map(|(k, v)| ((*v), k.clone()))
+                            .collect();
+                        let default_planes = self.engine.get_default_planes().read().await.clone();
+
+                        return Err(KclErrorWithOutputs::new(
+                            e,
+                            exec_state.global.operations.clone(),
+                            exec_state.global.artifact_commands.clone(),
+                            exec_state.global.artifact_graph.clone(),
+                            module_id_to_module_path,
+                            exec_state.global.id_to_source.clone(),
+                            default_planes,
+                        ));
                     }
                 }
             }
@@ -883,11 +898,12 @@ impl ExecutorContext {
             )
         })?;
 
-        if !self.is_mock() {
+        // TODO: fix this
+        /* if !self.is_mock() {
             let mut mem = exec_state.stack().deep_clone();
             mem.restore_env(env_ref);
             cache::write_old_memory((mem, exec_state.global.module_infos.clone())).await;
-        }
+        }*/
         let session_data = self.engine.get_session_data().await;
         Ok((env_ref, session_data))
     }