riz/mcp-memory
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(server): fix broadcast loops, avoid locking db on index, eliminate O(N*M) search, and fix db loss

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-09-22 06:33:32 +01:00
1 parent 1e39122e05
commit 27aba454e2
4 files changed
+69 -54
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/handlers/meta.rs
+21 -15
View File
@@ -361,14 +361,28 @@ impl McpTool for OmniSearchHandler {
serde_json::to_value(&kg_entities).unwrap_or_default() serde_json::to_value(&kg_entities).unwrap_or_default()
}); });
let mut matched_tasks = std::collections::HashSet::new();
let mut matched_snippets = std::collections::HashSet::new();
let mut matched_adrs = std::collections::HashSet::new();
for (id, typ, _, _, _) in &matches {
match typ.as_str() {
"task" => {
matched_tasks.insert(id.as_str());
}
"snippet" => {
matched_snippets.insert(id.as_str());
}
"adr" => {
matched_adrs.insert(id.as_str());
}
_ => {}
}
}
let tasks_json = state.tasks.read_with(|all_tasks| { let tasks_json = state.tasks.read_with(|all_tasks| {
let filtered: Vec<_> = all_tasks let filtered: Vec<_> = all_tasks
.iter() .iter()
.filter(|t| { .filter(|t| matched_tasks.contains(t.id.as_str()))
matches
.iter()
.any(|(id, typ, _, _, _)| id == &t.id && typ == "task")
})
.collect(); .collect();
serde_json::to_value(&filtered).unwrap_or_default() serde_json::to_value(&filtered).unwrap_or_default()
}); });
@@ -376,11 +390,7 @@ impl McpTool for OmniSearchHandler {
let snippets_json = state.snippets.read_with(|all_snippets| { let snippets_json = state.snippets.read_with(|all_snippets| {
let filtered: Vec<_> = all_snippets let filtered: Vec<_> = all_snippets
.iter() .iter()
.filter(|s| { .filter(|s| matched_snippets.contains(s.name.as_str()))
matches
.iter()
.any(|(id, typ, _, _, _)| id == &s.name && typ == "snippet")
})
.collect(); .collect();
serde_json::to_value(&filtered).unwrap_or_default() serde_json::to_value(&filtered).unwrap_or_default()
}); });
@@ -388,11 +398,7 @@ impl McpTool for OmniSearchHandler {
let adrs_json = state.adrs.read_with(|all_adrs| { let adrs_json = state.adrs.read_with(|all_adrs| {
let filtered: Vec<_> = all_adrs let filtered: Vec<_> = all_adrs
.iter() .iter()
.filter(|a| { .filter(|a| matched_adrs.contains(a.id.as_str()))
matches
.iter()
.any(|(id, typ, _, _, _)| id == &a.id && typ == "adr")
})
.collect(); .collect();
serde_json::to_value(&filtered).unwrap_or_default() serde_json::to_value(&filtered).unwrap_or_default()
}); });
server/src/main.rs
+16 -10
View File
@@ -115,16 +115,22 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
let app_state_clone = Arc::clone(&app_state); let app_state_clone = Arc::clone(&app_state);
let mut rx = state.activity_tx.subscribe(); let mut rx = state.activity_tx.subscribe();
tokio::spawn(async move { tokio::spawn(async move {
while let Ok(msg) = rx.recv().await { loop {
let senders: Vec<_> = app_state_clone match rx.recv().await {
.clients Ok(msg) => {
.read() let senders: Vec<_> = app_state_clone
.unwrap_or_else(|e| e.into_inner()) .clients
.values() .read()
.cloned() .unwrap_or_else(|e| e.into_inner())
.collect(); .values()
for client_tx in senders { .cloned()
let _ = client_tx.try_send(msg.clone()); .collect();
for client_tx in senders {
let _ = client_tx.try_send(msg.clone());
}
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
} }
} }
}); });
server/src/state.rs
+21 -20
View File
@@ -124,26 +124,27 @@ impl MemoryState {
let idx = new_idx.clone(); let idx = new_idx.clone();
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
state.graph.read_with(|g| { let entities: Vec<_> = state
for e in g.entities.values() { .graph
idx.add_entity_sync(e); .read_with(|g| g.entities.values().cloned().collect());
} for e in entities {
}); idx.add_entity_sync(&e);
state.tasks.read_with(|tasks| { }
for t in tasks {
idx.add_task_sync(t); let tasks = state.tasks.read_with(|t| t.clone());
} for t in tasks {
}); idx.add_task_sync(&t);
state.snippets.read_with(|snippets| { }
for s in snippets {
idx.add_snippet_sync(s); let snippets = state.snippets.read_with(|s| s.clone());
} for s in snippets {
}); idx.add_snippet_sync(&s);
state.adrs.read_with(|adrs| { }
for a in adrs {
idx.add_adr_sync(a); let adrs = state.adrs.read_with(|a| a.clone());
} for a in adrs {
}); idx.add_adr_sync(&a);
}
}) })
.await .await
.unwrap_or_else(|e| { .unwrap_or_else(|e| {
server/src/store.rs
+11 -9
View File
@@ -28,18 +28,20 @@ impl<T: DeserializeOwned + Default + Serialize + Clone + Send + Sync + 'static>
let key_inner = key_clone.clone(); let key_inner = key_clone.clone();
let json_data = { let json_data = {
let lock = cache_clone.read().unwrap_or_else(|e| e.into_inner()); let lock = cache_clone.read().unwrap_or_else(|e| e.into_inner());
serde_json::to_vec(&*lock).unwrap_or_default() serde_json::to_vec(&*lock).ok()
}; };
let _ = tokio::task::spawn_blocking(move || { if let Some(json_data) = json_data {
if let Ok(write_txn) = db_inner.begin_write() { let _ = tokio::task::spawn_blocking(move || {
if let Ok(mut table) = write_txn.open_table(STORE_TABLE) { if let Ok(write_txn) = db_inner.begin_write() {
let _ = table.insert(key_inner.as_str(), json_data.as_slice()); if let Ok(mut table) = write_txn.open_table(STORE_TABLE) {
let _ = table.insert(key_inner.as_str(), json_data.as_slice());
}
let _ = write_txn.commit();
} }
let _ = write_txn.commit(); })
} .await;
}) }
.await;
} }
}); });