Collect simple stats about engine usage (#5823)

Signed-off-by: Nick Cameron <nrc@ncameron.org>

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

@@ -18,7 +18,7 @@ use tokio::sync::{mpsc, oneshot, RwLock};
 use tokio_tungstenite::tungstenite::Message as WsMsg;
 use uuid::Uuid;
 
-use super::ExecutionKind;
+use super::{EngineStats, ExecutionKind};
 use crate::{
     engine::EngineManager,
     errors::{KclError, KclErrorDetails},
@@ -52,6 +52,7 @@ pub struct EngineConnection {
     session_data: Arc<RwLock<Option<ModelingSessionData>>>,
 
     execution_kind: Arc<RwLock<ExecutionKind>>,
+    stats: EngineStats,
 }
 
 pub struct TcpRead {
@@ -344,6 +345,7 @@ impl EngineConnection {
             default_planes: Default::default(),
             session_data,
             execution_kind: Default::default(),
+            stats: Default::default(),
         })
     }
 }
@@ -378,6 +380,10 @@ impl EngineManager for EngineConnection {
         original
     }
 
+    fn stats(&self) -> &EngineStats {
+        &self.stats
+    }
+
     fn get_default_planes(&self) -> Arc<RwLock<Option<DefaultPlanes>>> {
         self.default_planes.clone()
     }

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

@@ -16,7 +16,7 @@ use kittycad_modeling_cmds::{self as kcmc};
 use tokio::sync::RwLock;
 use uuid::Uuid;
 
-use super::ExecutionKind;
+use super::{EngineStats, ExecutionKind};
 use crate::{
     errors::KclError,
     exec::DefaultPlanes,
@@ -32,6 +32,7 @@ pub struct EngineConnection {
     execution_kind: Arc<RwLock<ExecutionKind>>,
     /// The default planes for the scene.
     default_planes: Arc<RwLock<Option<DefaultPlanes>>>,
+    stats: EngineStats,
 }
 
 impl EngineConnection {
@@ -42,6 +43,7 @@ impl EngineConnection {
             artifact_commands: Arc::new(RwLock::new(Vec::new())),
             execution_kind: Default::default(),
             default_planes: Default::default(),
+            stats: Default::default(),
         })
     }
 }
@@ -60,6 +62,10 @@ impl crate::engine::EngineManager for EngineConnection {
         Arc::new(RwLock::new(IndexMap::new()))
     }
 
+    fn stats(&self) -> &EngineStats {
+        &self.stats
+    }
+
     fn artifact_commands(&self) -> Arc<RwLock<Vec<ArtifactCommand>>> {
         self.artifact_commands.clone()
     }

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

@@ -11,7 +11,7 @@ use uuid::Uuid;
 use wasm_bindgen::prelude::*;
 
 use crate::{
-    engine::ExecutionKind,
+    engine::{EngineStats, ExecutionKind},
     errors::{KclError, KclErrorDetails},
     execution::{ArtifactCommand, DefaultPlanes, IdGenerator},
     SourceRange,
@@ -45,6 +45,7 @@ pub struct EngineConnection {
     execution_kind: Arc<RwLock<ExecutionKind>>,
     /// The default planes for the scene.
     default_planes: Arc<RwLock<Option<DefaultPlanes>>>,
+    stats: EngineStats,
 }
 
 // Safety: WebAssembly will only ever run in a single-threaded context.
@@ -62,6 +63,7 @@ impl EngineConnection {
             artifact_commands: Arc::new(RwLock::new(Vec::new())),
             execution_kind: Default::default(),
             default_planes: Default::default(),
+            stats: Default::default(),
         })
     }
 
@@ -141,6 +143,10 @@ impl crate::engine::EngineManager for EngineConnection {
         self.responses.clone()
     }
 
+    fn stats(&self) -> &EngineStats {
+        &self.stats
+    }
+
     fn artifact_commands(&self) -> Arc<RwLock<Vec<ArtifactCommand>>> {
         self.artifact_commands.clone()
     }

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

@@ -8,7 +8,13 @@ pub mod conn_mock;
 #[cfg(feature = "engine")]
 pub mod conn_wasm;
 
-use std::{collections::HashMap, sync::Arc};
+use std::{
+    collections::HashMap,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+};
 
 use indexmap::IndexMap;
 use kcmc::{
@@ -58,6 +64,21 @@ impl ExecutionKind {
     }
 }
 
+#[derive(Default, Debug)]
+pub struct EngineStats {
+    pub commands_batched: AtomicUsize,
+    pub batches_sent: AtomicUsize,
+}
+
+impl Clone for EngineStats {
+    fn clone(&self) -> Self {
+        Self {
+            commands_batched: AtomicUsize::new(self.commands_batched.load(Ordering::Relaxed)),
+            batches_sent: AtomicUsize::new(self.batches_sent.load(Ordering::Relaxed)),
+        }
+    }
+}
+
 #[async_trait::async_trait]
 pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
     /// Get the batch of commands to be sent to the engine.
@@ -97,6 +118,8 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
     /// Get the default planes.
     fn get_default_planes(&self) -> Arc<RwLock<Option<DefaultPlanes>>>;
 
+    fn stats(&self) -> &EngineStats;
+
     /// Get the default planes, creating them if they don't exist.
     async fn default_planes(
         &self,
@@ -254,6 +277,7 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
 
         // Add cmd to the batch.
         self.batch().write().await.push((req, source_range));
+        self.stats().commands_batched.fetch_add(1, Ordering::Relaxed);
 
         Ok(())
     }
@@ -277,6 +301,9 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
         for cmd in cmds {
             extended_cmds.push((WebSocketRequest::ModelingCmdReq(cmd.clone()), source_range));
         }
+        self.stats()
+            .commands_batched
+            .fetch_add(extended_cmds.len(), Ordering::Relaxed);
         self.batch().write().await.extend(extended_cmds);
 
         Ok(())
@@ -303,6 +330,7 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
 
         // Add cmd to the batch end.
         self.batch_end().write().await.insert(id, (req, source_range));
+        self.stats().commands_batched.fetch_add(1, Ordering::Relaxed);
         Ok(())
     }
 
@@ -405,6 +433,7 @@ pub trait EngineManager: std::fmt::Debug + Send + Sync + 'static {
         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.
         match final_req {