From 248799ca5d1d2dfe8bae474298fcb673c4b23897 Mon Sep 17 00:00:00 2001 From: Riz Ashraf Date: Sat, 12 Sep 2026 08:54:12 +0100 Subject: [PATCH] perf(server): offload sync disk/db IO to tokio blocking thread pool Makes apply_sync_write, write_to_local_delta, and store modify truly async, preventing the mcp-memory-server from locking the tokio executor during disk IO --- server/src/fix.py | 44 +++++++++++++++++++++++++++ server/src/handlers.rs | 16 +++++----- server/src/main.rs | 7 +++-- server/src/state.rs | 69 +++++++++++++++++++++++++----------------- server/src/store.rs | 9 ++++-- 5 files changed, 105 insertions(+), 40 deletions(-) create mode 100644 server/src/fix.py diff --git a/server/src/fix.py b/server/src/fix.py new file mode 100644 index 0000000..f0e8225 --- /dev/null +++ b/server/src/fix.py @@ -0,0 +1,44 @@ +import re + +def fix_handlers(filepath): + with open(filepath, "r") as f: + content = f.read() + + # Find the positions of self.state.write_to_local_delta( and self.state.apply_sync_write( + patterns = ["self.state.write_to_local_delta(|", "self.state.apply_sync_write(|"] + + # Actually, we can just find 'write_to_local_delta(|' and 'apply_sync_write(|' + + for method in ["write_to_local_delta", "apply_sync_write"]: + start_idx = 0 + while True: + idx = content.find(f"self.state.{method}(|", start_idx) + if idx == -1: + break + + # Find the closing bracket that matches the first `{` after `|` + brace_idx = content.find("{", idx) + + # count braces + brace_count = 1 + curr_idx = brace_idx + 1 + while brace_count > 0 and curr_idx < len(content): + if content[curr_idx] == '{': + brace_count += 1 + elif content[curr_idx] == '}': + brace_count -= 1 + curr_idx += 1 + + # curr_idx is now right after the matching '}' + # the original code was '});' + if content[curr_idx:curr_idx+2] == ");": + # we want to insert .await + content = content[:curr_idx+1] + ".await" + content[curr_idx+1:] + start_idx = curr_idx + 8 # move past `.await);` + else: + start_idx = curr_idx + + with open(filepath, "w") as f: + f.write(content) + +fix_handlers("C:\\Users\\reazul.ashraf\\workspace\\rust\\mcp-memory\\server\\src\\handlers.rs") diff --git a/server/src/handlers.rs b/server/src/handlers.rs index e74b781..bb26789 100644 --- a/server/src/handlers.rs +++ b/server/src/handlers.rs @@ -146,7 +146,7 @@ impl MemoryHandler { } } } - }); + }).await; Ok(vec!["Entities created".to_string()][0].clone()) } "create_relations" => { @@ -172,7 +172,7 @@ impl MemoryHandler { } } } - }); + }).await; Ok(vec!["Relations created".to_string()][0].clone()) } "add_observations" => { @@ -214,7 +214,7 @@ impl MemoryHandler { g.entities.insert(o.entity_name, e); } } - }); + }).await; Ok(vec!["Observations added".to_string()][0].clone()) } "delete_entities" => { @@ -235,7 +235,7 @@ impl MemoryHandler { master.relations.retain(|r| { !to_delete.contains(&r.from) && !to_delete.contains(&r.to) }); - }); + }).await; Ok(vec!["Entities deleted".to_string()][0].clone()) } "delete_observations" => { @@ -266,7 +266,7 @@ impl MemoryHandler { e.observations.retain(|o| !to_rem.contains(o)); } } - }); + }).await; Ok(vec!["Observations deleted".to_string()][0].clone()) } "delete_relations" => { @@ -299,7 +299,7 @@ impl MemoryHandler { r.from, r.to, r.relation_type, r.namespace )) }); - }); + }).await; Ok(vec!["Relations deleted".to_string()][0].clone()) } "read_graph" => { @@ -526,7 +526,7 @@ impl MemoryHandler { if let Some(e) = master.entities.get_mut(&req.entity_name) { e.observations = req.summarized_observations; } - }); + }).await; Ok(vec!["Entity condensed".to_string()][0].clone()) } "add_task" => { @@ -763,7 +763,7 @@ impl MemoryHandler { } } master.relations = MemoryState::unique_items(master.relations.clone()); - }); + }).await; Ok(vec!["Entities merged".to_string()][0].clone()) } "find_orphans" => { diff --git a/server/src/main.rs b/server/src/main.rs index bf29c44..f4fe6c8 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -91,8 +91,11 @@ async fn reconcile_worker(state: Arc) { }; let has_files = glob::glob(&pattern).map(|p| p.count() > 0).unwrap_or(false); if has_local || has_files { - state.apply_sync_write(|_master| {}); - state.rebuild_index(); + state.apply_sync_write(|_master| {}).await; + let state_clone = state.clone(); + let _ = tokio::task::spawn_blocking(move || { + state_clone.rebuild_index(); + }).await; } let now = SystemTime::now() diff --git a/server/src/state.rs b/server/src/state.rs index d4355b1..efd1c37 100644 --- a/server/src/state.rs +++ b/server/src/state.rs @@ -110,61 +110,74 @@ impl MemoryState { Self::merge_graphs(&mut master, &session_graph); master } - pub fn write_to_local_delta(&self, update_fn: F) { - let mut session_graph = self.session_graph.write().unwrap(); - update_fn(&mut session_graph); - let wal_path = self.base_dir.join("wal.jsonl"); - if let Ok(payload) = serde_json::to_string(&*session_graph) - && let Ok(mut file) = std::fs::OpenOptions::new() + pub async fn write_to_local_delta(&self, update_fn: F) { + let payload = { + let mut session_graph = self.session_graph.write().unwrap(); + update_fn(&mut session_graph); + serde_json::to_string(&*session_graph).ok() + }; + if let Some(payload) = payload { + let wal_path = self.base_dir.join("wal.jsonl"); + if let Ok(mut file) = tokio::fs::OpenOptions::new() .create(true) .append(true) .open(&wal_path) + .await { - use std::io::Write; - let _ = writeln!(file, "{}", payload); + use tokio::io::AsyncWriteExt; + let _ = file.write_all(payload.as_bytes()).await; + let _ = file.write_all(b"\n").await; } + } } - pub fn apply_sync_write(&self, update_fn: F) { + pub async fn apply_sync_write(&self, update_fn: F) { let lock_path = self.base_dir.join("master.lock"); let mut attempts = 0; loop { - if fs::OpenOptions::new() + if tokio::fs::OpenOptions::new() .create_new(true) .write(true) .open(&lock_path) + .await .is_ok() { break; } if attempts > 100 { - let _ = fs::remove_file(&lock_path); + let _ = tokio::fs::remove_file(&lock_path).await; } attempts += 1; - std::thread::sleep(Duration::from_millis(50)); + tokio::time::sleep(Duration::from_millis(50)).await; } let mut master = self.get_full_graph(); let wal_path = self.base_dir.join("wal.jsonl"); - let _ = fs::remove_file(&wal_path); + let _ = tokio::fs::remove_file(&wal_path).await; *self.session_graph.write().unwrap() = KnowledgeGraph::default(); update_fn(&mut master); - let write_json = |path: &std::path::Path, data: &KnowledgeGraph| -> std::io::Result<()> { - if path.exists() { - let bak_path = path.with_extension("json.bak"); - let _ = fs::copy(path, &bak_path); - } - let tmp_path = path.with_extension("json.tmp"); - let json_data = serde_json::to_string_pretty(data)?; - fs::write(&tmp_path, json_data)?; - fs::rename(&tmp_path, path) - }; - let _ = write_json(&self.master_path, &master); + let master_path = self.master_path.clone(); + let master_clone = master.clone(); + let _ = tokio::task::spawn_blocking(move || { + let write_json = |path: &std::path::Path, data: &KnowledgeGraph| -> std::io::Result<()> { + if path.exists() { + let bak_path = path.with_extension("json.bak"); + let _ = std::fs::copy(path, &bak_path); + } + let tmp_path = path.with_extension("json.tmp"); + let json_data = serde_json::to_string_pretty(data)?; + std::fs::write(&tmp_path, json_data)?; + std::fs::rename(&tmp_path, path) + }; + let _ = write_json(&master_path, &master_clone); + }).await; - let mut cache_lock = self.master_cache.write().unwrap(); - cache_lock.0 = master; - cache_lock.1 = self.master_mtime(); + { + let mut cache_lock = self.master_cache.write().unwrap(); + cache_lock.0 = master; + cache_lock.1 = self.master_mtime(); + } - let _ = fs::remove_file(&lock_path); + let _ = tokio::fs::remove_file(&lock_path).await; } pub fn rebuild_index(&self) { diff --git a/server/src/store.rs b/server/src/store.rs index 37efdf7..4f8d7cc 100644 --- a/server/src/store.rs +++ b/server/src/store.rs @@ -10,7 +10,7 @@ pub struct Store { pub cache: RwLock, } -impl Store { +impl Store { pub fn new(key: &str, db: Arc) -> Self { let initial_data = Self::load_from_db(key, &db); Self { @@ -50,6 +50,11 @@ impl Store { pub fn modify(&self, f: F) { let mut lock = self.cache.write().unwrap(); f(&mut lock); - Self::save_to_db(&self.key, &self.db, &lock); + let key = self.key.clone(); + let db = self.db.clone(); + let data = lock.clone(); + tokio::task::spawn_blocking(move || { + Self::save_to_db(&key, &db, &data); + }); } }