From 573c9586fd796c20f0b37eb597f58bf7ccd008c1 Mon Sep 17 00:00:00 2001 From: Riz Ashraf Date: Mon, 21 Sep 2026 08:22:25 +0100 Subject: [PATCH] fix: memory leaks and unoptimized websocket broadcasts --- server/src/main.rs | 4 ++-- stub/tests/e2e.rs | 35 ++++++++++++++++++++--------------- 2 files changed, 22 insertions(+), 17 deletions(-) diff --git a/server/src/main.rs b/server/src/main.rs index 36d6d20..cdab8ee 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -514,7 +514,7 @@ async fn handle_socket(socket: WebSocket, state: Arc, client_type: Str .collect(); for client_tx in senders { - let _ = client_tx.send(event.to_string()).await; + let _ = client_tx.try_send(event.to_string()); } } } // End if proxy @@ -642,7 +642,7 @@ async fn nvim_telemetry_handler( let msg_str = ws_msg.to_string(); let senders: Vec<_> = state.clients.read().unwrap().values().cloned().collect(); for tx in senders { - let _ = tx.send(msg_str.clone()).await; + let _ = tx.try_send(msg_str.clone()); } axum::Json(serde_json::json!({"status": "ok"})) diff --git a/stub/tests/e2e.rs b/stub/tests/e2e.rs index 96bdef4..080c7ed 100644 --- a/stub/tests/e2e.rs +++ b/stub/tests/e2e.rs @@ -1,8 +1,17 @@ use serde_json::{Value, json}; use std::io::{BufRead, BufReader, Write}; -use std::process::{Command, Stdio}; +use std::process::{Child, Command, Stdio}; use std::time::Duration; +struct ChildGuard(Child); + +impl Drop for ChildGuard { + fn drop(&mut self) { + let _ = self.0.kill(); + let _ = self.0.wait(); + } +} + fn send_message(stdin: &mut std::process::ChildStdin, msg: Value) { let s = serde_json::to_string(&msg).unwrap(); let payload = format!("{}\n", s); @@ -52,7 +61,7 @@ async fn test_full_system_e2e_performance() { assert!(stub_exe.exists(), "Stub not found at {:?}", stub_exe); // 1. Start Server - let mut server = Command::new(&server_exe) + let mut server = ChildGuard(Command::new(&server_exe) .arg("--daemon") .env("MCP_PORT", test_port) .env("RUST_LOG", "debug") @@ -62,7 +71,7 @@ async fn test_full_system_e2e_performance() { .stdout(Stdio::inherit()) .stderr(Stdio::inherit()) .spawn() - .expect("Failed to start server"); + .expect("Failed to start server")); // Give server time to generate TLS cert and start let client = reqwest::Client::builder() @@ -85,7 +94,7 @@ async fn test_full_system_e2e_performance() { assert!(started, "Server failed to start in time"); // 2. Start Stub - let mut stub = Command::new(&stub_exe) + let mut stub = ChildGuard(Command::new(&stub_exe) .arg("--target") .arg(format!("http://127.0.0.1:{}", test_port)) .env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap()) @@ -95,21 +104,21 @@ async fn test_full_system_e2e_performance() { .stdout(Stdio::piped()) .stderr(Stdio::inherit()) .spawn() - .expect("Failed to start stub"); + .expect("Failed to start stub")); - let mut stub_stdin = stub.stdin.take().unwrap(); - let mut stub_stdout = BufReader::new(stub.stdout.take().unwrap()); + let mut stub_stdin = stub.0.stdin.take().unwrap(); + let mut stub_stdout = BufReader::new(stub.0.stdout.take().unwrap()); // 3. Start Nvim Bridge - let mut nvim = Command::new(&nvim_exe) + let mut nvim = ChildGuard(Command::new(&nvim_exe) .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::inherit()) .spawn() - .expect("Failed to start nvim bridge"); + .expect("Failed to start nvim bridge")); - let mut nvim_stdin = nvim.stdin.take().unwrap(); - let mut nvim_stdout = BufReader::new(nvim.stdout.take().unwrap()); + let mut nvim_stdin = nvim.0.stdin.take().unwrap(); + let mut nvim_stdout = BufReader::new(nvim.0.stdout.take().unwrap()); println!("Server, stub, and nvim spawned successfully"); @@ -175,9 +184,5 @@ async fn test_full_system_e2e_performance() { println!("Stub 100 requests: {:?}", stub_duration); println!("Win-Nvim 100 requests: {:?}", nvim_duration); - // Cleanup - let _ = server.kill(); - let _ = stub.kill(); - let _ = nvim.kill(); let _ = std::fs::remove_dir_all(temp_dir); }