From 5454ae9d2f85aca7cf498b2ef05db7508d0d301c Mon Sep 17 00:00:00 2001 From: Beinan Date: Thu, 10 Sep 2026 21:37:09 +0000 Subject: [PATCH] fix: bound MemWAL merge memory with a byte budget --- .../lance-context-core/src/datagen_store.rs | 14 +- .../lance-context-core/src/generic_store.rs | 14 +- .../lance-context-core/src/rollout_store.rs | 34 ++++- crates/lance-context-core/src/store.rs | 16 ++- crates/lance-context-core/src/store_base.rs | 69 +++++----- .../tests/wal_merge_generation_cleanup.rs | 123 +++++++++++++++++- crates/lance-context-server/src/config.rs | 51 +++++--- .../src/routes/datagen.rs | 1 + .../src/routes/generic.rs | 1 + .../src/routes/rollouts.rs | 1 + crates/lance-context-server/src/state.rs | 11 +- crates/lance-context/src/unified_datagen.rs | 1 + crates/lance-context/src/unified_generic.rs | 2 + deploy/kubernetes/README.md | 5 + docs/src/specs/rollout-deployment.md | 13 +- 15 files changed, 270 insertions(+), 86 deletions(-) diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs index b967a46..cf78c7a 100644 --- a/crates/lance-context-core/src/datagen_store.rs +++ b/crates/lance-context-core/src/datagen_store.rs @@ -51,13 +51,14 @@ pub struct DatagenStoreOptions { /// threshold is reached. `None` or zero disables count-triggered merging. pub merge_after_generations: Option, /// Maximum flushed generations folded into the base table by one merge - /// pass. `None` uses the crate default (8); `Some(0)` means unbounded. - /// - /// A merge buffers every row of every generation it takes before appending, - /// so this caps peak merge memory. Leftover generations stay pending for the - /// next pass. Raise it only if merge commits are the bottleneck and the - /// rows are known to be small. + /// pass. `None` uses the crate default (8); `Some(0)` disables this cap. + /// The byte budget applies independently. Leftovers stay pending. pub merge_max_generations: Option, + /// Buffered Arrow array byte budget per merge pass. `None` uses 1 GiB; + /// `Some(0)` disables only this cap. Stops after the generation that + /// reaches the budget, or the generation-count cap, whichever comes first. + /// A generation is indivisible, so even an oversized one is fully merged. + pub merge_max_bytes: Option, /// Periodically merge this writer's pending generations. `None` or zero /// disables the timer. pub cleanup_interval_secs: Option, @@ -101,6 +102,7 @@ impl DatagenStore { shard_id: options.shard_id.clone(), merge_after_generations: options.merge_after_generations, merge_max_generations: options.merge_max_generations, + merge_max_bytes: options.merge_max_bytes, session: None, schema: Arc::new(datagen_log_schema()), // Datagen keys on `event_id`, not `id`: event ids are derived diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index 4c55247..f568aab 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -61,13 +61,14 @@ pub struct GenericStoreOptions { /// accumulated this many. `None`/`0` disables the count trigger. pub merge_after_generations: Option, /// Maximum flushed generations folded into the base table by one merge - /// pass. `None` uses the crate default (8); `Some(0)` means unbounded. - /// - /// A merge buffers every row of every generation it takes before appending, - /// so this caps peak merge memory. Leftover generations stay pending for the - /// next pass. Raise it only if merge commits are the bottleneck and the - /// rows are known to be small. + /// pass. `None` uses the crate default (8); `Some(0)` disables this cap. + /// The byte budget applies independently. Leftovers stay pending. pub merge_max_generations: Option, + /// Buffered Arrow array byte budget per merge pass. `None` uses 1 GiB; + /// `Some(0)` disables only this cap. Stops after the generation that + /// reaches the budget, or the generation-count cap, whichever comes first. + /// A generation is indivisible, so even an oversized one is fully merged. + pub merge_max_bytes: Option, /// Shared, capacity-bounded Lance session. pub session: Option>, /// Whether [`GenericStore::add`] seals before returning, making the rows it @@ -170,6 +171,7 @@ impl GenericStore { shard_id: options.shard_id, merge_after_generations: options.merge_after_generations, merge_max_generations: options.merge_max_generations, + merge_max_bytes: options.merge_max_bytes, session: options.session, schema: create_schema, // Always `id`: the LSM merge key, which `SchemaSpec::validate` diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 23c4149..313f0c0 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -375,13 +375,14 @@ pub struct RolloutStoreOptions { /// accumulate and are unioned at read time). pub merge_after_generations: Option, /// Maximum flushed generations folded into the base table by one merge - /// pass. `None` uses the crate default (8); `Some(0)` means unbounded. - /// - /// A merge buffers every row of every generation it takes before appending, - /// so this caps peak merge memory. Leftover generations stay pending for the - /// next pass. Raise it only if merge commits are the bottleneck and the - /// rows are known to be small. + /// pass. `None` uses the crate default (8); `Some(0)` disables this cap. + /// The byte budget applies independently. Leftovers stay pending. pub merge_max_generations: Option, + /// Buffered Arrow array byte budget per merge pass. `None` uses 1 GiB; + /// `Some(0)` disables only this cap. Stops after the generation that + /// reaches the budget, or the generation-count cap, whichever comes first. + /// A generation is indivisible, so even an oversized one is fully merged. + pub merge_max_bytes: Option, /// Shared Lance [`Session`] used to open this store's base dataset (and, /// transitively, every flushed MemWAL generation it reads — those inherit /// the base dataset's session). @@ -467,6 +468,7 @@ impl RolloutStore { shard_id, merge_after_generations, merge_max_generations, + merge_max_bytes, session, } = options; let base = StorageBase::open( @@ -476,6 +478,7 @@ impl RolloutStore { shard_id, merge_after_generations, merge_max_generations, + merge_max_bytes, session, schema: Arc::new(rollout_schema()), key_column: "id".to_string(), @@ -2584,6 +2587,7 @@ mod tests { // make this row visible, exactly as with flush interval 0. merge_after_generations: Some(0), merge_max_generations: None, + merge_max_bytes: None, }, ) .await @@ -2627,6 +2631,7 @@ mod tests { shard_id: Some("evicted-0".to_string()), merge_after_generations: None, merge_max_generations: None, + merge_max_bytes: None, }; { @@ -2675,6 +2680,7 @@ mod tests { shard_id: Some("observe-0".to_string()), merge_after_generations: None, merge_max_generations: None, + merge_max_bytes: None, }, ) .await @@ -2725,6 +2731,7 @@ mod tests { shard_id: Some(shard.to_string()), merge_after_generations: None, merge_max_generations: None, + merge_max_bytes: None, }; let instance_a = RolloutStore::open_with_options(&uri, options("rollout-0")) @@ -2945,6 +2952,7 @@ mod tests { shard_id: Some("refresh-writer".to_string()), merge_after_generations: Some(1), merge_max_generations: None, + merge_max_bytes: None, ..Default::default() }, ) @@ -3006,6 +3014,7 @@ mod tests { shard_id: Some("trajectory-test".to_string()), merge_after_generations: Some(1), merge_max_generations: None, + merge_max_bytes: None, ..Default::default() }, ) @@ -3105,6 +3114,7 @@ mod tests { shard_id: Some("rollout-0".to_string()), merge_after_generations: None, // no merge → epoch never reclaimed merge_max_generations: None, + merge_max_bytes: None, }, ) .await @@ -3162,6 +3172,7 @@ mod tests { // the following append always hits the reopen path. merge_after_generations: Some(1), merge_max_generations: None, + merge_max_bytes: None, }, ) .await @@ -3200,6 +3211,7 @@ mod tests { shard_id: Some("rollout-0".to_string()), merge_after_generations: None, merge_max_generations: None, + merge_max_bytes: None, }, ) .await @@ -3335,6 +3347,7 @@ mod tests { // Merge every append into base so each forms its own fragment. merge_after_generations: Some(1), merge_max_generations: None, + merge_max_bytes: None, }, ) .await @@ -3420,6 +3433,7 @@ mod tests { shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(1), merge_max_generations: None, + merge_max_bytes: None, }, ) .await @@ -3466,6 +3480,7 @@ mod tests { shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(1), merge_max_generations: None, + merge_max_bytes: None, }, ) .await @@ -3534,6 +3549,7 @@ mod tests { shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(3), merge_max_generations: None, + merge_max_bytes: None, }, ) .await @@ -3570,6 +3586,7 @@ mod tests { shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(3), merge_max_generations: None, + merge_max_bytes: None, }, ) .await @@ -3627,6 +3644,7 @@ mod tests { shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(2), merge_max_generations: None, + merge_max_bytes: None, }, ) .await @@ -3667,6 +3685,7 @@ mod tests { shard_id: Some("rollout-0".to_string()), merge_after_generations: None, // count trigger off merge_max_generations: None, + merge_max_bytes: None, }, ) .await @@ -4165,6 +4184,7 @@ mod tests { shard_id: Some("pagination-benchmark".to_string()), merge_after_generations: Some(1), merge_max_generations: None, + merge_max_bytes: None, ..Default::default() }, ) @@ -4262,6 +4282,7 @@ mod tests { shard_id: Some("rollout-0".to_string()), merge_after_generations: None, // disabled merge_max_generations: None, + merge_max_bytes: None, }, ) .await @@ -4293,6 +4314,7 @@ mod tests { shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(1), merge_max_generations: None, + merge_max_bytes: None, }, ) .await diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 2dc504b..22c316b 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -284,13 +284,14 @@ pub struct ContextStoreOptions { /// previously had no bound at all on this store. pub merge_after_generations: Option, /// Maximum flushed generations folded into the base table by one merge - /// pass. `None` uses the crate default (8); `Some(0)` means unbounded. - /// - /// A merge buffers every row of every generation it takes before appending, - /// so this caps peak merge memory. Leftover generations stay pending for the - /// next pass. Raise it only if merge commits are the bottleneck and the - /// rows are known to be small. + /// pass. `None` uses the crate default (8); `Some(0)` disables this cap. + /// The byte budget applies independently. Leftovers stay pending. pub merge_max_generations: Option, + /// Buffered Arrow array byte budget per merge pass. `None` uses 1 GiB; + /// `Some(0)` disables only this cap. Stops after the generation that + /// reaches the budget, or the generation-count cap, whichever comes first. + /// A generation is indivisible, so even an oversized one is fully merged. + pub merge_max_bytes: Option, /// Whether [`ContextStore::add`] seals the memtable before returning, so the /// rows it wrote are immediately readable. /// @@ -321,6 +322,7 @@ impl Default for ContextStoreOptions { shard_id: None, merge_after_generations: None, merge_max_generations: None, + merge_max_bytes: None, // Read-your-write by default; see the field docs. seal_on_add: true, } @@ -601,6 +603,7 @@ impl ContextStore { shard_id: options.shard_id.clone(), merge_after_generations: options.merge_after_generations, merge_max_generations: options.merge_max_generations, + merge_max_bytes: options.merge_max_bytes, session: None, schema: Arc::new(arrow_schema.clone()), key_column: "id".to_string(), @@ -2119,6 +2122,7 @@ impl ContextStore { shard_id: None, merge_after_generations: None, merge_max_generations: None, + merge_max_bytes: None, // A compactor never appends, so the seal mode is irrelevant to it; // deferring keeps it from ever emitting a generation. seal_on_add: false, diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index b534a22..3a62257 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -94,6 +94,9 @@ pub(crate) const DEFAULT_OBSERVE_CONCURRENCY: usize = 16; /// takes them. pub(crate) const DEFAULT_MERGE_MAX_GENERATIONS: usize = 8; +/// Buffered Arrow array bytes per merge pass, checked at generation boundaries. +pub(crate) const DEFAULT_MERGE_MAX_BYTES: usize = 1024 * 1024 * 1024; + /// Execute only the first `max_source_fragments` from a Lance compaction plan. /// /// Lance's built-in `max_source_fragments` stops before a whole planned task @@ -215,14 +218,14 @@ pub(crate) struct StorageBaseOptions { /// Count-triggered self-merge threshold; `None`/`0` disables it. pub merge_after_generations: Option, /// Maximum flushed generations folded into the base table by one merge - /// pass. `None`/`0` means unbounded (every pending generation at once). - /// - /// This bounds peak merge memory. `read_flushed_generations` buffers every - /// row of every generation it takes, and rollout rows carry `binary_payload` - /// inline, so an unbounded pass over a backlog materialises the full blob - /// volume at once -- the worker OOM this exists to prevent. Leftover - /// generations stay pending and the next pass takes them. + /// pass. `None` uses the crate default (8); `Some(0)` disables this cap. + /// The byte budget applies independently. Leftovers stay pending. pub merge_max_generations: Option, + /// Buffered Arrow array byte budget per merge pass. `None` uses 1 GiB; + /// `Some(0)` disables only this cap. Stops after the generation that + /// reaches the budget, or the generation-count cap, whichever comes first. + /// A generation is indivisible, so even an oversized one is fully merged. + pub merge_max_bytes: Option, /// Shared, capacity-bounded Lance session. `None` preserves Lance's /// per-open default (a fresh 6 GiB index + 1 GiB metadata session *per /// store*, which is the source of unbounded per-append RSS growth). @@ -283,6 +286,7 @@ pub(crate) struct StorageBase { /// Self-merge threshold; `0` disables it. merge_after_generations: usize, merge_max_generations: usize, + merge_max_bytes: usize, /// Timestamp of the last successful [`Self::compact`] on this handle. last_compaction: Option>, /// Number of successful compactions performed by this handle. @@ -336,6 +340,7 @@ impl StorageBase { shard_id, merge_after_generations, merge_max_generations, + merge_max_bytes, session, schema, key_column, @@ -365,6 +370,7 @@ impl StorageBase { shard_id, merge_after_generations, merge_max_generations, + merge_max_bytes, session, schema, key_column, @@ -388,6 +394,7 @@ impl StorageBase { shard_id, merge_after_generations, merge_max_generations, + merge_max_bytes, session, schema, key_column, @@ -412,6 +419,7 @@ impl StorageBase { seal_on_put, merge_after_generations: merge_after_generations.unwrap_or(0), merge_max_generations: merge_max_generations.unwrap_or(DEFAULT_MERGE_MAX_GENERATIONS), + merge_max_bytes: merge_max_bytes.unwrap_or(DEFAULT_MERGE_MAX_BYTES), last_compaction: None, total_compactions: 0, last_compaction_error: None, @@ -801,7 +809,7 @@ impl StorageBase { } /// The `&self` half of a merge: everything that can run while appends - /// continue — sealing the memtable and reading every flushed generation + /// continue — sealing the memtable and reading the budgeted generations /// into memory. /// /// # Concurrency: the expensive phase does not need exclusive access @@ -824,7 +832,7 @@ impl StorageBase { // close the writer, so appends continue throughout. observe_phase!("seal", self.flush().await)?; - // The expensive phase: pull every flushed generation out of object + // The expensive phase: pull a budgeted prefix of generations out of object // storage. Buffered in memory, so this is the part that must not hold an // exclusive lock. let (merged_generations, merged_paths, batches, merge_schema) = @@ -967,7 +975,7 @@ impl StorageBase { Ok(()) } - /// Read every flushed generation listed in `manifest` into memory, aligned + /// Read a budgeted prefix of flushed generations in `manifest`, aligned /// to the base table's current schema. /// /// Returns the merged generation ids, their on-storage folder names (needed @@ -984,32 +992,16 @@ impl StorageBase { let mut generation_batches: Vec<(u64, Vec)> = Vec::new(); let merge_schema: Arc = Arc::new(self.dataset.schema().into()); - // Read at most `merge_max_generations` generations per pass. - // - // Every row of every generation is buffered here and stays resident in - // `PreparedMerge.batches` until the commit appends it, so peak memory - // for one merge is the *total* size of the generations taken. Rollout - // rows carry `binary_payload` inline (blob-v2 offload reads back as - // `None` through the LSM scanner, so it cannot be used here), which - // means multi-MB artifacts are in these batches. Unbounded, one pass - // over a large backlog materialises hundreds of MB to several GiB; on - // glibc that memory is freed logically but retained in the allocator's - // arenas, so worker RSS ratchets up a step per merge and never returns. + // Stop when either cap binds. Count the aligned batches before dedup: + // all of their arrays (including inline blobs) are resident during this + // read phase, even if dedup will later discard some rows. // - // Capping the *count* is what bounds it, and it is safe because a merge - // of a subset is already a first-class case: `commit_merge` drains only - // the generation ids it actually merged (a relative, retain-not-in-set - // edit) and deletes only those directories, precisely so a concurrent - // flush is not clobbered. Whatever is left over stays pending and the - // next pass takes it -- the same incremental-progress shape as the - // master-side compaction budget in #229. - // - // Generations are the granularity because the drain removes whole ids; - // splitting one generation's rows across two passes would leave rows - // committed to the base table with the generation still listed. That is - // read-safe (the LSM dedups by key) but would re-read and re-append - // those rows on the next pass, so the budget stops at a generation - // boundary. + // The manifest drains whole generation ids, so a generation is the + // smallest indivisible unit. Finish the generation that reaches the + // byte budget and stop before opening the next one. This may overshoot + // by one generation, but always makes progress, including when the + // first generation alone exceeds the budget. Leftovers stay pending. + let mut buffered_bytes = 0usize; let budget = if self.merge_max_generations == 0 { manifest.flushed_generations.len() } else { @@ -1031,12 +1023,17 @@ impl StorageBase { let mut stream = gen_dataset.scan().try_into_stream().await?; while let Some(batch) = stream.try_next().await? { if batch.num_rows() > 0 { - current_batches.push(align_batch_to_schema(batch, merge_schema.clone())?); + let batch = align_batch_to_schema(batch, merge_schema.clone())?; + buffered_bytes = buffered_bytes.saturating_add(batch.get_array_memory_size()); + current_batches.push(batch); } } generation_batches.push((flushed.generation, current_batches)); merged_generations.insert(flushed.generation); merged_paths.push(flushed.path.clone()); + if self.merge_max_bytes != 0 && buffered_bytes >= self.merge_max_bytes { + break; + } } let batches = dedupe_merge_batches(generation_batches, &self.key_column, merge_schema.clone())?; diff --git a/crates/lance-context-core/tests/wal_merge_generation_cleanup.rs b/crates/lance-context-core/tests/wal_merge_generation_cleanup.rs index 0848448..729d605 100644 --- a/crates/lance-context-core/tests/wal_merge_generation_cleanup.rs +++ b/crates/lance-context-core/tests/wal_merge_generation_cleanup.rs @@ -246,10 +246,9 @@ async fn merge_pass_is_bounded_and_leftovers_survive() { ); } -/// `merge_max_generations: Some(0)` restores the unbounded behavior, so a -/// deployment can opt out without reverting. +/// Disabling both caps restores the unbounded behavior. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn zero_max_generations_merges_everything_in_one_pass() { +async fn zero_merge_caps_merge_everything_in_one_pass() { let tmp = tempfile::tempdir().unwrap(); let uri = tmp.path().to_string_lossy().to_string(); @@ -257,6 +256,7 @@ async fn zero_max_generations_merges_everything_in_one_pass() { shard_id: Some("unbounded".to_string()), merge_after_generations: None, merge_max_generations: Some(0), + merge_max_bytes: Some(0), ..Default::default() }; let mut store = RolloutStore::open_with_options(&uri, opts).await.unwrap(); @@ -273,3 +273,120 @@ async fn zero_max_generations_merges_everything_in_one_pass() { ); assert_eq!(store.observe().await.unwrap().pending_wal_generations, 0); } + +/// The same five generations must drain differently as either cap binds. +/// Small budgets exercise the large-blob behavior without allocating GiBs. +async fn assert_blob_merge_passes(max_generations: usize, max_bytes: usize, passes: &[usize]) { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let opts = RolloutStoreOptions { + shard_id: Some("byte-budget".to_string()), + merge_max_generations: Some(max_generations), + merge_max_bytes: Some(max_bytes), + ..Default::default() + }; + let mut store = RolloutStore::open_with_options(&uri, opts).await.unwrap(); + let mut expected = Vec::new(); + for i in 0..5 { + let mut record = rec(&format!("blob-{i}")); + record.binary_payload = Some(vec![i as u8; 1024 * 1024]); + store.add(std::slice::from_ref(&record)).await.unwrap(); + store.flush().await.unwrap(); + expected.push(record); + } + + let mut pending = expected.len(); + for &reclaimed in passes { + assert_eq!(store.cleanup_own_shard().await.unwrap(), reclaimed); + pending -= reclaimed; + assert_eq!( + store.observe().await.unwrap().pending_wal_generations, + pending as i64, + "only merged generations may be drained" + ); + assert_eq!(count_gen_dirs_on_disk(tmp.path()), pending); + // Check the union of base and pending generations after every pass, + // including the actual inline bytes, not just ids or row counts. + let mut listed = store.list(None, None).await.unwrap(); + listed.sort_by(|a, b| a.id.cmp(&b.id)); + assert_eq!(listed.len(), expected.len()); + for (actual, expected) in listed.iter().zip(&expected) { + assert_eq!(actual.id, expected.id); + assert_eq!( + store.get_blob(&actual.id).await.unwrap(), + expected.binary_payload + ); + } + } + assert_eq!(pending, 0); + assert_eq!(store.cleanup_own_shard().await.unwrap(), 0); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn byte_budget_binds_before_generation_cap() { + // Each generation is below 1.5 MiB, but two together exceed it. + assert_blob_merge_passes(8, 1536 * 1024, &[2, 2, 1]).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn byte_budget_applies_with_generation_cap_disabled() { + assert_blob_merge_passes(0, 1536 * 1024, &[2, 2, 1]).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn generation_cap_binds_before_byte_budget() { + assert_blob_merge_passes(2, 16 * 1024 * 1024, &[2, 2, 1]).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn zero_byte_budget_keeps_generation_cap() { + assert_blob_merge_passes(2, 0, &[2, 2, 1]).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn oversized_generation_is_merged_in_full_and_makes_progress() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let opts = RolloutStoreOptions { + shard_id: Some("oversized".to_string()), + merge_after_generations: Some(1), + merge_max_generations: Some(8), + merge_max_bytes: Some(1), + ..Default::default() + }; + let mut store = RolloutStore::open_with_options(&uri, opts).await.unwrap(); + for generation in 0..2 { + let records: Vec<_> = (0..3) + .map(|row| { + let mut record = rec(&format!("{generation}-{row}")); + record.binary_payload = Some(vec![row; 4096]); + record + }) + .collect(); + store.add(&records).await.unwrap(); + store.flush().await.unwrap(); + } + // Exercise the count-triggered prepare/commit path as well as cleanup. + for pending in [1, 0] { + assert_eq!(store.maybe_merge_own_shard().await.unwrap(), 1); + assert_eq!( + store.observe().await.unwrap().pending_wal_generations, + pending + ); + assert_eq!(count_gen_dirs_on_disk(tmp.path()), pending as usize); + let listed = store.list(None, None).await.unwrap(); + assert_eq!( + listed.len(), + 6, + "a generation must never be partially drained" + ); + for record in listed { + let row = record.id.as_bytes()[2] - b'0'; + assert_eq!( + store.get_blob(&record.id).await.unwrap(), + Some(vec![row; 4096]) + ); + } + } + assert_eq!(store.maybe_merge_own_shard().await.unwrap(), 0); +} diff --git a/crates/lance-context-server/src/config.rs b/crates/lance-context-server/src/config.rs index 8642086..b4ad4e4 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -37,27 +37,19 @@ pub struct ServerConfig { #[arg(long, env = "ROLLOUT_MERGE_AFTER_GENERATIONS", default_value = "0")] pub rollout_merge_after_generations: usize, - /// Maximum flushed MemWAL generations folded into the base table by a single - /// merge pass. `0` means unbounded (every pending generation at once). - /// - /// # Why this is capped - /// - /// A merge reads every row of every generation it takes into memory before - /// appending, and rollout rows carry `binary_payload` inline (blob-v2 - /// offload reads back as `None` through the MemWAL scanner, so it cannot be - /// used). Peak memory for one pass is therefore the total artifact volume of - /// the generations it took. Unbounded, a worker draining a backlog - /// materialised multiple GiB at once; glibc frees that logically but retains - /// it in its arenas, so RSS ratcheted up a step per merge and workers were - /// eventually OOMKilled. - /// - /// Leftover generations are not dropped -- they stay pending and the next - /// pass takes them. This bounds *per-pass* memory, so it is the knob that - /// matters for RSS; `--rollout-merge-after-generations` only controls when - /// the count trigger fires, and the time-based cleanup ignores it entirely. + /// Maximum flushed MemWAL generations folded into the base table by one + /// merge pass. `0` disables this cap; the byte budget still applies. + /// Leftovers stay pending for the next pass. Both caps apply to count- and + /// time-triggered merges of rollout, datagen and generic stores. #[arg(long, env = "ROLLOUT_MERGE_MAX_GENERATIONS", default_value = "8")] pub rollout_merge_max_generations: usize, + /// Buffered Arrow array byte budget per merge pass; `0` disables this cap. + /// Checked after each whole generation, so one oversized generation + /// still makes progress. The generation-count cap applies independently. + #[arg(long, env = "ROLLOUT_MERGE_MAX_BYTES", default_value = "1073741824")] + pub rollout_merge_max_bytes: usize, + /// Interval, in seconds, for the periodic per-shard WAL cleanup task. When /// non-zero, the global sweeper folds this instance's flushed MemWAL /// generations into the base table on a schedule — the *time* half of the @@ -220,7 +212,30 @@ mod tests { fn parses_with_only_a_data_dir() { let config = ServerConfig::try_parse_from(["lance-context-server", "--data-dir", "/tmp/x"]) .expect("the documented minimal invocation must parse"); + assert_eq!(config.rollout_merge_max_bytes, 1024 * 1024 * 1024); assert_eq!(config.rollout_flush_interval_secs, 30); assert_eq!(config.rollout_cleanup_interval_secs, 0); } + + #[test] + fn merge_byte_budget_flag_and_env_binding() { + let command = ServerConfig::command(); + let arg = command + .get_arguments() + .find(|arg| arg.get_id() == "rollout_merge_max_bytes") + .unwrap(); + assert_eq!(arg.get_env().unwrap(), "ROLLOUT_MERGE_MAX_BYTES"); + for (value, expected) in [("0", 0), ("1048576", 1048576)] { + let config = ServerConfig::try_parse_from([ + "lance-context-server", + "--rollout-merge-max-bytes", + value, + "--rollout-merge-max-generations", + "3", + ]) + .unwrap(); + assert_eq!(config.rollout_merge_max_bytes, expected); + assert_eq!(config.rollout_merge_max_generations, 3); + } + } } diff --git a/crates/lance-context-server/src/routes/datagen.rs b/crates/lance-context-server/src/routes/datagen.rs index 98ec618..f8bb19e 100644 --- a/crates/lance-context-server/src/routes/datagen.rs +++ b/crates/lance-context-server/src/routes/datagen.rs @@ -46,6 +46,7 @@ pub async fn create_datagen_store( shard_id: state.instance_id.clone(), merge_after_generations: None, merge_max_generations: Some(state.rollout_merge_max_generations), + merge_max_bytes: Some(state.rollout_merge_max_bytes), cleanup_interval_secs: None, }; let store = DatagenStore::open_with_options(&uri, options) diff --git a/crates/lance-context-server/src/routes/generic.rs b/crates/lance-context-server/src/routes/generic.rs index 87477a1..c9c4f43 100644 --- a/crates/lance-context-server/src/routes/generic.rs +++ b/crates/lance-context-server/src/routes/generic.rs @@ -57,6 +57,7 @@ pub async fn create_generic_store( shard_id: state.instance_id.clone(), merge_after_generations: None, merge_max_generations: Some(state.rollout_merge_max_generations), + merge_max_bytes: Some(state.rollout_merge_max_bytes), session: None, seal_on_add: req.seal_on_add, }; diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index 64ad090..c36583b 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -180,6 +180,7 @@ pub async fn create_rollout_store( merge_after_generations: (state.rollout_merge_after_generations > 0) .then_some(state.rollout_merge_after_generations), merge_max_generations: Some(state.rollout_merge_max_generations), + merge_max_bytes: Some(state.rollout_merge_max_bytes), session: state.rollout_session.clone(), }; diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index 19ab7af..e0a6e27 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -62,9 +62,11 @@ pub struct AppState { /// disables it. See `RolloutStoreOptions::merge_after_generations`. pub rollout_merge_after_generations: usize, /// Per-pass cap on flushed generations folded into the base table. `0` is - /// unbounded. Bounds peak merge memory; see - /// `Config::rollout_merge_max_generations`. + /// disabled; the byte budget still applies. See + /// `ServerConfig::rollout_merge_max_generations`. pub rollout_merge_max_generations: usize, + /// Buffered Arrow array byte budget per merge pass; `0` disables this cap. + pub rollout_merge_max_bytes: usize, /// Periodic per-shard WAL-cleanup interval in seconds; `0` disables the /// global sweeper. See [`Self::spawn_global_sweeper`]. pub rollout_cleanup_interval_secs: u64, @@ -232,6 +234,7 @@ impl AppState { instance_id, rollout_merge_after_generations: config.rollout_merge_after_generations, rollout_merge_max_generations: config.rollout_merge_max_generations, + rollout_merge_max_bytes: config.rollout_merge_max_bytes, rollout_cleanup_interval_secs: config.rollout_cleanup_interval_secs, rollout_flush_interval_secs: config.rollout_flush_interval_secs, blob_budget, @@ -293,6 +296,7 @@ impl AppState { instance_id, rollout_merge_after_generations: 0, rollout_merge_max_generations: 8, + rollout_merge_max_bytes: 1024 * 1024 * 1024, rollout_cleanup_interval_secs: 0, rollout_flush_interval_secs: 0, blob_budget: None, @@ -325,6 +329,7 @@ impl AppState { merge_after_generations: (self.rollout_merge_after_generations > 0) .then_some(self.rollout_merge_after_generations), merge_max_generations: Some(self.rollout_merge_max_generations), + merge_max_bytes: Some(self.rollout_merge_max_bytes), session: self.rollout_session.clone(), } } @@ -477,6 +482,7 @@ impl AppState { merge_after_generations: (self.rollout_merge_after_generations > 0) .then_some(self.rollout_merge_after_generations), merge_max_generations: Some(self.rollout_merge_max_generations), + merge_max_bytes: Some(self.rollout_merge_max_bytes), cleanup_interval_secs: None, } } @@ -581,6 +587,7 @@ impl AppState { merge_after_generations: (self.rollout_merge_after_generations > 0) .then_some(self.rollout_merge_after_generations), merge_max_generations: Some(self.rollout_merge_max_generations), + merge_max_bytes: Some(self.rollout_merge_max_bytes), session: self.rollout_session.clone(), seal_on_add, } diff --git a/crates/lance-context/src/unified_datagen.rs b/crates/lance-context/src/unified_datagen.rs index e8caef8..daf6a21 100644 --- a/crates/lance-context/src/unified_datagen.rs +++ b/crates/lance-context/src/unified_datagen.rs @@ -37,6 +37,7 @@ impl DatagenStore { shard_id: None, merge_after_generations: None, merge_max_generations: None, + merge_max_bytes: None, cleanup_interval_secs: None, }; let store = LocalStore::open_with_options(uri, options) diff --git a/crates/lance-context/src/unified_generic.rs b/crates/lance-context/src/unified_generic.rs index f5f2815..9268554 100644 --- a/crates/lance-context/src/unified_generic.rs +++ b/crates/lance-context/src/unified_generic.rs @@ -40,6 +40,7 @@ impl GenericStore { shard_id: None, merge_after_generations: None, merge_max_generations: None, + merge_max_bytes: None, session: None, seal_on_add, }; @@ -60,6 +61,7 @@ impl GenericStore { shard_id: None, merge_after_generations: None, merge_max_generations: None, + merge_max_bytes: None, session: None, seal_on_add, }; diff --git a/deploy/kubernetes/README.md b/deploy/kubernetes/README.md index 7177f53..620791f 100644 --- a/deploy/kubernetes/README.md +++ b/deploy/kubernetes/README.md @@ -50,3 +50,8 @@ Two periodic sweeps run on the master and feed the shared scheduler queue: The master cannot merge a shard it does not own without fencing the live writer, so the merge itself always runs on the owning worker. Set the interval to `0` to disable it; the manual "Merge WAL" / "Optimize" UI actions still work. + On workers, `ROLLOUT_MERGE_MAX_BYTES` (default `1073741824`, 1 GiB) and + `ROLLOUT_MERGE_MAX_GENERATIONS` (default `8`) independently cap each pass. + The first cap reached ends the pass after a whole generation; an oversized + generation still merges in full, and leftovers drain on later passes. + Setting either cap to `0` disables only that cap. diff --git a/docs/src/specs/rollout-deployment.md b/docs/src/specs/rollout-deployment.md index 276197d..ae9533e 100644 --- a/docs/src/specs/rollout-deployment.md +++ b/docs/src/specs/rollout-deployment.md @@ -168,11 +168,18 @@ To bound this, an instance can **merge its own shard back into the base table** 2. appends their rows to the **base table** (`Dataset::append`), and 3. `commit_update`s the shard manifest to drain **exactly the generations it merged** out of `flushed_generations` — leaving `replay_after_wal_entry_position` untouched, so a reopened writer never re-replays already-merged WAL entries. -**Bounding merge memory (`--rollout-merge-max-generations`, env `ROLLOUT_MERGE_MAX_GENERATIONS`, default 8).** Step 1 buffers the rows it is about to append **entirely in memory** — `Dataset::append` takes a synchronous `RecordBatchReader`, so the batches cannot be streamed lazily off object storage without pushing that IO into the writer's commit window. Peak merge memory is therefore proportional to the bytes in the generations folded by one pass, and rollout rows carry inline `binary_payload` blobs. Left unbounded, a shard that accumulated many large generations makes a single merge allocate the whole shard at once — the failure mode that OOM-killed workers in production. +**Bounding merge memory.** Each pass buffers the selected generations' rows in memory before committing them to the base table. Inline `binary_payload` blobs can dominate this memory, and generation sizes vary with the flush window and workload. Two independent caps apply to both count-triggered merges and time-triggered cleanup: -`M` caps that: one pass folds at most `M` generations and leaves the rest pending for the next pass. This is safe because the step-3 drain is *surgical* — it removes only the generations actually merged, rather than clearing the list — so a partial merge is just a smaller version of a full one, with the same crash-safety argument below. Repeated passes drain the backlog incrementally at bounded peak memory. `M = 0` restores the unbounded "fold everything in one pass" behavior. +| Setting | Default | Effect | +|---|---|---| +| `--rollout-merge-max-generations` / `ROLLOUT_MERGE_MAX_GENERATIONS` | `8` | Maximum generations folded per pass. | +| `--rollout-merge-max-bytes` / `ROLLOUT_MERGE_MAX_BYTES` | `1073741824` (1 GiB) | Stop after the generation that brings buffered Arrow array memory to or above this budget. | + +The caps form an **OR**: whichever binds first ends the pass. Bytes are measured with `RecordBatch::get_array_memory_size()` after schema alignment and before deduplication, including inline blobs and other columns. Embedded Rust callers can set `merge_max_bytes` on their store options; `None` uses the same 1 GiB default. Setting either cap to `0` disables only that cap; both must be `0` to fold every pending generation in one pass. + +The manifest reclaims whole generations, so a generation is the smallest indivisible merge unit. A pass always finishes the generation that reaches the budget, even if its first generation alone is oversized. Thus the read buffer can exceed the budget by up to one generation; this is not a hard process-RSS limit, and scan, deduplication and commit allocations need additional headroom. -Note that `M` binds independently of `N`: the time-triggered cleanup path merges at a threshold of 1 generation, so it does not consult `N` at all, but it is still capped by `M`. +Only merged generation ids and directories are removed. Leftovers stay pending and drain on subsequent passes. The caps apply independently of the count-trigger threshold (`--rollout-merge-after-generations`), including when time-triggered cleanup merges at a threshold of one generation. Server-managed rollout, datagen and generic stores share these settings. This is the "external compactor" path that Lance's MemWAL LSM design explicitly anticipates. Two properties make it safe under the §2 deployment model: