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:
1 parent
1fd1d119e6
commit
248799ca5d
5 files changed
+92
-27
No files matched your search
@@ -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")
|
||||||
@@ -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
@@ -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
@@ -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
@@ -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);
|
||||||
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Reference in new issue
Block a user