[mysma-9403]

Summary

memory_tree::tree::flush::flush_stale_buffers previously called store::get_tree(config, &buf.tree_id)? once per stale buffer returned by list_stale_buffers — i.e. one SQLite round-trip per stale row before the cascade even started. On hosts with many stale L0 buffers (e.g. after a long backfill, a cron miss, or a process restart during ingest) every flush pass re-paid that cost.

Replace the per-buffer loop with one upfront batched read:

• memory_store::trees::store::get_trees_batch(config, &[ids]) -> Result<HashMap<String, Tree>>

flush_stale_buffers now dedups buffer tree_ids into a HashSet, issues one SELECT … WHERE id IN (?,?,…) (so 1 round-trip), then per-id HashMap lookup routes each buffer to its own Tree for cascade_all_from. The level = 0 filter in list_stale_buffers keeps the same tree_id at most once today, but the dedup is defensive against future widenings of that filter that would otherwise re-pay the cost.

Defensive chunking — TREES_MAX_FETCH_BATCH = 500

get_trees_batch iterates via chunks(TREES_MAX_FETCH_BATCH).

• SQLite's compile-time SQLITE_MAX_VARIABLE_NUMBER has been ≥ 32 766 since 3.32 → 500 leaves a ~65× safety margin.
• The flush_stale_buffers call-site passes the deduped distinct tree_id set (typically << 500 in practice) → runs as one query today.
• The loop is purely defensive: future callers passing larger id lists will 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.

Semantic preservation — no behavioural change

• get_trees_batch reuses row_to_tree so decoded Tree values stay bit-identical to the per-row path.
• Missing ids stay silently absent from the HashMap, mirroring the per-row get_tree Ok(None) contract — so the existing orphan-buffer warn-and-skip path keeps working unchanged.
• Per-buffer iteration order preserved by walking stale directly; HashMap lookup keys by &buf.tree_id (not by enumerate() position).
• cascade_all_from failure still bubbles up with the same error semantics as before.

Test plan

New unit tests (3 total)

trees store (src/openhuman/memory_store/trees/store_tests.rs):

• get_trees_batch_returns_present_ids_in_map — both ids present in the map, decoded Tree equal to the per-row path bit-for-bit (same row_to_tree decoder).
• get_trees_batch_empty_input_and_missing_ids — empty input → empty map (no SQL); missing id silently absent so callers keep their Ok(None) skip path.

flush (src/openhuman/memory_tree/tree/flush.rs):

• flush_seals_multiple_distinct_trees_via_batched_lookup — two real trees each carry a stale L0 buffer; one flush_stale_buffers pass seals both via the batched fetch. Pins the HashMap-keyed-by-id contract: swapping the lookup key to enumerate() position would route buffers to the wrong Tree and mis-write summaries against the wrong root_id — the test asserts count_summaries(tree_a) == 1 && count_summaries(tree_b) == 1.

Existing surface — verified no regression

• All existing tests in memory_store::trees:: and memory_tree::tree::flush:: continue to pass (25 + 4 = 29 tests, 0 failures).
• cargo fmt --check — clean.

Companion to #2927 / #2941 / #2948

