Add 2nd server test and make Mutex more granular

bragaigor · bragaigor · commit 6ef41c902683 · 2026-01-27T11:39:16.000-05:00
Signed-off-by: Igor Braga &lt;5835477+bragaigor@users.noreply.github.com&gt;
diff --git a/crates/validation/src/lib.rs b/crates/validation/src/lib.rs
@@ -8,7 +8,7 @@ use std::{
 };
 
 /// Counterpart to Go `validator.GoGlobalState`.
-#[derive(Clone, Debug, Serialize, Deserialize)]
+#[derive(Clone, Debug, Serialize, Deserialize, Default)]
 #[serde(rename_all = "PascalCase")]
 pub struct GoGlobalState {
     #[serde(with = "As::<DisplayFromStr>")]
diff --git a/crates/validator/Cargo.toml b/crates/validator/Cargo.toml
@@ -25,7 +25,7 @@ tracing-subscriber = { workspace = true, features = [
     "json",
     "env-filter",
 ] }
+validation = { workspace = true }
 
 [dev-dependencies]
 reqwest = "0.13.1"
-validation = { workspace = true }
diff --git a/crates/validator/src/config.rs b/crates/validator/src/config.rs
@@ -13,7 +13,6 @@ use clap::{Args, Parser, ValueEnum};
 use std::fs::read_to_string;
 use std::net::SocketAddr;
 use std::path::PathBuf;
-use tokio::sync::Mutex;
 
 use crate::engine::config::JitMachineConfig;
 use crate::engine::machine::JitMachine;
@@ -23,7 +22,9 @@ pub struct ServerState {
     pub mode: InputMode,
     pub binary: PathBuf,
     pub module_root: Bytes32,
-    pub jit_machine: Option<Mutex<JitMachine>>,
+    /// jit machine responsible for computing next GlobalState. Not wrapped
+    /// in Arc<> since the caller of ServerState is wrapped in Arc<>
+    pub jit_machine: Option<JitMachine>,
     pub available_workers: usize,
 }
 
@@ -37,7 +38,7 @@ impl ServerState {
 
                 let jit_machine = JitMachine::new(&config, Some(module_root))?;
 
-                Some(Mutex::new(jit_machine))
+                Some(jit_machine)
             }
             InputMode::Native => None,
         };
diff --git a/crates/validator/src/engine/execution.rs b/crates/validator/src/engine/execution.rs
@@ -70,16 +70,16 @@ pub async fn validate_continuous(
         ));
     }
 
-    let mut locked_jit_machine = server_state.jit_machine.as_ref().unwrap().lock().await;
+    let jit_machine = server_state.jit_machine.as_ref().unwrap();
 
