From 2e254262e9090634b0c3fdfae10267959ceb05cd Mon Sep 17 00:00:00 2001 From: mysma-9403 <64923976+mysma-9403@users.noreply.github.com> Date: Mon, 22 Jun 2026 15:01:54 +0200 Subject: [PATCH] perf(memory_tree): concurrently seal document-subtree sibling batches (#3416) Co-authored-by: M3gA-Mind --- src/openhuman/memory_tree/tree/bucket_seal.rs | 78 ++++++++++-- .../memory_tree/tree/bucket_seal_tests.rs | 115 ++++++++++++++++++ 2 files changed, 180 insertions(+), 13 deletions(-) diff --git a/src/openhuman/memory_tree/tree/bucket_seal.rs b/src/openhuman/memory_tree/tree/bucket_seal.rs index dfc717f65..b78505fcd 100644 --- a/src/openhuman/memory_tree/tree/bucket_seal.rs +++ b/src/openhuman/memory_tree/tree/bucket_seal.rs @@ -38,6 +38,7 @@ use std::sync::Arc; use anyhow::{Context, Result}; use chrono::{DateTime, Utc}; +use futures::stream::{StreamExt, TryStreamExt}; use rusqlite::Transaction; use crate::openhuman::config::Config; @@ -1012,6 +1013,17 @@ pub const MERGE_LEVEL_BASE: u32 = 1_000; /// children. Independent of [`SUMMARY_FANOUT`] (which gates the merge tier). const DOC_SUBTREE_MAX_FANIN: usize = 32; +/// Bound on how many sibling batches within a single document-subtree level +/// are sealed concurrently. Each [`seal_explicit_children`] call is one +/// independent summarise + embed round-trip pair, so overlapping them with a +/// small pool collapses the serial `N × (LLM + embed RTT)` wall-time toward +/// `N/CONCURRENCY × RTT` on the ingest path. Kept low: the SQLite write tx +/// serialises on a single writer anyway, and networked summarise/embed +/// providers (Voyage, cloud LLMs) rate-limit — so a large pool buys nothing +/// and risks 429s. `buffered` (ordered) preserves the serial cascade's batch +/// order, so the resulting tree shape is identical. +const DOC_SUBTREE_SEAL_CONCURRENCY: usize = 8; + /// Build (or re-build, for a new version) one document's subtree and merge /// its doc-root into the connection tree. /// @@ -1067,21 +1079,61 @@ pub async fn seal_document_subtree( batch_by_count(¤t_ids, DOC_SUBTREE_MAX_FANIN) }; - let mut next_ids: Vec = Vec::with_capacity(batches.len()); - for batch in &batches { - let node = seal_explicit_children( - config, - tree, + // Within one level the batches are disjoint child sets, and + // `seal_explicit_children` touches no shared `(tree, level)` buffer + // and never advances the tree root (see its docstring) — so the + // sibling seals are independent and their summarise + embed + // round-trips can overlap. Run them with bounded concurrency. + // + // `buffered` (ordered), NOT `buffer_unordered`: the next level groups + // `next_ids` positionally via `batch_by_count`, so the output order + // must match the serial cascade to keep the tree shape identical. + // `try_collect` short-circuits on the first error, preserving the + // old `?`-on-first-failure behaviour. Levels stay sequential — level + // N+1 consumes level N's output — so only siblings overlap, not the + // cross-level fan-in. The blocking SQLite write inside each seal + // still serialises (single writer), which is fine: the win is the + // overlapped network RTTs, not the DB writes. + // Build the per-batch futures eagerly into a `Vec` (they don't run until + // polled) rather than lazily inside a `stream::iter(...).map(closure)`: + // the lazy form makes rustc try to infer a higher-ranked `FnOnce` over + // the borrowed `batch` lifetime, which fails ("implementation of + // `FnOnce` is not general enough") once this async fn is spawned. Eager + // collection ties each future to this stack frame's concrete lifetime. + if batches.len() > 1 { + log::debug!( + "[tree::bucket_seal] doc-subtree concurrent seal tree_id={} doc_id_hash={} level={} batches={} concurrency={}", + tree.id, + redact(doc_id), current_level, - batch, - Some(doc_id), - version_ms, - strategy, - ) - .await?; - next_ids.push(node.id.clone()); - doc_root = Some(node); + batches.len(), + DOC_SUBTREE_SEAL_CONCURRENCY + ); } + let batch_futures: Vec<_> = batches + .iter() + .map(|batch| { + seal_explicit_children( + config, + tree, + current_level, + batch, + Some(doc_id), + version_ms, + strategy, + ) + }) + .collect(); + let nodes: Vec = futures::stream::iter(batch_futures) + .buffered(DOC_SUBTREE_SEAL_CONCURRENCY) + .try_collect() + .await?; + + let next_ids: Vec = nodes.iter().map(|n| n.id.clone()).collect(); + // The terminal level produces exactly one node (the loop breaks once + // `current_ids.len() <= 1`); preserving batch order means the last + // node is the doc-root, identical to the serial assignment. + doc_root = nodes.into_iter().next_back(); current_level += 1; current_ids = next_ids; diff --git a/src/openhuman/memory_tree/tree/bucket_seal_tests.rs b/src/openhuman/memory_tree/tree/bucket_seal_tests.rs index 01b7ea264..bed658d32 100644 --- a/src/openhuman/memory_tree/tree/bucket_seal_tests.rs +++ b/src/openhuman/memory_tree/tree/bucket_seal_tests.rs @@ -351,6 +351,121 @@ async fn seal_document_subtree_single_chunk_is_verbatim_passthrough_no_llm() { ); } +/// Multi-batch document subtree: a doc whose chunks span several level-0 +/// batches must seal its siblings **concurrently** (overlapping the +/// summarise + embed round-trips) while still producing an identical +/// doc-root. A `ChatProvider` probe records the peak number of in-flight +/// summarise calls; with four independent level-0 batches the peak must +/// exceed 1 (proving overlap) yet stay within `DOC_SUBTREE_SEAL_CONCURRENCY` +/// (proving the bound is respected). This is the regression guard for the +/// serial → bounded-concurrency change in `seal_document_subtree`. +#[tokio::test] +async fn seal_document_subtree_seals_sibling_batches_concurrently() { + use crate::openhuman::memory::chat::ChatPrompt; + use crate::openhuman::memory_store::chunks::store::upsert_chunks; + use crate::openhuman::memory_store::chunks::types::{ + chunk_id, Chunk, Metadata, SourceKind, SourceRef, + }; + use std::sync::atomic::{AtomicUsize, Ordering}; + + /// Records concurrent in-flight `chat_for_json` calls and the peak; a + /// short sleep forces would-be-concurrent calls to actually overlap so + /// the peak reflects real scheduling, not luck. + struct ConcurrencyProbe { + in_flight: AtomicUsize, + peak: AtomicUsize, + calls: AtomicUsize, + } + #[async_trait::async_trait] + impl ChatProvider for ConcurrencyProbe { + fn name(&self) -> &str { + "test:concurrency-probe" + } + async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { + let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; + self.peak.fetch_max(now, Ordering::SeqCst); + self.calls.fetch_add(1, Ordering::SeqCst); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + self.in_flight.fetch_sub(1, Ordering::SeqCst); + Ok("doc batch summary".to_string()) + } + } + + let (_tmp, cfg) = test_config(); + let tree = get_or_create_source_tree(&cfg, "notion:conn1").unwrap(); + let doc_id = "notion:conn1:pageM"; + + // Four chunks, each ~60% of the input budget, so the greedy token-budget + // batcher puts each in its OWN level-0 batch → four independent + // summarise calls that should overlap under bounded concurrency. + let per_chunk = INPUT_TOKEN_BUDGET * 6 / 10; + let ts = Utc::now(); + let mut ids = Vec::new(); + let mut chunks = Vec::new(); + for seq in 0..4u32 { + let content = format!("substantive document body number {seq}"); + let c = Chunk { + id: chunk_id(SourceKind::Document, doc_id, seq, &content), + content, + metadata: Metadata { + source_kind: SourceKind::Document, + source_id: doc_id.to_string(), + owner: "notion:conn1".into(), + timestamp: ts, + time_range: (ts, ts), + tags: vec!["notion".into()], + source_ref: Some(SourceRef::new("notion://page/pageM")), + path_scope: Some("notion:conn1".into()), + }, + token_count: per_chunk, + seq_in_source: seq, + created_at: ts, + partial_message: false, + }; + ids.push(c.id.clone()); + chunks.push(c); + } + upsert_chunks(&cfg, &chunks).unwrap(); + stage_test_chunks(&cfg, &chunks); + + let probe = Arc::new(ConcurrencyProbe { + in_flight: AtomicUsize::new(0), + peak: AtomicUsize::new(0), + calls: AtomicUsize::new(0), + }); + let provider: Arc = probe.clone(); + + let doc_root_id = test_override::with_provider(provider, async { + seal_document_subtree(&cfg, &tree, doc_id, Some(7), &ids, &LabelStrategy::Empty) + .await + .unwrap() + }) + .await; + + // Correctness: the four chunks roll up under a single doc-root tagged + // with the document id + version — tree shape unchanged by batching. + let root = store::get_summary(&cfg, &doc_root_id).unwrap().unwrap(); + assert_eq!(root.doc_id.as_deref(), Some(doc_id)); + assert_eq!(root.version_ms, Some(7)); + assert_eq!( + root.child_ids.len(), + 4, + "all four level-0 siblings must fan into the doc-root" + ); + + // Concurrency: at least two level-0 siblings overlapped, and the pool + // bound was respected. + let peak = probe.peak.load(Ordering::SeqCst); + assert!( + peak >= 2, + "expected overlapping sibling seals (peak in-flight = {peak})" + ); + assert!( + peak <= DOC_SUBTREE_SEAL_CONCURRENCY, + "concurrency must stay within the pool bound (peak = {peak})" + ); +} + #[tokio::test] async fn crossing_budget_triggers_seal() { use crate::openhuman::memory_store::chunks::store::upsert_chunks;