Same playbook as the topic-rerank batch (#2927), fetch_leaves batch (#2941), and hydrate_summary_inputs batch (#2948): partition / batch / decorate. Different table, different call-site, same pattern. get_trees_batch is a building block several remaining N+1 hot-paths in the graph layer will reuse.

Summary by CodeRabbit

Release Notes

• Performance

  • Optimized memory tree operations by consolidating multiple database queries into single batch requests, reducing round-trips and improving efficiency when processing multiple trees.

• Tests

  • Added regression test coverage for bulk tree flushing operations to ensure correct behavior across multiple trees.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces a batched tree-fetch optimization: a new get_trees_batch API fetches multiple trees in one SQL round-trip with placeholder-aware chunking, and flush_stale_buffers is refactored to use it instead of per-buffer individual queries.

Changes

Batch tree-fetch and flush optimization

Layer / File(s)	Summary
Batch tree-fetch API src/openhuman/memory_store/trees/store.rs, src/openhuman/memory_store/trees/store_tests.rs	get_trees_batch fetches multiple trees in batched IN-clause queries, chunking ids via TREES_MAX_FETCH_BATCH to avoid SQLite placeholder overflow, and returns a HashMap keyed by tree id. Tests cover nominal retrieval and edge cases (empty input, missing ids).
Flush stale buffers optimization src/openhuman/memory_tree/tree/flush.rs	flush_stale_buffers now collects distinct tree ids from stale buffers, fetches all trees in one batch, and looks up each buffer's tree from the map, replacing per-buffer individual queries. Regression test verifies correct tree lookup across multiple distinct trees.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Suggested labels

memory, working

Suggested reviewers

• M3gA-Mind

Poem

🐰 Behold! Batched queries hop with glee,
One round-trip fetches trees, not three per tree.
Stale buffers now seal swift and right—
Each tree_id finds its match in flight! 🌳✨

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title directly reflects the main change: introducing batched tree fetching in flush_stale_buffers as a performance optimization.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

src/openhuman/memory_store/trees/store.rs (1)

130-160: ⚡ Quick win

Add debug logging to the new batched fetch flow.

Logic is correct: numbered placeholders bind positionally, missing ids stay absent from the map (preserving the prior Ok(None) skip semantics), and the chunking keeps results identical to a single query. However, this new external-call path emits no diagnostics. A debug/trace line at entry/exit (counts and chunk windows) would aid grep-based debugging of the cron flush path.

♻️ Suggested logging

     with_connection(config, |conn| {
         let mut out: HashMap<String, Tree> = HashMap::with_capacity(tree_ids.len());
         for window in tree_ids.chunks(TREES_MAX_FETCH_BATCH) {
+            log::trace!(
+                "[tree::store] get_trees_batch window_len={} total_ids={}",
+                window.len(),
+                tree_ids.len()
+            );
             let placeholders = (1..=window.len())

As per coding guidelines: "add substantial, development-oriented logs on new/changed flows at entry/exit points, branch decisions, external calls ... use structured, grep-friendly context (stable prefixes like [domain]...)".

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/openhuman/memory_store/trees/store.rs` around lines 130 - 160, Add
development-oriented debug/trace logging to the batched fetch flow in
get_trees_batch: log at entry the incoming tree_ids count and relevant
config/Context, log each chunk window (index, size, and the placeholder range or
first/last id) before executing the SELECT, and log exit with total found rows
and out.len(); emit logs using the project's structured logger (e.g.,
tracing::debug! or log::debug!) so they include a stable tag like "[mem-tree]"
and reference TREES_MAX_FETCH_BATCH, mem_tree_trees, and row_to_tree for
grepability. Ensure logs are inexpensive (only format when enabled) and do not
change behavior or return values from with_connection/get_trees_batch.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@src/openhuman/memory_store/trees/store.rs`:
- Around line 130-160: Add development-oriented debug/trace logging to the
batched fetch flow in get_trees_batch: log at entry the incoming tree_ids count
and relevant config/Context, log each chunk window (index, size, and the
placeholder range or first/last id) before executing the SELECT, and log exit
with total found rows and out.len(); emit logs using the project's structured
logger (e.g., tracing::debug! or log::debug!) so they include a stable tag like
"[mem-tree]" and reference TREES_MAX_FETCH_BATCH, mem_tree_trees, and
row_to_tree for grepability. Ensure logs are inexpensive (only format when
enabled) and do not change behavior or return values from
with_connection/get_trees_batch.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: a2ef0810-50b7-440a-a725-46946a92dd65

📥 Commits

Reviewing files that changed from the base of the PR and between f895013 and 9a3f1d2.

📒 Files selected for processing (3)

• src/openhuman/memory_store/trees/store.rs
• src/openhuman/memory_store/trees/store_tests.rs
• src/openhuman/memory_tree/tree/flush.rs

[oxoxDev]

LGTM — clean, behavior-preserving perf fix. get_trees_batch replaces the per-id get_tree loop in flush_stale_buffers with one batched SELECT … WHERE id IN (…) per 500-id window, mirroring the existing get_chunks_batch/get_scores_batch pattern. De-dups tree_ids defensively, and missing rows stay absent from the map so the orphan-buffer warn-and-skip contract is preserved bit-for-bit. Tests cover present/empty/missing. Approving.

[staimoorulhassan]

CodeRabbit-style Review — PR #2976

Walkthrough

Replaces N per-buffer get_tree() SQLite calls in flush_stale_buffers with one batched SELECT … WHERE id IN (…) via new get_trees_batch(). Adds 3 targeted unit tests. Companion to the batch-fetch series (#2927, #2941, #2948).

---

Findings

🟡 MINOR — Dynamic SQL construction via format!() — add a safety comment

get_trees_batch builds the WHERE id IN ({placeholders}) clause with format!(). The values are correctly bound via typed rusqlite::ToSql params (no interpolation), but a future maintainer who doesn't notice this distinction could be misled into thinking the code is SQL-injectable. A one-line comment is all that's needed:

// placeholders are positional bind params (?1,?2,…) — values are never interpolated
let sql = format!(
    "SELECT … WHERE id IN ({placeholders})"
);

---

🟡 MINOR — Prefer params_from_iter over Vec<&dyn rusqlite::ToSql>

The current approach:

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)?;

is valid but fragile (raw trait-object lifetime). The idiomatic rusqlite pattern for dynamic params is:

let rows = stmt.query_map(
    rusqlite::params_from_iter(window.iter().map(String::as_str)),
    row_to_tree,
)?;

params_from_iter was introduced specifically for this use-case and avoids the manual &dyn dance.

---

🟡 MINOR — Document that the dedup is for DB efficiency, not semantic ordering

let distinct_tree_ids: Vec<String> = {
    let mut seen = std::collections::HashSet::new();
    // …
};

The insertion-order-preserving dedup is correct and deterministic (because list_stale_buffers returns rows in ORDER BY order). But a future reader simplifying this to .sort().dedup() would still be correct today — the per-buffer loop walks stale directly for ordering, not distinct_tree_ids. A short comment prevents the "why not just sort+dedup?" question:

// Preserve list_stale_buffers order so the HashMap lookup below
// stays predictable; dedup is purely to avoid redundant DB params.

---

🔵 NITPICK — Consider splitting the large flush regression test

flush_seals_multiple_distinct_trees_via_batched_lookup does two things: verifies that both trees are sealed (happy path) and pins the HashMap-keyed-by-id contract (regression). Splitting into:

• flush_seals_multiple_trees_batched — happy path
• flush_keyed_by_tree_id_not_position — regression assertion

would make the test intent clearer and failure messages more diagnostic.

---

Verdict

No blockers. Four minor/nitpick comments. The algorithm, chunking strategy, semantic-preservation argument, and test coverage are all excellent — this is a well-engineered perf improvement. Approve once comments are considered.