From 24f2b502be2debaa9e21f7fee65fc4b557f4e50e Mon Sep 17 00:00:00 2001 From: mysma-9403 <64923976+mysma-9403@users.noreply.github.com> Date: Sat, 30 May 2026 18:13:06 +0200 Subject: [PATCH] perf(memory_tree): batch tree fetch in flush_stale_buffers (#2976) Co-authored-by: Steven Enamakel --- src/openhuman/memory_store/trees/store.rs | 66 +++++++++ .../memory_store/trees/store_tests.rs | 45 ++++++ src/openhuman/memory_tree/tree/flush.rs | 135 ++++++++++++++++-- 3 files changed, 235 insertions(+), 11 deletions(-) diff --git a/src/openhuman/memory_store/trees/store.rs b/src/openhuman/memory_store/trees/store.rs index b74041a10..1d4ad71a1 100644 --- a/src/openhuman/memory_store/trees/store.rs +++ b/src/openhuman/memory_store/trees/store.rs @@ -107,6 +107,72 @@ pub fn get_tree(config: &Config, id: &str) -> Result> { }) } +/// Defensive upper bound on the number of `?` placeholders per batched +/// `SELECT … WHERE id IN (?,?,…)` query. SQLite's compile-time +/// `SQLITE_MAX_VARIABLE_NUMBER` has been ≥ 32 766 since 3.32 — 500 leaves +/// a ~65× safety margin. The current call-site +/// (`memory_tree::tree::flush::flush_stale_buffers`) passes one tree_id +/// per stale L0 buffer, so a typical N (tens to low hundreds across all +/// connected sources) runs the loop exactly once. The window exists so +/// future callers with larger id slices do not blow up against a host +/// with a lower compile-time SQLite cap. No volume reduction: all input +/// ids in → all matching rows out; the merged `HashMap` is byte- +/// identical to one giant query. +const TREES_MAX_FETCH_BATCH: usize = 500; + +/// Fetch many trees by id in a single SQL round-trip per +/// [`TREES_MAX_FETCH_BATCH`] window. Replaces the per-id `get_tree` +/// loop inside paths like the cron-driven `flush_stale_buffers`, where +/// the previous code did one `SELECT … WHERE id = ?` per stale buffer. +/// Missing ids are silently absent from the map so callers can preserve +/// the existing "row missing → warn-and-skip" contract without an extra +/// `Ok(None)` sentinel per id. +pub fn get_trees_batch(config: &Config, tree_ids: &[String]) -> Result> { + if tree_ids.is_empty() { + return Ok(HashMap::new()); + } + log::debug!( + "[tree::store] get_trees_batch: ids={} max_batch={TREES_MAX_FETCH_BATCH}", + tree_ids.len() + ); + with_connection(config, |conn| { + let mut out: HashMap = HashMap::with_capacity(tree_ids.len()); + for window in tree_ids.chunks(TREES_MAX_FETCH_BATCH) { + // SAFETY (SQL injection): only the *count* of placeholders is + // interpolated into the query string — `?1,?2,…` are numbered + // bind slots, never the id values. Every id is bound via typed + // `rusqlite::ToSql` params below; nothing user-controlled is + // ever formatted into `sql`. Do NOT inline id values here. + let placeholders = (1..=window.len()) + .map(|i| format!("?{i}")) + .collect::>() + .join(","); + let sql = format!( + "SELECT id, kind, scope, root_id, max_level, status, + created_at_ms, last_sealed_at_ms + FROM mem_tree_trees + WHERE id IN ({placeholders})" + ); + let mut stmt = conn.prepare(&sql)?; + let params: Vec<&dyn rusqlite::ToSql> = + window.iter().map(|s| s as &dyn rusqlite::ToSql).collect(); + let rows = stmt + .query_map(params.as_slice(), row_to_tree)? + .collect::>>() + .context("Failed to collect trees batch")?; + for t in rows { + out.insert(t.id.clone(), t); + } + } + log::debug!( + "[tree::store] get_trees_batch: requested={} found={}", + tree_ids.len(), + out.len() + ); + Ok(out) + }) +} + /// List every tree of a given kind. Used by the global digest to enumerate /// source trees, and by diagnostics. Rows come back ordered by `created_at_ms` /// ASC so callers see a stable iteration order. diff --git a/src/openhuman/memory_store/trees/store_tests.rs b/src/openhuman/memory_store/trees/store_tests.rs index 2e700b7ad..296f6a9b2 100644 --- a/src/openhuman/memory_store/trees/store_tests.rs +++ b/src/openhuman/memory_store/trees/store_tests.rs @@ -288,6 +288,51 @@ fn list_stale_buffers_orders_by_age() { assert_eq!(only_oldest[0].tree_id, "tree-1"); } +// ── get_trees_batch ──────────────────────────────────────────────────── +// +// Same shape as `chunks::store::get_chunks_batch` / +// `score::store::get_scores_batch`: present ids decode through the same +// `row_to_tree` path as the per-id `get_tree` and land in a `HashMap` +// keyed by id; missing ids are silently absent so the +// `flush_stale_buffers` orphan-buffer warn-and-skip path keeps working +// without an extra Ok(None) sentinel per id. + +#[test] +fn get_trees_batch_returns_present_ids_in_map() { + let (_tmp, cfg) = test_config(); + let a = sample_tree("tree-a", "slack:#eng"); + let b = sample_tree("tree-b", "slack:#design"); + insert_tree(&cfg, &a).unwrap(); + insert_tree(&cfg, &b).unwrap(); + + let ids = vec!["tree-a".to_string(), "tree-b".to_string()]; + let map = get_trees_batch(&cfg, &ids).unwrap(); + assert_eq!(map.len(), 2); + // Each decoded row must match the per-id `get_tree` path bit-for-bit + // — same `row_to_tree` decoder under the hood, so the structs are + // equal including the parsed `kind` / `status` enums. + assert_eq!(map.get("tree-a").unwrap(), &a); + assert_eq!(map.get("tree-b").unwrap(), &b); +} + +#[test] +fn get_trees_batch_empty_input_and_missing_ids() { + // Empty input: empty map (no SQL issued). + let (_tmp, cfg) = test_config(); + let empty = get_trees_batch(&cfg, &[]).unwrap(); + assert!(empty.is_empty()); + + // Missing ids: silently absent so `flush_stale_buffers` can warn + // + skip without an extra `Ok(None)` sentinel per id. + let a = sample_tree("tree-a", "slack:#eng"); + insert_tree(&cfg, &a).unwrap(); + let ids = vec!["tree-a".to_string(), "ghost:no-such".to_string()]; + let map = get_trees_batch(&cfg, &ids).unwrap(); + assert_eq!(map.len(), 1); + assert_eq!(map.get("tree-a").unwrap(), &a); + assert!(map.get("ghost:no-such").is_none()); +} + // ── get_summaries_batch ──────────────────────────────────────────────── // // Same shape as `chunks::store::get_chunks_batch` / diff --git a/src/openhuman/memory_tree/tree/flush.rs b/src/openhuman/memory_tree/tree/flush.rs index 8a1fde6e7..f93a32462 100644 --- a/src/openhuman/memory_tree/tree/flush.rs +++ b/src/openhuman/memory_tree/tree/flush.rs @@ -27,6 +27,8 @@ pub async fn flush_stale_buffers( max_age: Duration, strategy: &LabelStrategy, ) -> Result { + use crate::openhuman::memory_store::trees::store::get_trees_batch; + let now = Utc::now(); let cutoff = now - max_age; let stale = store::list_stale_buffers(config, cutoff)?; @@ -36,20 +38,38 @@ pub async fn flush_stale_buffers( max_age ); + // One batched `SELECT … WHERE id IN (?,?,…)` over the distinct + // tree_ids instead of N per-buffer `get_tree` round-trips. `stale` + // is filtered to L0 by `list_stale_buffers` so the same tree_id + // typically appears at most once, but we still de-dup defensively + // before hitting the DB — future widenings of the filter (e.g. + // dropping the `level = 0` clause) would otherwise re-pay the cost. + // Missing rows stay silently absent from the map; the + // `Some(t) => ... / None => warn-and-skip` orphan-buffer path below + // preserves the per-id `get_tree` `Ok(None)` semantics. + let distinct_tree_ids: Vec = { + let mut seen = std::collections::HashSet::new(); + let mut out = Vec::new(); + for buf in &stale { + if seen.insert(buf.tree_id.clone()) { + out.push(buf.tree_id.clone()); + } + } + out + }; + let tree_by_id = get_trees_batch(config, &distinct_tree_ids)?; + let mut seals: usize = 0; for buf in stale { - let tree = match store::get_tree(config, &buf.tree_id)? { - Some(t) => t, - None => { - log::warn!( - "[tree::flush] orphan buffer tree_id={} level={}", - buf.tree_id, - buf.level - ); - continue; - } + let Some(tree) = tree_by_id.get(&buf.tree_id) else { + log::warn!( + "[tree::flush] orphan buffer tree_id={} level={}", + buf.tree_id, + buf.level + ); + continue; }; - let sealed = cascade_all_from(config, &tree, buf.level, Some(now), strategy).await?; + let sealed = cascade_all_from(config, tree, buf.level, Some(now), strategy).await?; seals += sealed.len(); } Ok(seals) @@ -274,4 +294,97 @@ mod tests { assert_eq!(seals, 0); assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0); } + + /// Regression for the `get_trees_batch` refactor: two distinct + /// trees each carry a stale L0 buffer; `flush_stale_buffers` must + /// seal both via a single batched tree-fetch. Pre-refactor this + /// path issued N per-buffer `get_tree(...)` round-trips; post- + /// refactor it's one `SELECT … WHERE id IN (?,?)` and a per-id + /// HashMap lookup. The interesting bit is the HashMap lookup + /// keying by tree_id (not by `enumerate()` position over the input + /// slice) — with two trees, swapping the lookup key would leak one + /// tree's `Tree` into the other tree's seal and either error out + /// or write summaries against the wrong `root_id`. + #[tokio::test] + async fn flush_seals_multiple_distinct_trees_via_batched_lookup() { + let (_tmp, cfg) = test_config(); + let tree_a = get_or_create_source_tree(&cfg, "slack:#eng").unwrap(); + let tree_b = get_or_create_source_tree(&cfg, "slack:#design").unwrap(); + assert_ne!( + tree_a.id, tree_b.id, + "test prerequisite: distinct source_ids must yield distinct tree_ids" + ); + let provider: Arc = + Arc::new(StaticChatProvider::new("test summary content")); + + let old_ts = Utc::now() - Duration::days(10); + + // Plant a stale L0 leaf in each tree, backed by real chunks so + // the cascade can seal end-to-end. + for (src_id, content) in [ + ("slack:#eng", "engineering content"), + ("slack:#design", "design content"), + ] { + let c = Chunk { + id: chunk_id(SourceKind::Chat, src_id, 0, content), + content: content.into(), + metadata: Metadata { + source_kind: SourceKind::Chat, + source_id: src_id.into(), + owner: "alice".into(), + timestamp: old_ts, + time_range: (old_ts, old_ts), + tags: vec![], + source_ref: Some(SourceRef::new(&format!("slack://{src_id}"))), + }, + token_count: 100, + seq_in_source: 0, + created_at: old_ts, + partial_message: false, + }; + upsert_chunks(&cfg, &[c.clone()]).unwrap(); + stage_test_chunks(&cfg, &[c.clone()]); + let leaf = LeafRef { + chunk_id: c.id.clone(), + token_count: 100, + timestamp: old_ts, + content: c.content.clone(), + entities: vec![], + topics: vec![], + score: 0.5, + }; + let target_tree = if src_id == "slack:#eng" { + &tree_a + } else { + &tree_b + }; + test_override::with_provider(Arc::clone(&provider), async { + append_leaf(&cfg, target_tree, &leaf, &LabelStrategy::Empty) + .await + .unwrap(); + }) + .await; + } + + let seals = test_override::with_provider(provider, async { + flush_stale_buffers(&cfg, Duration::days(7), &LabelStrategy::Empty) + .await + .unwrap() + }) + .await; + + // Both trees sealed via a single batched fetch — the HashMap + // lookup correctly routed each buffer to its own `Tree`. + assert_eq!(seals, 2, "both stale trees must seal in one flush pass"); + assert_eq!( + store::count_summaries(&cfg, &tree_a.id).unwrap(), + 1, + "tree_a must own its summary (HashMap keyed by id, not position)" + ); + assert_eq!( + store::count_summaries(&cfg, &tree_b.id).unwrap(), + 1, + "tree_b must own its summary (HashMap keyed by id, not position)" + ); + } }