perf(graph): eliminate heavy allocations by substituting KnowledgeGraph clones with BorrowedGraph for json serialization

This commit is contained in:
Riz Ashraf committed 2026-09-21 16:33:48 +01:00
1 parent 1d5e4bd00d
commit 8934857635
2 files changed
+38 -18

No files matched your search

+32 -18
View File
@@ -7,6 +7,21 @@ use serde_json::Value;
use std::collections::HashSet; use std::collections::HashSet;
use std::sync::Arc; use std::sync::Arc;
#[derive(serde::Serialize)]
struct BorrowedGraph<'a> {
entities: std::collections::HashMap<&'a String, &'a crate::models::Entity>,
relations: Vec<&'a crate::models::Relation>,
}
impl<'a> Default for BorrowedGraph<'a> {
fn default() -> Self {
Self {
entities: std::collections::HashMap::new(),
relations: Vec::new(),
}
}
}
pub struct QueryGraphPathHandler; pub struct QueryGraphPathHandler;
#[async_trait] #[async_trait]
@@ -267,15 +282,15 @@ impl McpTool for ReadGraphHandler {
let req: ReadGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let req: ReadGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let data = state.read_graph(|full| { let data = state.read_graph(|full| {
if let Some(ns) = req.namespace { if let Some(ns) = req.namespace {
let mut filtered = KnowledgeGraph::default(); let mut filtered = BorrowedGraph::default();
for (k, v) in &full.entities { for (k, v) in &full.entities {
if v.namespace == ns { if v.namespace == ns {
filtered.entities.insert(k.clone(), v.clone()); filtered.entities.insert(k, v);
} }
} }
for r in &full.relations { for r in &full.relations {
if r.namespace == ns { if r.namespace == ns {
filtered.relations.push(r.clone()); filtered.relations.push(r);
} }
} }
serde_json::to_string(&filtered).unwrap_or_default() serde_json::to_string(&filtered).unwrap_or_default()
@@ -308,18 +323,18 @@ impl McpTool for SearchNodesHandler {
vec![] vec![]
}; };
let mut result = KnowledgeGraph::default(); let data = state.read_graph(|full| {
state.read_graph(|full| { let mut result = BorrowedGraph::default();
for (id, doc_type, _, _, _) in matches { for (id, doc_type, _, _, _) in &matches {
if doc_type == "entity" if doc_type == "entity"
&& let Some(e) = full.entities.get(&id) && let Some(e) = full.entities.get(id)
{ {
result.entities.insert(id, e.clone()); result.entities.insert(id, e);
} }
} }
serde_json::to_string(&result).unwrap_or_default()
}); });
let data = serde_json::to_string(&result).unwrap_or_default(); Ok(data)
Ok(data.to_string())
} }
} }
@@ -337,28 +352,27 @@ impl McpTool for OpenNodesHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
let req: OpenNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let req: OpenNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let result = state.read_graph(|full| { let data = state.read_graph(|full| {
let targets: HashSet<&str> = req.names.iter().map(|s| s.as_str()).collect(); let targets: HashSet<&str> = req.names.iter().map(|s| s.as_str()).collect();
let mut result = KnowledgeGraph::default(); let mut result = BorrowedGraph::default();
let mut connected: HashSet<&str> = HashSet::new(); let mut connected: HashSet<&str> = HashSet::new();
for r in &full.relations { for r in &full.relations {
if targets.contains(r.from.as_str()) { if targets.contains(r.from.as_str()) {
connected.insert(r.to.as_str()); connected.insert(r.to.as_str());
result.relations.push(r.clone()); result.relations.push(r);
} else if targets.contains(r.to.as_str()) { } else if targets.contains(r.to.as_str()) {
connected.insert(r.from.as_str()); connected.insert(r.from.as_str());
result.relations.push(r.clone()); result.relations.push(r);
} }
} }
for (name, e) in &full.entities { for (name, e) in &full.entities {
if targets.contains(name.as_str()) || connected.contains(name.as_str()) { if targets.contains(name.as_str()) || connected.contains(name.as_str()) {
result.entities.insert(name.clone(), e.clone()); result.entities.insert(name, e);
} }
} }
result serde_json::to_string(&result).unwrap_or_default()
}); });
let data = serde_json::to_string(&result).unwrap_or_default(); Ok(data)
Ok(data.to_string())
} }
} }
+6
View File
@@ -0,0 +1,6 @@
local buf = vim.fn.bufnr('server/src/handlers_v2/graph.rs')
if buf == -1 then
vim.cmd('e server/src/handlers_v2/graph.rs')
buf = vim.api.nvim_get_current_buf()
end
return vim.api.nvim_buf_get_name(buf)