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
This commit is contained in:
Riz Ashraf committed 2026-09-12 08:54:12 +01:00
1 parent 1fd1d119e6
commit 248799ca5d
5 files changed
+92 -27

No files matched your search

+44
View File
@@ -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")
+8 -8
View File
@@ -146,7 +146,7 @@ impl MemoryHandler {
} }
} }
} }
}); }).await;
Ok(vec!["Entities created".to_string()][0].clone()) Ok(vec!["Entities created".to_string()][0].clone())
} }
"create_relations" => { "create_relations" => {
@@ -172,7 +172,7 @@ impl MemoryHandler {
} }
} }
} }
}); }).await;
Ok(vec!["Relations created".to_string()][0].clone()) Ok(vec!["Relations created".to_string()][0].clone())
} }
"add_observations" => { "add_observations" => {
@@ -214,7 +214,7 @@ impl MemoryHandler {
g.entities.insert(o.entity_name, e); g.entities.insert(o.entity_name, e);
} }
} }
}); }).await;
Ok(vec!["Observations added".to_string()][0].clone()) Ok(vec!["Observations added".to_string()][0].clone())
} }
"delete_entities" => { "delete_entities" => {
@@ -235,7 +235,7 @@ impl MemoryHandler {
master.relations.retain(|r| { master.relations.retain(|r| {
!to_delete.contains(&r.from) && !to_delete.contains(&r.to) !to_delete.contains(&r.from) && !to_delete.contains(&r.to)
}); });
}); }).await;
Ok(vec!["Entities deleted".to_string()][0].clone()) Ok(vec!["Entities deleted".to_string()][0].clone())
} }
"delete_observations" => { "delete_observations" => {
@@ -266,7 +266,7 @@ impl MemoryHandler {
e.observations.retain(|o| !to_rem.contains(o)); e.observations.retain(|o| !to_rem.contains(o));
} }
} }
}); }).await;
Ok(vec!["Observations deleted".to_string()][0].clone()) Ok(vec!["Observations deleted".to_string()][0].clone())
} }
"delete_relations" => { "delete_relations" => {
@@ -299,7 +299,7 @@ impl MemoryHandler {
r.from, r.to, r.relation_type, r.namespace r.from, r.to, r.relation_type, r.namespace
)) ))
}); });
}); }).await;
Ok(vec!["Relations deleted".to_string()][0].clone()) Ok(vec!["Relations deleted".to_string()][0].clone())
} }
"read_graph" => { "read_graph" => {
@@ -526,7 +526,7 @@ impl MemoryHandler {
if let Some(e) = master.entities.get_mut(&req.entity_name) { if let Some(e) = master.entities.get_mut(&req.entity_name) {
e.observations = req.summarized_observations; e.observations = req.summarized_observations;
} }
}); }).await;
Ok(vec!["Entity condensed".to_string()][0].clone()) Ok(vec!["Entity condensed".to_string()][0].clone())
} }
"add_task" => { "add_task" => {
@@ -763,7 +763,7 @@ impl MemoryHandler {
} }
} }
master.relations = MemoryState::unique_items(master.relations.clone()); master.relations = MemoryState::unique_items(master.relations.clone());
}); }).await;
Ok(vec!["Entities merged".to_string()][0].clone()) Ok(vec!["Entities merged".to_string()][0].clone())
} }
"find_orphans" => { "find_orphans" => {
+5 -2
View File
@@ -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); let has_files = glob::glob(&pattern).map(|p| p.count() > 0).unwrap_or(false);
if has_local || has_files { if has_local || has_files {
state.apply_sync_write(|_master| {}); state.apply_sync_write(|_master| {}).await;
state.rebuild_index(); let state_clone = state.clone();
let _ = tokio::task::spawn_blocking(move || {
state_clone.rebuild_index();
}).await;
} }
let now = SystemTime::now() let now = SystemTime::now()
+28 -15
View File
@@ -110,61 +110,74 @@ impl MemoryState {
Self::merge_graphs(&mut master, &session_graph); Self::merge_graphs(&mut master, &session_graph);
master master
} }
pub fn write_to_local_delta<F: FnOnce(&mut KnowledgeGraph)>(&self, update_fn: F) { 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(); let mut session_graph = self.session_graph.write().unwrap();
update_fn(&mut session_graph); 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"); let wal_path = self.base_dir.join("wal.jsonl");
if let Ok(payload) = serde_json::to_string(&*session_graph) if let Ok(mut file) = tokio::fs::OpenOptions::new()
&& let Ok(mut file) = std::fs::OpenOptions::new()
.create(true) .create(true)
.append(true) .append(true)
.open(&wal_path) .open(&wal_path)
.await
{ {
use std::io::Write; use tokio::io::AsyncWriteExt;
let _ = writeln!(file, "{}", payload); 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 lock_path = self.base_dir.join("master.lock");
let mut attempts = 0; let mut attempts = 0;
loop { loop {
if fs::OpenOptions::new() if tokio::fs::OpenOptions::new()
.create_new(true) .create_new(true)
.write(true) .write(true)
.open(&lock_path) .open(&lock_path)
.await
.is_ok() .is_ok()
{ {
break; break;
} }
if attempts > 100 { if attempts > 100 {
let _ = fs::remove_file(&lock_path); let _ = tokio::fs::remove_file(&lock_path).await;
} }
attempts += 1; 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 mut master = self.get_full_graph();
let wal_path = self.base_dir.join("wal.jsonl"); 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(); *self.session_graph.write().unwrap() = KnowledgeGraph::default();
update_fn(&mut master); update_fn(&mut 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<()> { let write_json = |path: &std::path::Path, data: &KnowledgeGraph| -> std::io::Result<()> {
if path.exists() { if path.exists() {
let bak_path = path.with_extension("json.bak"); let bak_path = path.with_extension("json.bak");
let _ = fs::copy(path, &bak_path); let _ = std::fs::copy(path, &bak_path);
} }
let tmp_path = path.with_extension("json.tmp"); let tmp_path = path.with_extension("json.tmp");
let json_data = serde_json::to_string_pretty(data)?; let json_data = serde_json::to_string_pretty(data)?;
fs::write(&tmp_path, json_data)?; std::fs::write(&tmp_path, json_data)?;
fs::rename(&tmp_path, path) std::fs::rename(&tmp_path, path)
}; };
let _ = write_json(&self.master_path, &master); let _ = write_json(&master_path, &master_clone);
}).await;
{
let mut cache_lock = self.master_cache.write().unwrap(); let mut cache_lock = self.master_cache.write().unwrap();
cache_lock.0 = master; cache_lock.0 = master;
cache_lock.1 = self.master_mtime(); 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) { pub fn rebuild_index(&self) {
+7 -2
View File
@@ -10,7 +10,7 @@ pub struct Store<T> {
pub cache: RwLock<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 { pub fn new(key: &str, db: Arc<Database>) -> Self {
let initial_data = Self::load_from_db(key, &db); let initial_data = Self::load_from_db(key, &db);
Self { Self {
@@ -50,6 +50,11 @@ impl<T: DeserializeOwned + Default + Serialize + Clone> Store<T> {
pub fn modify<F: FnOnce(&mut T)>(&self, f: F) { pub fn modify<F: FnOnce(&mut T)>(&self, f: F) {
let mut lock = self.cache.write().unwrap(); let mut lock = self.cache.write().unwrap();
f(&mut lock); 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);
});
} }
} }