diff --git a/src/openhuman/memory/ingest_pipeline.rs b/src/openhuman/memory/ingest_pipeline.rs index 614bd63d1b..df3a7e68e2 100644 --- a/src/openhuman/memory/ingest_pipeline.rs +++ b/src/openhuman/memory/ingest_pipeline.rs @@ -290,7 +290,7 @@ async fn persist( let source_id_for_store = source_id.to_string(); let raw_refs_for_store = raw_refs_for_chunks.clone(); let written = tokio::task::spawn_blocking(move || -> Result> { - use std::collections::{HashMap, HashSet}; + use std::collections::HashSet; chunk_store::with_connection(&config_owned, |conn| { // IMMEDIATE, not the default DEFERRED: this transaction reads // (get_chunk_lifecycle_status_tx) before it writes @@ -362,11 +362,16 @@ async fn persist( // insert a new row that picks up the column DEFAULT — so reading // post-upsert can't distinguish "brand new" from // "already-admitted-from-prior-ingest". - let mut prior: HashMap> = HashMap::new(); - for s in &staged_for_store { - let status = chunk_store::get_chunk_lifecycle_status_tx(&tx, &s.chunk.id)?; - prior.insert(s.chunk.id.clone(), status); - } + // Batch the prior-lifecycle read: one `SELECT … WHERE id IN (…)` + // instead of one round-trip per staged chunk. Keeping this short + // matters because it runs inside the write transaction, and SQLite + // serialises writers. `prior` maps a chunk id to its pre-upsert + // status; ids with no prior row are simply absent (a batch miss == + // the old per-id `Ok(None)`), so the lookup below treats a miss as + // "genuinely new". + let staged_ids: Vec = + staged_for_store.iter().map(|s| s.chunk.id.clone()).collect(); + let prior = chunk_store::get_chunk_lifecycle_statuses_tx(&tx, &staged_ids)?; let n = chunk_store::upsert_staged_chunks_tx(&tx, &staged_for_store)?; @@ -392,7 +397,7 @@ async fn persist( // leave the lifecycle alone and skip the extract enqueue. let mut to_schedule: HashSet = HashSet::new(); for s in &staged_for_store { - let pre = prior.get(&s.chunk.id).cloned().flatten(); + let pre = prior.get(&s.chunk.id).cloned(); let needs_processing = matches!( pre.as_deref(), None | Some(chunk_store::CHUNK_STATUS_PENDING_EXTRACTION), diff --git a/src/openhuman/memory_store/chunks/store.rs b/src/openhuman/memory_store/chunks/store.rs index 7b068c657c..aed7c565be 100644 --- a/src/openhuman/memory_store/chunks/store.rs +++ b/src/openhuman/memory_store/chunks/store.rs @@ -807,11 +807,64 @@ pub fn get_chunk_lifecycle_status(config: &Config, chunk_id: &str) -> Result, - chunk_id: &str, -) -> Result> { - get_chunk_lifecycle_status_conn(tx, chunk_id) + chunk_ids: &[String], +) -> Result> { + get_chunk_lifecycle_statuses_conn(tx, chunk_ids) +} + +fn get_chunk_lifecycle_statuses_conn( + conn: &Connection, + chunk_ids: &[String], +) -> Result> { + let mut out: HashMap = HashMap::with_capacity(chunk_ids.len()); + if chunk_ids.is_empty() { + return Ok(out); + } + for window in chunk_ids.chunks(MAX_FETCH_BATCH) { + // Positional placeholders `?1, ?2, …, ?n` — rusqlite binds 1..n in the + // order values are passed. Same shape as `get_chunks_batch`. + let placeholders = (1..=window.len()) + .map(|i| format!("?{i}")) + .collect::>() + .join(","); + let sql = format!( + "SELECT id, lifecycle_status FROM mem_tree_chunks WHERE id IN ({placeholders})" + ); + let mut stmt = conn + .prepare(&sql) + .context("prepare get_chunk_lifecycle_statuses")?; + let params: Vec<&dyn rusqlite::ToSql> = + window.iter().map(|id| id as &dyn rusqlite::ToSql).collect(); + let rows = stmt + .query_map(params.as_slice(), |r| { + Ok((r.get::<_, String>(0)?, r.get::<_, Option>(1)?)) + }) + .context("query get_chunk_lifecycle_statuses")?; + for row in rows { + let (id, status) = row.context("decode get_chunk_lifecycle_statuses row")?; + if let Some(status) = status { + out.insert(id, status); + } + } + } + Ok(out) } fn get_chunk_lifecycle_status_conn(conn: &Connection, chunk_id: &str) -> Result> { diff --git a/src/openhuman/memory_store/chunks/store_tests.rs b/src/openhuman/memory_store/chunks/store_tests.rs index c50052d42a..166e9211e6 100644 --- a/src/openhuman/memory_store/chunks/store_tests.rs +++ b/src/openhuman/memory_store/chunks/store_tests.rs @@ -1518,6 +1518,69 @@ fn get_chunks_batch_empty_input_and_missing_ids() { assert!(map.get("ghost:no-such-2").is_none()); } +// ---------- get_chunk_lifecycle_statuses_tx ---------- +// +// Contract: batched, transaction-scoped form of the per-chunk lifecycle read. +// Equivalent to looping the single-id `get_chunk_lifecycle_status` but in +// `O(ceil(n / MAX_FETCH_BATCH))` round-trips. Only ids with a present, +// non-null row appear in the map; misses are silently absent (same as the +// single-id `Ok(None)`) — which is what the ingest pipeline relies on to +// distinguish genuinely-new chunks from ones already past extraction. + +#[test] +fn lifecycle_statuses_batch_matches_single_reads() { + let (_tmp, cfg) = test_config(); + let c1 = sample_chunk("slack:#eng", 0, 1_700_000_000_000); + let c2 = sample_chunk("slack:#eng", 1, 1_700_000_000_000); + let c3 = sample_chunk("slack:#ops", 0, 1_700_000_000_000); + upsert_chunks(&cfg, &[c1.clone(), c2.clone(), c3.clone()]).unwrap(); + set_chunk_lifecycle_status(&cfg, &c1.id, "admitted").unwrap(); + set_chunk_lifecycle_status(&cfg, &c2.id, "sealed").unwrap(); + // c3 keeps whatever the insert default is — the batch must still agree + // with the single-id read for it, without the test hardcoding that value. + + let ids = vec![ + c1.id.clone(), + c2.id.clone(), + c3.id.clone(), + "ghost:no-such".to_string(), + ]; + let map = with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + let m = get_chunk_lifecycle_statuses_tx(&tx, &ids)?; + tx.commit()?; + Ok(m) + }) + .unwrap(); + + assert_eq!(map.get(&c1.id).map(String::as_str), Some("admitted")); + assert_eq!(map.get(&c2.id).map(String::as_str), Some("sealed")); + // Batch result must equal looping the single-id read, for every id. + for id in [&c1.id, &c2.id, &c3.id] { + let single = get_chunk_lifecycle_status(&cfg, id).unwrap(); + assert_eq!( + map.get(id), + single.as_ref(), + "batch must match single-id read for {id}" + ); + } + // Missing ids are silently absent (mirrors the single-id `Ok(None)`). + assert!(map.get("ghost:no-such").is_none()); +} + +#[test] +fn lifecycle_statuses_batch_empty_input_issues_no_query() { + let (_tmp, cfg) = test_config(); + let map = with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + let m = get_chunk_lifecycle_statuses_tx(&tx, &[])?; + tx.commit()?; + Ok(m) + }) + .unwrap(); + assert!(map.is_empty()); +} + // ---------- get_chunk_embeddings_for_signature_batch ---------- // // Contract: equivalent to looping `get_chunk_embedding_for_signature`