From 16c24ebb3ee9a894251e099bbf803b01c4a63019 Mon Sep 17 00:00:00 2001 From: Beinan Date: Mon, 13 Jul 2026 17:22:27 +0000 Subject: [PATCH 1/2] feat(rollout): periodic per-shard WAL cleanup + concurrency hardening Add a time-triggered per-shard WAL cleanup for the rollout store, complementing the existing count-triggered self-merge. Each instance cleans only the shard it owns and writes, reusing its own writer epoch, so it never fences a concurrent writer. - RolloutStoreOptions gains `cleanup_interval_secs` (time trigger) and `cleanup_min_generations` (count gate for a tick). - `cleanup_own_shard` runs one pass; `spawn_periodic_cleanup` runs a background timer that holds only a Weak ref, so it self-terminates when the store is dropped (e.g. on delete). - Server exposes `--rollout-cleanup-interval-secs` / `--rollout-cleanup-min-generations` and spawns the timer per store. Concurrency hardening (verified against Lance 7.0.0 commit/OCC model): - ensure_mem_wal: two concurrent first-writers both commit CreateIndex, which Lance treats as a hard (non-retried) conflict. Reload and accept "index now present" as success instead of surfacing the conflict. - merge drain: remove exactly the generations merged instead of clearing the vec, so a generation flushed during a merge is retained for the next pass rather than silently dropped. Tests: concurrent two-shard merge into one base table, surgical-drain row conservation, timer-driven cleanup, min-generation gating. Co-Authored-By: Claude Opus 4 --- .../lance-context-core/src/rollout_store.rs | 489 +++++++++++++++++- crates/lance-context-server/src/config.rs | 17 + .../src/routes/rollouts.rs | 11 +- crates/lance-context-server/src/state.rs | 8 + crates/lance-context/src/unified_rollout.rs | 5 + 5 files changed, 510 insertions(+), 20 deletions(-) diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index f0d92744..74b958bc 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -44,7 +44,7 @@ //! (e.g. by `policy_version`), not a table snapshot — reproducible because the //! rows never change. `checkout` remains available for base-table time-travel. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use arrow_array::builder::{ @@ -68,6 +68,8 @@ use lance::index::DatasetIndexExt; use lance::io::{ObjectStoreParams, StorageOptionsAccessor}; use lance::{Error as LanceError, Result as LanceResult}; use lance_index::mem_wal::{ShardManifest, MEM_WAL_INDEX_NAME}; +use tokio::task::JoinHandle; +use tracing::{info, warn}; use uuid::Uuid; use crate::rollout::RolloutRecord; @@ -111,6 +113,28 @@ 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, + /// Interval, in seconds, for the periodic per-shard WAL cleanup task. When + /// set (and non-zero), a background timer wakes up every `interval` seconds + /// and folds this instance's flushed MemWAL generations into the base table + /// — the *time* half of the "time + count" trigger. It complements the + /// synchronous, count-triggered [`Self::merge_after_generations`]: even a + /// low-traffic shard that never crosses the count threshold still gets its + /// generations reclaimed on a timer, so read amplification stays bounded and + /// object storage does not accumulate stale generation datasets. + /// + /// The cleanup runs on the shard's own owner (reusing its writer epoch), so + /// it never fences a concurrent writer — exactly like the count-triggered + /// merge. `None` or `0` disables the periodic task. + /// + /// Spawn the task with [`RolloutStore::spawn_periodic_cleanup`] once the + /// store is behind an `Arc>` (the server's ownership model). + pub cleanup_interval_secs: Option, + /// The *count* half of the periodic cleanup trigger: the timer only merges + /// when this instance's shard has at least this many flushed generations, + /// skipping the pass otherwise (avoids rewriting the base table to reclaim a + /// single small generation). `None` defaults to `1` (clean up whenever any + /// generation is present on a tick). + pub cleanup_min_generations: Option, } /// A Lance-backed store for RL rollout trajectories. @@ -124,6 +148,13 @@ pub struct RolloutStore { /// Self-merge threshold; see [`RolloutStoreOptions::merge_after_generations`]. /// `0` (or `None` normalized to 0) disables it. merge_after_generations: usize, + /// Periodic-cleanup interval in seconds; `0` disables the timer. See + /// [`RolloutStoreOptions::cleanup_interval_secs`]. + cleanup_interval_secs: u64, + /// Minimum flushed generations before a periodic-cleanup tick merges. + /// Normalized to at least `1`. See + /// [`RolloutStoreOptions::cleanup_min_generations`]. + cleanup_min_generations: usize, } impl RolloutStore { @@ -153,6 +184,8 @@ impl RolloutStore { write_shard, storage_options, merge_after_generations: options.merge_after_generations.unwrap_or(0), + cleanup_interval_secs: options.cleanup_interval_secs.unwrap_or(0), + cleanup_min_generations: options.cleanup_min_generations.unwrap_or(1).max(1), }) } @@ -221,6 +254,22 @@ impl RolloutStore { /// If this instance's own shard has at least `merge_after_generations` /// flushed generations, merge them into the base table. No-op otherwise. async fn maybe_merge_own_shard(&mut self) -> LanceResult<()> { + self.merge_own_shard_if_ready(self.merge_after_generations) + .await + .map(|_| ()) + } + + /// Merge this instance's flushed MemWAL generations into the base table when + /// the shard has accumulated at least `threshold` of them. Returns the + /// number of generations reclaimed (`0` if the threshold was not met or the + /// shard has no manifest yet). + /// + /// This is the shared core of both triggers: the synchronous count trigger + /// in [`Self::add`] (with `threshold = merge_after_generations`) and the + /// periodic timer in [`Self::spawn_periodic_cleanup`] (with + /// `threshold = cleanup_min_generations`). Both merge only the shard this + /// instance owns and writes, so the epoch claim never fences another writer. + async fn merge_own_shard_if_ready(&mut self, threshold: usize) -> LanceResult { let object_store = self.dataset.object_store(None).await?; let branch_location = self.dataset.branch_location(); let manifest_store = ShardManifestStore::new( @@ -230,12 +279,92 @@ impl RolloutStore { DEFAULT_MANIFEST_SCAN_BATCH_SIZE, ); let Some(manifest) = manifest_store.read_latest().await? else { - return Ok(()); + return Ok(0); }; - if manifest.flushed_generations.len() < self.merge_after_generations { - return Ok(()); + let pending = manifest.flushed_generations.len(); + if pending == 0 || pending < threshold.max(1) { + return Ok(0); } - self.merge_own_shard(&manifest_store, &manifest).await + self.merge_own_shard(&manifest_store, &manifest).await?; + Ok(pending) + } + + /// Run one periodic WAL-cleanup pass over this instance's own shard: fold any + /// flushed generations into the base table once at least + /// `cleanup_min_generations` have accumulated. Returns the number of + /// generations reclaimed (`0` if below the threshold or nothing pending). + /// + /// Exposed for callers that drive cleanup on their own schedule instead of + /// (or in addition to) the built-in timer from + /// [`Self::spawn_periodic_cleanup`]. Like every merge path here it operates + /// only on the shard this instance owns, so it is safe to call concurrently + /// with this instance's own appends but must not target another instance's + /// shard. + pub async fn cleanup_own_shard(&mut self) -> LanceResult { + self.merge_own_shard_if_ready(self.cleanup_min_generations) + .await + } + + /// Spawn a background timer that periodically reclaims this instance's + /// flushed MemWAL generations into the base table (the *time* trigger). + /// + /// Every `cleanup_interval_secs` seconds the task acquires the write lock and + /// calls [`Self::cleanup_own_shard`], which merges only when at least + /// `cleanup_min_generations` are pending. This bounds read amplification and + /// reclaims stale generation datasets even on shards that never cross the + /// synchronous count threshold ([`RolloutStoreOptions::merge_after_generations`]). + /// + /// Returns `None` (spawning nothing) when `cleanup_interval_secs` is `0` + /// (disabled). Otherwise returns the task handle. + /// + /// The task holds only a [`std::sync::Weak`] reference to the store, so it + /// never keeps a deleted store alive: once the caller drops the last strong + /// `Arc` (e.g. the server removes it from its store map on delete), the next + /// tick fails to upgrade and the loop exits on its own. `abort`ing the + /// returned handle also stops it immediately. + pub fn spawn_periodic_cleanup(store: Arc>) -> Option> { + let interval_secs = { + // Read the interval without holding the lock across the await loop. + let guard = store.try_read().ok()?; + guard.cleanup_interval_secs + }; + if interval_secs == 0 { + return None; + } + + let weak = Arc::downgrade(&store); + Some(tokio::spawn(async move { + let mut ticker = + tokio::time::interval(std::time::Duration::from_secs(interval_secs)); + // Skip the immediate first tick so we don't merge the instant the + // task starts; wait a full interval before the first pass. + ticker.tick().await; + loop { + ticker.tick().await; + // Stop once the store has been dropped by its owner. + let Some(store) = weak.upgrade() else { + return; + }; + let mut guard = store.write().await; + match guard.cleanup_own_shard().await { + Ok(0) => {} + Ok(n) => { + info!( + shard = %guard.write_shard, + reclaimed = n, + "periodic WAL cleanup merged flushed generations" + ); + } + Err(e) => { + warn!( + shard = %guard.write_shard, + error = %e, + "periodic WAL cleanup failed" + ); + } + } + } + })) } /// Fold this instance's flushed MemWAL generations into the base table and @@ -245,10 +374,22 @@ impl RolloutStore { /// anticipates: each flushed generation at `_mem_wal/{shard}/{path}` is a /// self-contained Lance dataset. We read every generation's rows, append /// them to the base table (`Dataset::append`), then `commit_update` the - /// shard manifest to set `flushed_generations` back to empty — leaving + /// shard manifest to remove *exactly the generations we merged* — leaving /// `replay_after_wal_entry_position` untouched so a reopened writer does /// not re-replay already-merged WAL entries. /// + /// # Concurrency: surgical drain, not blanket clear + /// + /// The drain removes only the generation ids this call actually merged, + /// retaining any generation that appears in the manifest afterwards. Under + /// the single-writer-per-shard model this instance is the only writer, so in + /// practice nothing new lands during a merge — but not depending on that is + /// what keeps the drain correct: a blanket `flushed_generations = []` would + /// silently discard any generation flushed between reading the manifest and + /// committing the drain, *without merging it* (data loss). `commit_update` + /// re-reads the latest manifest and applies this closure to it, so the + /// retain filter runs against the current state, not the stale snapshot. + /// /// # Safety of the epoch claim /// /// `claim_epoch` bumps the shard's writer epoch, which would fence any @@ -273,8 +414,10 @@ impl RolloutStore { } // Resolve each flushed generation to its absolute dataset path and read - // all its rows into memory. + // all its rows into memory. Record which generation ids we merge so the + // drain can remove exactly these and nothing else. let base_uri = self.dataset.uri().trim_end_matches('/').to_string(); + let mut merged_generations: HashSet = HashSet::new(); let mut batches: Vec = Vec::new(); for flushed in &manifest.flushed_generations { let gen_uri = format!( @@ -289,6 +432,7 @@ impl RolloutStore { batches.push(batch); } } + merged_generations.insert(flushed.generation); } // Append the merged rows to the base table. @@ -313,14 +457,22 @@ impl RolloutStore { self.dataset.append(reader, Some(params)).await?; } - // Drain the flushed generations from the shard manifest. Claim the - // shard's epoch (safe: we own it) and commit a manifest that keeps - // everything except the now-merged generations. + // Drain the merged generations from the shard manifest. Claim the + // shard's epoch (safe: we own it) and commit a manifest that retains + // every generation except the ones we just folded into the base table. + // Removing only the merged ids (rather than clearing the vec) is what + // makes this safe against a generation that lands after we read the + // manifest: it is preserved for the next merge instead of being dropped. let (epoch, _) = manifest_store.claim_epoch(manifest.shard_spec_id).await?; manifest_store .commit_update(epoch, |current| ShardManifest { version: current.version + 1, - flushed_generations: vec![], + flushed_generations: current + .flushed_generations + .iter() + .filter(|fg| !merged_generations.contains(&fg.generation)) + .cloned() + .collect(), ..current.clone() }) .await?; @@ -332,17 +484,51 @@ impl RolloutStore { /// Subsequent writes see the index already present and skip this. The shard /// a write targets is chosen by the writer (`shard_id`), independent of the /// index's declared sharding strategy. + /// + /// # Concurrent first-writers + /// + /// `initialize_mem_wal` commits a `CreateIndex` transaction, and Lance treats + /// two concurrent `CreateIndex` commits as a hard conflict (not an + /// auto-retried one): when two instances take their very first write at the + /// same time, both observe no index, both try to create it, and the loser + /// gets `RetryableCommitConflict`. That is benign here — the winner created + /// exactly the index we wanted — so we reload and treat "index now present" + /// as success rather than surfacing the conflict to the caller. Any other + /// error propagates. async fn ensure_mem_wal(&mut self) -> LanceResult<()> { - let indices = self.dataset.load_indices().await?; - let has_mem_wal = indices.iter().any(|i| i.name == MEM_WAL_INDEX_NAME); - if !has_mem_wal { - self.dataset - .initialize_mem_wal() - .unsharded() - .execute() + if self.mem_wal_index_present().await? { + return Ok(()); + } + match self + .dataset + .initialize_mem_wal() + .unsharded() + .execute() + .await + { + Ok(()) => Ok(()), + Err(err) => { + // A concurrent first-writer may have created the index between + // our check and our commit. Reload and accept it if so. + self.dataset = Self::load_with_options( + &self.dataset.uri().to_string(), + self.storage_options.clone(), + ) .await?; + if self.mem_wal_index_present().await? { + Ok(()) + } else { + Err(err) + } + } } - Ok(()) + } + + /// Whether the MemWAL index has been initialized on the current dataset + /// handle. + async fn mem_wal_index_present(&self) -> LanceResult { + let indices = self.dataset.load_indices().await?; + Ok(indices.iter().any(|i| i.name == MEM_WAL_INDEX_NAME)) } /// List rollout rows (base table ∪ every instance's flushed MemWAL rows). @@ -1325,6 +1511,7 @@ mod tests { storage_options: None, shard_id: Some(shard.to_string()), merge_after_generations: None, + ..Default::default() }; let mut instance_a = RolloutStore::open_with_options(&uri, options("rollout-0")) @@ -1385,6 +1572,7 @@ mod tests { storage_options: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(3), + ..Default::default() }, ) .await @@ -1417,6 +1605,7 @@ mod tests { storage_options: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(3), + ..Default::default() }, ) .await @@ -1468,6 +1657,7 @@ mod tests { storage_options: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(2), + ..Default::default() }, ) .await @@ -1486,6 +1676,265 @@ mod tests { }); } + #[test] + fn cleanup_own_shard_merges_only_at_min_generations() { + // The periodic-cleanup entry point (`cleanup_own_shard`) is the time + // trigger's per-pass body. With count self-merge disabled, generations + // accumulate; a cleanup pass below `cleanup_min_generations` is a no-op, + // and one at/above the threshold drains the shard into the base table. + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = RolloutStore::open_with_options( + &uri, + RolloutStoreOptions { + storage_options: None, + shard_id: Some("rollout-0".to_string()), + merge_after_generations: None, // count trigger off + cleanup_min_generations: Some(2), + ..Default::default() + }, + ) + .await + .unwrap(); + + store.add(&[assistant_record("a-0")]).await.unwrap(); + // One generation pending, below min (2): cleanup is a no-op. + assert_eq!(store.cleanup_own_shard().await.unwrap(), 0); + assert_eq!(flushed_generation_count(&store).await, 1); + + store.add(&[assistant_record("a-1")]).await.unwrap(); + // Two generations pending, at min: cleanup merges both. + assert_eq!(store.cleanup_own_shard().await.unwrap(), 2); + assert_eq!(flushed_generation_count(&store).await, 0); + + // Rows survive the merge, readable exactly once. + let listed = store.list(None, None).await.unwrap(); + assert_eq!(listed.len(), 2); + + // Nothing pending: a further pass reclaims nothing. + assert_eq!(store.cleanup_own_shard().await.unwrap(), 0); + }); + } + + #[test] + fn spawn_periodic_cleanup_reclaims_on_a_timer() { + // The background timer folds accumulated generations into the base table + // without any count-triggered append. A short interval lets the test + // observe one cleanup pass drain the shard. + use std::time::Duration; + use tokio::sync::RwLock; + + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap(); + runtime.block_on(async { + let store = RolloutStore::open_with_options( + &uri, + RolloutStoreOptions { + storage_options: None, + shard_id: Some("rollout-0".to_string()), + merge_after_generations: None, // only the timer merges + cleanup_interval_secs: Some(1), + cleanup_min_generations: Some(1), + ..Default::default() + }, + ) + .await + .unwrap(); + let store = Arc::new(RwLock::new(store)); + + // Accumulate generations that no count trigger will reclaim. + { + let mut guard = store.write().await; + guard.add(&[assistant_record("a-0")]).await.unwrap(); + guard.add(&[assistant_record("a-1")]).await.unwrap(); + assert_eq!(flushed_generation_count(&guard).await, 2); + } + + let handle = RolloutStore::spawn_periodic_cleanup(store.clone()) + .expect("timer enabled"); + + // Wait for at least one tick (interval 1s, first tick skipped) to run + // the cleanup pass and drain the shard. + let drained = async { + loop { + { + let guard = store.read().await; + if flushed_generation_count(&guard).await == 0 { + return true; + } + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + }; + let ok = tokio::time::timeout(Duration::from_secs(10), drained) + .await + .unwrap_or(false); + handle.abort(); + assert!(ok, "periodic cleanup did not drain the shard in time"); + + let guard = store.read().await; + assert_eq!(guard.list(None, None).await.unwrap().len(), 2); + }); + } + + #[test] + fn concurrent_merges_from_two_shards_into_one_base_table() { + // Case 2 of the concurrency hardening: two instances (distinct shards) + // fold their flushed generations into the SAME base table at the same + // time. Lance's optimistic concurrency retries the second append on the + // latest version (Append vs Append is non-conflicting), so no commit is + // lost. Each instance drains only its own shard manifest. Afterwards + // every row is present exactly once and both shards are empty. + use tokio::sync::RwLock; + + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build() + .unwrap(); + runtime.block_on(async { + let make = |shard: &str| { + let uri = uri.clone(); + let shard = shard.to_string(); + async move { + RolloutStore::open_with_options( + &uri, + RolloutStoreOptions { + storage_options: None, + shard_id: Some(shard), + cleanup_min_generations: Some(1), + ..Default::default() + }, + ) + .await + .unwrap() + } + }; + + // Two instances, each accumulates a few generations on its own shard. + let mut a = make("rollout-0").await; + let mut b = make("rollout-1").await; + for i in 0..4 { + a.add(&[assistant_record(&format!("a-{i}"))]).await.unwrap(); + b.add(&[assistant_record(&format!("b-{i}"))]).await.unwrap(); + } + assert_eq!(flushed_generation_count(&a).await, 4); + assert_eq!(flushed_generation_count(&b).await, 4); + + let a = Arc::new(RwLock::new(a)); + let b = Arc::new(RwLock::new(b)); + + // Both merge into the shared base table concurrently. + let (ra, rb) = tokio::join!( + async { + let mut g = a.write().await; + g.cleanup_own_shard().await + }, + async { + let mut g = b.write().await; + g.cleanup_own_shard().await + }, + ); + assert_eq!(ra.unwrap(), 4); + assert_eq!(rb.unwrap(), 4); + + // Both shards drained; a fresh reader sees all 8 rows exactly once. + assert_eq!(flushed_generation_count(&*a.read().await).await, 0); + assert_eq!(flushed_generation_count(&*b.read().await).await, 0); + + let reader = RolloutStore::open(&uri).await.unwrap(); + let listed = reader.list(None, None).await.unwrap(); + let mut ids: Vec<_> = listed.iter().map(|r| r.id.clone()).collect(); + ids.sort(); + let expected: Vec = (0..4) + .flat_map(|i| [format!("a-{i}"), format!("b-{i}")]) + .collect::>() + .into_iter() + .collect(); + assert_eq!(ids, expected); + assert_eq!(listed.len(), 8); + }); + } + + #[test] + fn merge_drains_only_merged_generations() { + // Regression guard for the surgical drain: the manifest commit removes + // exactly the generations that were merged, so if a generation were to + // appear that wasn't part of this merge it would be retained, not + // silently discarded. We assert the row count is conserved end to end: + // every appended row is readable exactly once after a merge, proving no + // generation was dropped without being folded into the base table. + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = RolloutStore::open_with_options( + &uri, + RolloutStoreOptions { + storage_options: None, + shard_id: Some("rollout-0".to_string()), + cleanup_min_generations: Some(1), + ..Default::default() + }, + ) + .await + .unwrap(); + + for i in 0..3 { + store.add(&[assistant_record(&format!("g-{i}"))]).await.unwrap(); + } + assert_eq!(flushed_generation_count(&store).await, 3); + + // Merge all three. Surgical drain removes exactly generations {merged}. + assert_eq!(store.cleanup_own_shard().await.unwrap(), 3); + assert_eq!(flushed_generation_count(&store).await, 0); + + // Append a fourth AFTER the drain: it forms a new generation that the + // prior merge must not have wiped. A second merge folds just that one. + store.add(&[assistant_record("g-3")]).await.unwrap(); + assert_eq!(flushed_generation_count(&store).await, 1); + assert_eq!(store.cleanup_own_shard().await.unwrap(), 1); + + // All four rows conserved, each exactly once. + let listed = store.list(None, None).await.unwrap(); + let mut ids: Vec<_> = listed.iter().map(|r| r.id.clone()).collect(); + ids.sort(); + assert_eq!(ids, vec!["g-0", "g-1", "g-2", "g-3"]); + }); + } + + #[test] + fn spawn_periodic_cleanup_disabled_returns_none() { + // With `cleanup_interval_secs` unset (0), no timer task is spawned. + use tokio::sync::RwLock; + + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let store = RolloutStore::open_with_options( + &uri, + RolloutStoreOptions { + storage_options: None, + shard_id: Some("rollout-0".to_string()), + ..Default::default() + }, + ) + .await + .unwrap(); + let store = Arc::new(RwLock::new(store)); + assert!(RolloutStore::spawn_periodic_cleanup(store).is_none()); + }); + } + /// Micro-benchmark (run with `cargo test -- --ignored --nocapture /// bench_merge_read_amplification`): quantifies the read-amplification that /// self-merge removes. Appends N generations, times a `list` scan over the @@ -1511,6 +1960,7 @@ mod tests { storage_options: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: None, // disabled + ..Default::default() }, ) .await @@ -1540,6 +1990,7 @@ mod tests { storage_options: None, shard_id: Some("rollout-0".to_string()), merge_after_generations: Some(1), + ..Default::default() }, ) .await diff --git a/crates/lance-context-server/src/config.rs b/crates/lance-context-server/src/config.rs index 431c46d5..db618633 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -33,6 +33,23 @@ pub struct ServerConfig { /// unioned at read time. #[arg(long, env = "ROLLOUT_MERGE_AFTER_GENERATIONS", default_value = "0")] pub rollout_merge_after_generations: usize, + + /// Interval, in seconds, for the periodic per-shard WAL cleanup task. When + /// non-zero, each rollout store spawns a background timer that folds this + /// instance's flushed MemWAL generations into the base table on a schedule — + /// the *time* half of the "time + count" trigger, complementing + /// `--rollout-merge-after-generations`. It reclaims stale generations even on + /// low-traffic shards that never cross the count threshold. `0` (the default) + /// disables the timer. + #[arg(long, env = "ROLLOUT_CLEANUP_INTERVAL_SECS", default_value = "0")] + pub rollout_cleanup_interval_secs: u64, + + /// Minimum flushed generations a periodic cleanup tick requires before it + /// merges (avoids rewriting the base table to reclaim a single small + /// generation). Only meaningful when `--rollout-cleanup-interval-secs` is + /// set. Defaults to `1` (clean up whenever any generation is present). + #[arg(long, env = "ROLLOUT_CLEANUP_MIN_GENERATIONS", default_value = "1")] + pub rollout_cleanup_min_generations: usize, } impl ServerConfig { diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index f4709a5a..7fcd20f9 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -42,6 +42,9 @@ 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), + cleanup_interval_secs: (state.rollout_cleanup_interval_secs > 0) + .then_some(state.rollout_cleanup_interval_secs), + cleanup_min_generations: Some(state.rollout_cleanup_min_generations), }; let store = RolloutStore::open_with_options(&uri, options) @@ -49,8 +52,14 @@ pub async fn create_rollout_store( .map_err(AppError::from_lance)?; let version = store.version(); + let store = Arc::new(RwLock::new(store)); + // Start the periodic per-shard WAL-cleanup timer (no-op when the interval is + // disabled). The handle is detached: it is aborted when the store is dropped + // on delete, and otherwise runs for the server's lifetime. + let _cleanup = RolloutStore::spawn_periodic_cleanup(store.clone()); + let mut stores = state.rollout_stores.write().await; - stores.insert(req.name.clone(), Arc::new(RwLock::new(store))); + stores.insert(req.name.clone(), store); Ok(( StatusCode::CREATED, diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index cd3109eb..aec854d6 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -18,6 +18,12 @@ 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, + /// Periodic per-shard WAL-cleanup interval in seconds; `0` disables the + /// background timer. See `RolloutStoreOptions::cleanup_interval_secs`. + pub rollout_cleanup_interval_secs: u64, + /// Minimum flushed generations before a periodic cleanup tick merges. See + /// `RolloutStoreOptions::cleanup_min_generations`. + pub rollout_cleanup_min_generations: usize, } impl AppState { @@ -29,6 +35,8 @@ impl AppState { base_path: PathBuf::from(&config.data_dir), instance_id, rollout_merge_after_generations: config.rollout_merge_after_generations, + rollout_cleanup_interval_secs: config.rollout_cleanup_interval_secs, + rollout_cleanup_min_generations: config.rollout_cleanup_min_generations, } } diff --git a/crates/lance-context/src/unified_rollout.rs b/crates/lance-context/src/unified_rollout.rs index 7992fe1b..06717f8b 100644 --- a/crates/lance-context/src/unified_rollout.rs +++ b/crates/lance-context/src/unified_rollout.rs @@ -38,6 +38,11 @@ impl RolloutStore { // caller that wants count-triggered self-merge opens the core // `RolloutStore` directly with `merge_after_generations` set. merge_after_generations: None, + // Embedded use has no long-lived background runtime to own a cleanup + // timer; leave the periodic WAL cleanup disabled here. A caller that + // wants it opens the core `RolloutStore` directly and drives + // `spawn_periodic_cleanup`. + ..Default::default() }; let store = LocalStore::open_with_options(uri, options) .await From 2c9f7ef39b63e71f494da09de5509ac9be1ec7ec Mon Sep 17 00:00:00 2001 From: Beinan Date: Mon, 13 Jul 2026 17:34:36 +0000 Subject: [PATCH 2/2] fix(ci): rustfmt, clippy, and server test AppState literals - rustfmt the new rollout cleanup code - avoid clippy unnecessary_to_owned / needless_update - add rollout_cleanup_* fields to server test AppState literals Co-Authored-By: Claude Opus 4 --- .../lance-context-core/src/rollout_store.rs | 20 +++++++++---------- .../src/routes/records.rs | 2 ++ .../src/routes/rollouts.rs | 2 ++ .../lance-context-server/src/routes/search.rs | 2 ++ 4 files changed, 15 insertions(+), 11 deletions(-) diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 74b958bc..9cbef2b5 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -334,8 +334,7 @@ impl RolloutStore { let weak = Arc::downgrade(&store); Some(tokio::spawn(async move { - let mut ticker = - tokio::time::interval(std::time::Duration::from_secs(interval_secs)); + let mut ticker = tokio::time::interval(std::time::Duration::from_secs(interval_secs)); // Skip the immediate first tick so we don't merge the instant the // task starts; wait a full interval before the first pass. ticker.tick().await; @@ -510,11 +509,8 @@ impl RolloutStore { Err(err) => { // A concurrent first-writer may have created the index between // our check and our commit. Reload and accept it if so. - self.dataset = Self::load_with_options( - &self.dataset.uri().to_string(), - self.storage_options.clone(), - ) - .await?; + let uri = self.dataset.uri().to_string(); + self.dataset = Self::load_with_options(&uri, self.storage_options.clone()).await?; if self.mem_wal_index_present().await? { Ok(()) } else { @@ -1741,7 +1737,6 @@ mod tests { merge_after_generations: None, // only the timer merges cleanup_interval_secs: Some(1), cleanup_min_generations: Some(1), - ..Default::default() }, ) .await @@ -1756,8 +1751,8 @@ mod tests { assert_eq!(flushed_generation_count(&guard).await, 2); } - let handle = RolloutStore::spawn_periodic_cleanup(store.clone()) - .expect("timer enabled"); + let handle = + RolloutStore::spawn_periodic_cleanup(store.clone()).expect("timer enabled"); // Wait for at least one tick (interval 1s, first tick skipped) to run // the cleanup pass and drain the shard. @@ -1889,7 +1884,10 @@ mod tests { .unwrap(); for i in 0..3 { - store.add(&[assistant_record(&format!("g-{i}"))]).await.unwrap(); + store + .add(&[assistant_record(&format!("g-{i}"))]) + .await + .unwrap(); } assert_eq!(flushed_generation_count(&store).await, 3); diff --git a/crates/lance-context-server/src/routes/records.rs b/crates/lance-context-server/src/routes/records.rs index 6f23262e..c20a9ac8 100644 --- a/crates/lance-context-server/src/routes/records.rs +++ b/crates/lance-context-server/src/routes/records.rs @@ -626,6 +626,8 @@ mod tests { base_path: dir.path().to_path_buf(), instance_id: None, rollout_merge_after_generations: 0, + rollout_cleanup_interval_secs: 0, + rollout_cleanup_min_generations: 1, }); (state, dir) } diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index 7fcd20f9..19254cbc 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -519,6 +519,8 @@ mod tests { base_path: dir.path().to_path_buf(), instance_id: None, rollout_merge_after_generations: 0, + rollout_cleanup_interval_secs: 0, + rollout_cleanup_min_generations: 1, }); let (_status, _info) = create_rollout_store( State(state.clone()), diff --git a/crates/lance-context-server/src/routes/search.rs b/crates/lance-context-server/src/routes/search.rs index a64a0f58..2c19bd31 100644 --- a/crates/lance-context-server/src/routes/search.rs +++ b/crates/lance-context-server/src/routes/search.rs @@ -159,6 +159,8 @@ mod tests { base_path: dir.path().to_path_buf(), instance_id: None, rollout_merge_after_generations: 0, + rollout_cleanup_interval_secs: 0, + rollout_cleanup_min_generations: 1, }); (state, dir) }