-    if !locked_jit_machine.is_active() {
+    if !jit_machine.is_active().await {
         return Err(format!(
             "Jit machine is not active. Maybe it received a shutdown signal? Requested module root: {}",
             server_state.module_root
         ));
     }
 
-    let new_state = locked_jit_machine
+    let new_state = jit_machine
         .feed_machine(&request)
         .await
         .map_err(|error| format!("{error:?}"))?;
diff --git a/crates/validator/src/engine/machine.rs b/crates/validator/src/engine/machine.rs
@@ -31,13 +31,12 @@ use tokio::{
     io::{AsyncRead, AsyncReadExt, AsyncWriteExt},
     net::{TcpListener, TcpStream},
     process::{Child, ChildStdin, Command},
+    sync::Mutex,
 };
 use tracing::{error, warn};
+use validation::{GoGlobalState, ValidationInput};
 
-use crate::{
-    engine::{config::JitMachineConfig, execution::ValidationRequest},
-    spawner_endpoints::{local_target, GlobalState},
-};
+use crate::{engine::config::JitMachineConfig, spawner_endpoints::local_target};
 
 const SUCCESS_BYTE: u8 = 0x0;
 const FAILURE_BYTE: u8 = 0x1;
@@ -80,10 +79,12 @@ async fn read_bytes_with_len<R: AsyncRead + Unpin>(reader: &mut R) -> Result<Vec
 
 #[derive(Debug)]
 pub struct JitMachine {
-    pub process_stdin: Option<ChildStdin>,
-    pub process: Child,
+    /// Handler to jit binary stdin. Instead of using Mutex<> for the entire
+    /// JitMachine we chose to use a more granular Mutex<> to avoid contention
+    pub process_stdin: Mutex<Option<ChildStdin>>,
+    /// Handler to jit binary process. Needs a Mutex<> to force quit on server shutdown
+    pub process: Mutex<Child>,
     pub wasm_memory_usage_limit: u64,
-    is_active: bool,
 }
 
 impl JitMachine {
@@ -136,39 +137,41 @@ impl JitMachine {
             .ok_or_else(|| anyhow!("failed to open stdin to jit process"))?;
 
         Ok(Self {
-            process_stdin: Some(stdin),
-            process: child,
+            process_stdin: Mutex::new(Some(stdin)),
+            process: Mutex::new(child),
             wasm_memory_usage_limit: config.wasm_memory_usage_limit,
-            is_active: true,
         })
     }
 
-    pub fn is_active(&self) -> bool {
-        self.is_active && self.process_stdin.is_some()
+    pub async fn is_active(&self) -> bool {
+        self.process_stdin.lock().await.is_some()
     }
 
-    pub async fn feed_machine(&mut self, request: &ValidationRequest) -> Result<GlobalState> {
+    pub async fn feed_machine(&self, request: &ValidationInput) -> Result<GoGlobalState> {
         // 1. Create new TCP connection
         // Binding with a port number of 0 will request that the OS assigns a port to this listener.
         let listener = TcpListener::bind("127.0.0.1:0")
             .await
             .context("failed to create TCP listener")?;
 
-        let mut state = GlobalState::default();
+        let mut state = GoGlobalState::default();
 
         let addr = listener.local_addr().context("failed to get local addr")?;
 
         // 2. Format the address string (Go: "%v\n")
         let address_str = format!("{addr}\n");
 
         // 3. Send TCP connection via stdin pipe
-        if let Some(stdin) = &mut self.process_stdin {
-            stdin
-                .write_all(address_str.as_bytes())
-                .await
-                .context("failed to write address to jit stdin")?;
-        } else {
-            return Err(anyhow!("JIT machine stdin is not available"));
+        {
+            let mut locked_process_stdin = self.process_stdin.lock().await;
+            if let Some(stdin) = locked_process_stdin.as_mut() {
+                stdin
+                    .write_all(address_str.as_bytes())
+                    .await
+                    .context("failed to write address to jit stdin")?;
+            } else {
+                return Err(anyhow!("JIT machine stdin is not available"));
+            }
         }
 
         // 4. Wait for the child to call us back
@@ -195,7 +198,7 @@ impl JitMachine {
         // 7. Send Delayed Inbox
         if request.has_delayed_msg {
             write_u8(&mut conn, ANOTHER_BYTE).await?;
-            write_u64(&mut conn, request.delayed_msg_number).await?;
+            write_u64(&mut conn, request.delayed_msg_nr).await?;
             write_bytes(&mut conn, &request.delayed_msg).await?;
         }
         write_u8(&mut conn, SUCCESS_BYTE).await?;
@@ -232,7 +235,7 @@ impl JitMachine {
         write_u32(&mut conn, local_user_wasm.len() as u32).await?;
         for (module_hash, program) in local_user_wasm {
             write_exact(&mut conn, &module_hash.0).await?;
-            write_bytes(&mut conn, program).await?;
+            write_bytes(&mut conn, &program.as_vec()).await?;
         }
 
         // 10. Signal that we are done sending global state
@@ -272,20 +275,21 @@ impl JitMachine {
         }
     }
 
-    pub async fn complete_machine(&mut self) -> Result<()> {
+    pub async fn complete_machine(&self) -> Result<()> {
         // Close stdin. This sends EOF to the child process, signaling it to stop.
         // We take the Option to ensure it's dropped and cannot be used again.
-        if let Some(stdin) = self.process_stdin.take() {
+
+        let mut locked_process_stdin = self.process_stdin.lock().await;
+        if let Some(stdin) = locked_process_stdin.take() {
             drop(stdin);
         }
 
-        self.process
+        let mut locked_process = self.process.lock().await;
+        locked_process
             .kill()
             .await
             .context("failed to kill jit process")?;
 
-        self.is_active = false;
-
         Ok(())
     }
 }
diff --git a/crates/validator/src/server.rs b/crates/validator/src/server.rs
@@ -8,7 +8,7 @@ use tracing::info;
 use crate::config::ServerState;
 use crate::router::create_router;
 
-pub(crate) async fn run_server(listener: TcpListener, state: Arc<ServerState>) -> Result<()> {
+pub async fn run_server(listener: TcpListener, state: Arc<ServerState>) -> Result<()> {
     run_server_internal(listener, state, shutdown_signal()).await
 }
 
@@ -27,8 +27,7 @@ where
     info!("Shutdown signal received. Running cleanup...");
 
     if let Some(jit_machine) = state.jit_machine.as_ref() {
-        let mut locked_jit_machine = jit_machine.lock().await;
-        locked_jit_machine.complete_machine().await?;
+        jit_machine.complete_machine().await?;
     }
 
     Ok(())
@@ -70,25 +69,27 @@ pub(crate) async fn shutdown_signal() {
 mod tests {
     use anyhow::Result;
     use clap::Parser;
-    use std::sync::Arc;
-    use tokio::{net::TcpListener, sync::oneshot};
+    use std::{net::SocketAddr, sync::Arc};
+    use tokio::{
+        net::TcpListener,
+        sync::oneshot::{self, Sender},
+        task::JoinHandle,
+    };
 
     use crate::{
         config::{ServerConfig, ServerState},
         server::run_server_internal,
     };
 
-    #[tokio::test]
-    async fn test_server_lifecycle() -> Result<()> {
-        // 1. Setup Config and State. Use dummy module root is okay.
-        let config = ServerConfig::try_parse_from([
-            "server",
-            "--module-root",
-            "0x0000000000000000000000000000000000000000000000000000000000000000",
-        ])
-        .unwrap();
-        let state = Arc::new(ServerState::new(&config)?);
+    struct TestServerConfig {
+        sender: Sender<()>,
+        server_handle: JoinHandle<Result<()>>,
+        addr: SocketAddr,
+        state: Arc<ServerState>,
+    }
 
+    async fn spinup_server(config: &ServerConfig) -> Result<TestServerConfig> {
+        let state = Arc::new(ServerState::new(config)?);
         // 2. Bind to random free port
         let listener = TcpListener::bind("127.0.0.1:0").await?;
         let addr = listener.local_addr()?;
@@ -107,16 +108,19 @@ mod tests {
             run_server_internal(listener, state_for_server, shutdown_signal).await
         });
 
-        // 5. Check that jit machine is active
-        if let Some(jit) = state.jit_machine.as_ref() {
-            let locked_jit_machine = jit.lock().await;
-            assert!(locked_jit_machine.is_active());
-        }
+        Ok(TestServerConfig {
+            sender: tx,
+            server_handle,
+            addr,
+            state,
+        })
+    }
 
-        // 6. Make a real request here to prove the server is up
+    async fn verify_and_shutdown_server(test_config: TestServerConfig) -> Result<()> {
+        // 5. Make a real request here to prove the server is up
         let client = reqwest::Client::new();
         let resp = client
-            .get(format!("http://{}/validation_capacity", addr))
+            .get(format!("http://{}/validation_capacity", test_config.addr))
             .send()
             .await;
 
@@ -126,28 +130,71 @@ mod tests {
         );
         assert_eq!(resp.unwrap().status(), 200);
 
-        // 7. Trigger Shutdown
+        // 6. Trigger Shutdown
         println!("Sending shutdown signal...");
-        let _ = tx.send(());
+        let _ = test_config.sender.send(());
 
-        // 8. Wait for the server to finish (this ensures cleanup ran)
-        let result = server_handle.await?;
+        // 7. Wait for the server to finish (this ensures cleanup ran)
+        let result = test_config.server_handle.await?;
         assert!(result.is_ok(), "Server should exit successfully");
 
-        // 9. Verify Cleanup
-        if let Some(jit) = state.jit_machine.as_ref() {
-            let locked_jit_machine = jit.lock().await;
-            assert!(!locked_jit_machine.is_active());
+        // 8. Verify jit_machine Cleanup
+        if let Some(jit) = test_config.state.jit_machine.as_ref() {
+            assert!(!jit.is_active().await);
         }
 
-        // 10. Verify same request from above fails expectadly
+        // 9. Verify same request from above fails expectadly
         let resp = client
-            .get(format!("http://{}/validation_capacity", addr))
+            .get(format!("http://{}/validation_capacity", test_config.addr))
             .send()
             .await;
 
         assert!(resp.is_err(), "server should not be up");
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_server_lifecycle_native_mode() -> Result<()> {
+        // 1. Setup Config and State. Use dummy module root is okay.
+        let config = ServerConfig::try_parse_from([
+            "server",
+            "--module-root",
+            "0x0000000000000000000000000000000000000000000000000000000000000000",
+        ])
+        .unwrap();
+        let test_config = spinup_server(&config).await?;
+
+        // Since we're running in native mode there should not be an active jit_machine
+        assert!(test_config.state.jit_machine.is_none());
+
+        verify_and_shutdown_server(test_config).await.unwrap();
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_server_lifecycle_continuous_mode() -> Result<()> {
+        // 1. Setup Config and State. Use dummy module root is okay.
+        let config = ServerConfig::try_parse_from([
+            "server",
+            "--module-root",
+            "0x0000000000000000000000000000000000000000000000000000000000000000",
+            "--mode",
+            "continuous",
+        ])
+        .unwrap();
+        let test_config = spinup_server(&config).await?;
+
+        assert!(test_config.state.jit_machine.is_some());
+
+        // Check that jit machine is active
+        if let Some(jit) = test_config.state.jit_machine.as_ref() {
+            assert!(jit.is_active().await);
+        }
+
+        verify_and_shutdown_server(test_config).await.unwrap();
+
+        Ok(())
+    }
 }
diff --git a/crates/validator/src/spawner_endpoints.rs b/crates/validator/src/spawner_endpoints.rs
@@ -8,7 +8,7 @@
 //! field names).
 
 use crate::engine::config::{TARGET_AMD_64, TARGET_ARM_64, TARGET_HOST};
-use crate::engine::execution::{validate_continuous, validate_native, ValidationRequest};
+use crate::engine::execution::{validate_continuous, validate_native};
 use crate::{config::InputMode, ServerState};
 use axum::extract::State;
 use axum::response::IntoResponse;
@@ -31,7 +31,7 @@ pub async fn validate(
     Json(request): Json<ValidationInput>,
 ) -> Result<Json<GoGlobalState>, String> {
     match state.mode {
-        InputMode::Native => validate_native(request).await,
+        InputMode::Native => validate_native(&state, request).await,
         InputMode::Continuous => validate_continuous(&state, request).await,
     }
 }
