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

@@ -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")

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" => {

server/src/main.rs

@@ -91,8 +91,11 @@ async fn reconcile_worker(state: Arc<MemoryState>) {
         };
         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()

server/src/state.rs

@@ -110,61 +110,74 @@ impl MemoryState {
         Self::merge_graphs(&mut master, &session_graph);
         master
     }
-    pub fn write_to_local_delta<F: FnOnce(&mut KnowledgeGraph)>(&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<F: FnOnce(&mut KnowledgeGraph)>(&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<F: FnOnce(&mut KnowledgeGraph)>(&self, update_fn: F) {
+    pub async fn apply_sync_write<F: FnOnce(&mut KnowledgeGraph)>(&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) {

server/src/store.rs

@@ -10,7 +10,7 @@ pub struct Store<T> {
     pub cache: RwLock<T>,
 }
 
-impl<T: DeserializeOwned + Default + Serialize + Clone> Store<T> {
+impl<T: DeserializeOwned + Default + Serialize + Clone + Send + 'static> Store<T> {
     pub fn new(key: &str, db: Arc<Database>) -> Self {
         let initial_data = Self::load_from_db(key, &db);
         Self {
@@ -50,6 +50,11 @@ impl<T: DeserializeOwned + Default + Serialize + Clone> Store<T> {
     pub fn modify<F: FnOnce(&mut T)>(&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);
+        });
     }
 }