From 382f1e2afcc7801af563785de0599e037afbb591 Mon Sep 17 00:00:00 2001 From: Beinan Date: Sun, 16 Aug 2026 22:28:44 +0000 Subject: [PATCH] fix(core): bound flushed generations folded per merge pass Worker OOM: prepare_merge materialized an entire shard's flushed generations into memory at once. 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 (the stop-the-world append stall that wal_merge_concurrency.rs pins). Peak merge memory was therefore proportional to the whole shard -- and rollout rows carry inline binary_payload blobs, because blob-v2 offload reads back as None through the MemWAL LSM scanner. In production this OOMKilled 8 of 20 workers, with worker-2 at 23.3 GiB RSS. Cap one merge pass at merge_max_generations (default 8, env ROLLOUT_MERGE_MAX_GENERATIONS); leftovers stay pending for the next pass. A subset merge is safe because commit_merge's drain is already surgical -- it filters out exactly the generations that were merged rather than clearing the list -- so a partial merge is a smaller version of a full one with the same crash-safety argument. Generations are the granularity because each is a self-contained Lance dataset and the manifest tracks them individually. The default binds even where the count trigger does not: the time-triggered cleanup path merges at a hardcoded threshold of 1, so it never consults ROLLOUT_MERGE_AFTER_GENERATIONS (50 in the deployment that OOMed). 0 opts out, restoring the unbounded behavior. Tests: merge_pass_is_bounded_and_leftovers_survive (10 generations, cap 3 -- asserts exactly 3 reclaimed per pass, leftovers stay pending, >=4 passes to drain, all 10 rows survive, 0 leaked generation dirs) and zero_max_generations_merges_everything_in_one_pass. Verified negatively: removing the bound fails the first with left: 10, right: 3. Co-Authored-By: Claude Opus 5 --- .../lance-context-core/src/datagen_store.rs | 9 ++ .../lance-context-core/src/generic_store.rs | 9 ++ .../lance-context-core/src/rollout_store.rs | 29 ++++ crates/lance-context-core/src/store.rs | 11 ++ crates/lance-context-core/src/store_base.rs | 65 ++++++++- .../tests/wal_merge_generation_cleanup.rs | 130 ++++++++++++++++++ crates/lance-context-server/src/config.rs | 21 +++ .../src/routes/datagen.rs | 1 + .../src/routes/generic.rs | 1 + .../src/routes/rollouts.rs | 1 + crates/lance-context-server/src/state.rs | 9 ++ crates/lance-context/src/unified_datagen.rs | 1 + crates/lance-context/src/unified_generic.rs | 2 + docs/src/specs/rollout-deployment.md | 10 +- 14 files changed, 296 insertions(+), 3 deletions(-) diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs index a1abaf9..84140dc 100644 --- a/crates/lance-context-core/src/datagen_store.rs +++ b/crates/lance-context-core/src/datagen_store.rs @@ -50,6 +50,14 @@ pub struct DatagenStoreOptions { /// Merge this writer's flushed generations into the base table after the /// 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. + pub merge_max_generations: Option, /// Periodically merge this writer's pending generations. `None` or zero /// disables the timer. pub cleanup_interval_secs: Option, @@ -92,6 +100,7 @@ impl DatagenStore { storage_options: options.storage_options.clone(), shard_id: options.shard_id.clone(), merge_after_generations: options.merge_after_generations, + merge_max_generations: options.merge_max_generations, 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 a9e8fdf..6e45b7b 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -60,6 +60,14 @@ pub struct GenericStoreOptions { /// Fold this instance's flushed generations into the base table once it has /// 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. + pub merge_max_generations: Option, /// Shared, capacity-bounded Lance session. pub session: Option>, /// Whether [`GenericStore::add`] seals before returning, making the rows it @@ -161,6 +169,7 @@ impl GenericStore { storage_options: options.storage_options, shard_id: options.shard_id, merge_after_generations: options.merge_after_generations, + merge_max_generations: options.merge_max_generations, 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 f1419df..1fa7141 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -337,6 +337,14 @@ pub struct RolloutStoreOptions { /// `None` or `0` disables self-merge (the 0.6.0 behavior: generations /// 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. + pub merge_max_generations: 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). @@ -421,6 +429,7 @@ impl RolloutStore { storage_options, shard_id, merge_after_generations, + merge_max_generations, session, } = options; let base = StorageBase::open( @@ -429,6 +438,7 @@ impl RolloutStore { storage_options, shard_id, merge_after_generations, + merge_max_generations, session, schema: Arc::new(rollout_schema()), key_column: "id".to_string(), @@ -2463,6 +2473,7 @@ mod tests { // Count trigger disabled: cleanup is the only path that can // make this row visible, exactly as with flush interval 0. merge_after_generations: Some(0), + merge_max_generations: None, }, ) .await @@ -2505,6 +2516,7 @@ mod tests { session: None, shard_id: Some("evicted-0".to_string()), merge_after_generations: None, + merge_max_generations: None, }; { @@ -2552,6 +2564,7 @@ mod tests { session: None, shard_id: Some("observe-0".to_string()), merge_after_generations: None, + merge_max_generations: None, }, ) .await @@ -2601,6 +2614,7 @@ mod tests { session: None, shard_id: Some(shard.to_string()), merge_after_generations: None, + merge_max_generations: None, }; let instance_a = RolloutStore::open_with_options(&uri, options("rollout-0")) @@ -2820,6 +2834,7 @@ mod tests { RolloutStoreOptions { shard_id: Some("refresh-writer".to_string()), merge_after_generations: Some(1), + merge_max_generations: None, ..Default::default() }, ) @@ -2880,6 +2895,7 @@ mod tests { RolloutStoreOptions { shard_id: Some("trajectory-test".to_string()), merge_after_generations: Some(1), + merge_max_generations: None, ..Default::default() }, ) @@ -2978,6 +2994,7 @@ mod tests { session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: None, // no merge → epoch never reclaimed + merge_max_generations: None, }, ) .await @@ -3034,6 +3051,7 @@ mod tests { // Merge on every append → epoch is reclaimed each time, so // the following append always hits the reopen path. merge_after_generations: Some(1), + merge_max_generations: None, }, ) .await @@ -3071,6 +3089,7 @@ mod tests { session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: None, + merge_max_generations: None, }, ) .await @@ -3205,6 +3224,7 @@ mod tests { shard_id: Some("rollout-0".to_string()), // Merge every append into base so each forms its own fragment. merge_after_generations: Some(1), + merge_max_generations: None, }, ) .await @@ -3289,6 +3309,7 @@ mod tests { session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(1), + merge_max_generations: None, }, ) .await @@ -3334,6 +3355,7 @@ mod tests { session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(1), + merge_max_generations: None, }, ) .await @@ -3401,6 +3423,7 @@ mod tests { session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(3), + merge_max_generations: None, }, ) .await @@ -3436,6 +3459,7 @@ mod tests { session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(3), + merge_max_generations: None, }, ) .await @@ -3492,6 +3516,7 @@ mod tests { session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(2), + merge_max_generations: None, }, ) .await @@ -3531,6 +3556,7 @@ mod tests { session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: None, // count trigger off + merge_max_generations: None, }, ) .await @@ -4028,6 +4054,7 @@ mod tests { RolloutStoreOptions { shard_id: Some("pagination-benchmark".to_string()), merge_after_generations: Some(1), + merge_max_generations: None, ..Default::default() }, ) @@ -4124,6 +4151,7 @@ mod tests { session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: None, // disabled + merge_max_generations: None, }, ) .await @@ -4154,6 +4182,7 @@ mod tests { session: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(1), + merge_max_generations: None, }, ) .await diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 932d1c5..f3a5d18 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -283,6 +283,14 @@ pub struct ContextStoreOptions { /// and every read unions all of them — the read amplification that /// 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. + pub merge_max_generations: Option, /// Whether [`ContextStore::add`] seals the memtable before returning, so the /// rows it wrote are immediately readable. /// @@ -312,6 +320,7 @@ impl Default for ContextStoreOptions { distance_metric: None, shard_id: None, merge_after_generations: None, + merge_max_generations: None, // Read-your-write by default; see the field docs. seal_on_add: true, } @@ -591,6 +600,7 @@ impl ContextStore { storage_options, shard_id: options.shard_id.clone(), merge_after_generations: options.merge_after_generations, + merge_max_generations: options.merge_max_generations, session: None, schema: Arc::new(arrow_schema.clone()), key_column: "id".to_string(), @@ -2102,6 +2112,7 @@ impl ContextStore { distance_metric: Some(self.distance_metric), shard_id: None, merge_after_generations: None, + merge_max_generations: 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 33d3740..bd13d1e 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -74,6 +74,22 @@ pub(crate) const DEFAULT_MANIFEST_SCAN_BATCH_SIZE: usize = 16; /// concurrently while collecting observability metrics. pub(crate) const DEFAULT_OBSERVE_CONCURRENCY: usize = 16; +/// Flushed generations folded into the base table by one merge pass, by default. +/// +/// A merge buffers every row of every generation it takes before appending, and +/// rollout rows store `binary_payload` inline, so the pass's peak memory is the +/// total blob volume of the generations it took. Unbounded, a worker merging a +/// backlog materialised several GiB at once; freed to glibc but retained in its +/// arenas, that ratcheted RSS up a step per merge until the pod was OOMKilled. +/// +/// 8 is deliberately well under the deployed `ROLLOUT_MERGE_AFTER_GENERATIONS` +/// (50 in the deployment that OOMed), so the cap actually binds there, while +/// staying high enough that a merge still amortises its fixed costs -- manifest +/// CAS, base-table commit, directory deletes -- over a useful number of +/// generations. Leftovers are not dropped: they stay pending and the next pass +/// takes them. +pub(crate) const DEFAULT_MERGE_MAX_GENERATIONS: usize = 8; + /// 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 @@ -194,6 +210,15 @@ pub(crate) struct StorageBaseOptions { pub shard_id: Option, /// 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. + pub merge_max_generations: 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). @@ -253,6 +278,7 @@ pub(crate) struct StorageBase { seal_on_put: bool, /// Self-merge threshold; `0` disables it. merge_after_generations: usize, + merge_max_generations: usize, /// Timestamp of the last successful [`Self::compact`] on this handle. last_compaction: Option>, /// Number of successful compactions performed by this handle. @@ -305,6 +331,7 @@ impl StorageBase { storage_options, shard_id, merge_after_generations, + merge_max_generations, session, schema, key_column, @@ -333,6 +360,7 @@ impl StorageBase { storage_options, shard_id, merge_after_generations, + merge_max_generations, session, schema, key_column, @@ -355,6 +383,7 @@ impl StorageBase { storage_options, shard_id, merge_after_generations, + merge_max_generations, session, schema, key_column, @@ -378,6 +407,7 @@ impl StorageBase { latest_schema, seal_on_put, merge_after_generations: merge_after_generations.unwrap_or(0), + merge_max_generations: merge_max_generations.unwrap_or(DEFAULT_MERGE_MAX_GENERATIONS), last_compaction: None, total_compactions: 0, last_compaction_error: None, @@ -933,7 +963,40 @@ impl StorageBase { let mut merged_paths: Vec = Vec::new(); let mut batches: Vec = Vec::new(); let merge_schema: Arc = Arc::new(self.dataset.schema().into()); - for flushed in &manifest.flushed_generations { + + // 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. + // + // 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. + let budget = if self.merge_max_generations == 0 { + manifest.flushed_generations.len() + } else { + self.merge_max_generations + }; + + for flushed in manifest.flushed_generations.iter().take(budget) { let gen_uri = format!( "{}/_mem_wal/{}/{}", base_uri, self.write_shard, flushed.path 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 d13351d..0848448 100644 --- a/crates/lance-context-core/tests/wal_merge_generation_cleanup.rs +++ b/crates/lance-context-core/tests/wal_merge_generation_cleanup.rs @@ -143,3 +143,133 @@ async fn serial_merge_deletes_merged_generation_dirs() { leaked their blob directories" ); } + +/// One merge pass must fold at most `merge_max_generations`, and the leftovers +/// must survive to be merged by later passes. +/// +/// A merge buffers every row of every generation it takes before appending, and +/// rollout rows carry `binary_payload` inline, so an unbounded pass over a +/// backlog materialises the whole artifact volume at once. That is what drove +/// worker RSS up a step per merge (glibc retains the freed bulk allocation in +/// its arenas) until pods were OOMKilled. +/// +/// Capping the pass is only safe because merging a *subset* is already a +/// first-class case: the drain removes just the generation ids it merged and +/// deletes just those directories. This test pins both halves of that -- the +/// cap binds, and nothing is lost to it. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn merge_pass_is_bounded_and_leftovers_survive() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + + // Count trigger off: drive merges explicitly so the assertions are about + // one pass, not about when a pass fires. + let opts = RolloutStoreOptions { + shard_id: Some("bounded".to_string()), + merge_after_generations: None, + merge_max_generations: Some(3), + ..Default::default() + }; + + let mut store = RolloutStore::open_with_options(&uri, opts).await.unwrap(); + + // Ten generations pending, well over the cap of 3. + let n = 10; + for i in 0..n { + store.add(&[rec(&format!("row-{i}"))]).await.unwrap(); + store.flush().await.unwrap(); + } + let pending_before = store.observe().await.unwrap().pending_wal_generations; + assert_eq!( + pending_before, n as i64, + "each append should seal one generation" + ); + + // `cleanup_own_shard` is the time-triggered path: threshold 1, so without a + // per-pass cap it would take all ten at once. It must take exactly 3. + let reclaimed = store.cleanup_own_shard().await.unwrap(); + assert_eq!( + reclaimed, 3, + "one pass must fold at most merge_max_generations (3), not the whole backlog" + ); + + // `cleanup_own_shard` seals first, which can add a generation; what matters + // is that the cap removed exactly 3 and the rest are still pending. + let pending_after = store.observe().await.unwrap().pending_wal_generations; + assert_eq!( + pending_after, + pending_before - 3, + "leftover generations must stay pending, not be dropped" + ); + + // Draining takes several passes, and every row survives all of them. + let mut passes = 1; + loop { + let reclaimed = store.cleanup_own_shard().await.unwrap(); + if reclaimed == 0 { + break; + } + assert!( + reclaimed <= 3, + "every pass must respect the cap; got {reclaimed}" + ); + passes += 1; + assert!(passes < 20, "merge failed to converge"); + } + assert!( + passes >= 4, + "10 generations at 3 per pass must take at least 4 passes; took {passes}" + ); + + assert_eq!( + store.observe().await.unwrap().pending_wal_generations, + 0, + "repeated passes must fully drain the backlog" + ); + + // No row was lost or duplicated across the multi-pass drain. + let listed = store.list(None, None).await.unwrap(); + let mut ids: Vec = listed.iter().map(|r| r.id.clone()).collect(); + ids.sort(); + ids.dedup(); + assert_eq!( + ids.len(), + n, + "every appended row must survive a bounded merge" + ); + + // The subset drain must still delete what it merged. + let on_disk = count_gen_dirs_on_disk(Path::new(&uri)); + assert_eq!( + on_disk, 0, + "a bounded merge must still delete merged generation dirs" + ); +} + +/// `merge_max_generations: Some(0)` restores the unbounded behavior, so a +/// deployment can opt out without reverting. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn zero_max_generations_merges_everything_in_one_pass() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + + let opts = RolloutStoreOptions { + shard_id: Some("unbounded".to_string()), + merge_after_generations: None, + merge_max_generations: Some(0), + ..Default::default() + }; + let mut store = RolloutStore::open_with_options(&uri, opts).await.unwrap(); + + for i in 0..6 { + store.add(&[rec(&format!("row-{i}"))]).await.unwrap(); + store.flush().await.unwrap(); + } + + let reclaimed = store.cleanup_own_shard().await.unwrap(); + assert_eq!( + reclaimed, 6, + "0 must mean unbounded: one pass takes all six" + ); + assert_eq!(store.observe().await.unwrap().pending_wal_generations, 0); +} diff --git a/crates/lance-context-server/src/config.rs b/crates/lance-context-server/src/config.rs index 9571e2e..2d36818 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -37,6 +37,27 @@ 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. + #[arg(long, env = "ROLLOUT_MERGE_MAX_GENERATIONS", default_value = "8")] + pub rollout_merge_max_generations: 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 diff --git a/crates/lance-context-server/src/routes/datagen.rs b/crates/lance-context-server/src/routes/datagen.rs index b80ca2b..98ec618 100644 --- a/crates/lance-context-server/src/routes/datagen.rs +++ b/crates/lance-context-server/src/routes/datagen.rs @@ -45,6 +45,7 @@ pub async fn create_datagen_store( storage_options: req.storage_options, shard_id: state.instance_id.clone(), merge_after_generations: None, + merge_max_generations: Some(state.rollout_merge_max_generations), 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 2e56580..87477a1 100644 --- a/crates/lance-context-server/src/routes/generic.rs +++ b/crates/lance-context-server/src/routes/generic.rs @@ -56,6 +56,7 @@ pub async fn create_generic_store( storage_options: req.storage_options, shard_id: state.instance_id.clone(), merge_after_generations: None, + merge_max_generations: Some(state.rollout_merge_max_generations), 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 feab6be..f9c7c5a 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -140,6 +140,7 @@ pub async fn create_rollout_store( shard_id: state.instance_id.clone(), 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), session: state.rollout_session.clone(), }; diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index b718589..dce945e 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -61,6 +61,10 @@ pub struct AppState { /// Count-triggered self-merge threshold for rollout MemWAL shards; `0` /// 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`. + pub rollout_merge_max_generations: 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, @@ -224,6 +228,7 @@ impl AppState { base_uri, instance_id, rollout_merge_after_generations: config.rollout_merge_after_generations, + rollout_merge_max_generations: config.rollout_merge_max_generations, rollout_cleanup_interval_secs: config.rollout_cleanup_interval_secs, rollout_flush_interval_secs: config.rollout_flush_interval_secs, blob_budget, @@ -283,6 +288,7 @@ impl AppState { base_uri, instance_id, rollout_merge_after_generations: 0, + rollout_merge_max_generations: 8, rollout_cleanup_interval_secs: 0, rollout_flush_interval_secs: 0, blob_budget: None, @@ -313,6 +319,7 @@ impl AppState { shard_id: self.instance_id.clone(), 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), session: self.rollout_session.clone(), } } @@ -464,6 +471,7 @@ impl AppState { shard_id: self.instance_id.clone(), 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), cleanup_interval_secs: None, } } @@ -567,6 +575,7 @@ impl AppState { shard_id: self.instance_id.clone(), 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), 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 19533b5..e8caef8 100644 --- a/crates/lance-context/src/unified_datagen.rs +++ b/crates/lance-context/src/unified_datagen.rs @@ -36,6 +36,7 @@ impl DatagenStore { // the core `DatagenStore` directly. shard_id: None, merge_after_generations: None, + merge_max_generations: 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 7a090a7..f5f2815 100644 --- a/crates/lance-context/src/unified_generic.rs +++ b/crates/lance-context/src/unified_generic.rs @@ -39,6 +39,7 @@ impl GenericStore { // the core `GenericStore` directly. shard_id: None, merge_after_generations: None, + merge_max_generations: None, session: None, seal_on_add, }; @@ -58,6 +59,7 @@ impl GenericStore { storage_options, shard_id: None, merge_after_generations: None, + merge_max_generations: None, session: None, seal_on_add, }; diff --git a/docs/src/specs/rollout-deployment.md b/docs/src/specs/rollout-deployment.md index 7f5bdec..276197d 100644 --- a/docs/src/specs/rollout-deployment.md +++ b/docs/src/specs/rollout-deployment.md @@ -164,9 +164,15 @@ By default (`--rollout-merge-after-generations 0`) there is no compaction step: To bound this, an instance can **merge its own shard back into the base table** on a size trigger. Set `--rollout-merge-after-generations N` (env `ROLLOUT_MERGE_AFTER_GENERATIONS`, or `RolloutStoreOptions::merge_after_generations`). After an append flushes a generation, if this instance's shard has accumulated ≥ N un-merged generations, the same `add` call synchronously: -1. reads every flushed generation (each is a self-contained Lance dataset under `_mem_wal/{shard}/`), +1. reads up to `--rollout-merge-max-generations M` flushed generations, oldest first (each is a self-contained Lance dataset under `_mem_wal/{shard}/`), 2. appends their rows to the **base table** (`Dataset::append`), and -3. `commit_update`s the shard manifest to drain `flushed_generations` back to empty — leaving `replay_after_wal_entry_position` untouched, so a reopened writer never re-replays already-merged WAL entries. +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. + +`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. + +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`. This is the "external compactor" path that Lance's MemWAL LSM design explicitly anticipates. Two properties make it safe under the §2 deployment model: