From 5e4c102f9be994bcc00207adb3c0c92e6e54188b Mon Sep 17 00:00:00 2001 From: Beinan Date: Wed, 15 Jul 2026 08:04:41 +0000 Subject: [PATCH] fix(core): delete merged MemWAL generation dirs after WAL merge MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit merge_own_shard appended each flushed generation into the base table and drained it from the shard manifest, but never deleted the underlying _mem_wal/{shard}/{gen}/ blob directory. Every merged generation therefore left a zombie directory (data.lance + bloom_filter + manifest) on object storage forever — a pure storage leak that grows without bound on any long-running store. Data stayed correct (the drained manifest no longer references the directory, so reads never see it), which is why this went unnoticed: only _mem_wal/ disk usage was affected. Delete each merged generation's directory after the manifest drain (best-effort, warn-on-failure). Ordering is deliberate: the drain removes the ids first so a reader can no longer resolve them, then the data is deleted — crash-safe in either order. Adds a serial (no concurrency, no fence) regression test asserting on-disk generation dirs match the manifest's pending count after count-triggered merges. Verified it fails without the fix (30 leaked dirs) and passes with it. Co-Authored-By: Claude Opus 4 --- .../lance-context-core/src/rollout_store.rs | 38 ++++++ .../tests/wal_merge_generation_cleanup.rs | 121 ++++++++++++++++++ 2 files changed, 159 insertions(+) create mode 100644 crates/lance-context-core/tests/wal_merge_generation_cleanup.rs diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 1791fd9..e59d986 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -558,6 +558,9 @@ impl RolloutStore { // 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(); + // Remember each merged generation's on-storage folder name so we can + // delete the blob directory after the manifest drain (see below). + let mut merged_paths: Vec = Vec::new(); let mut batches: Vec = Vec::new(); for flushed in &manifest.flushed_generations { let gen_uri = format!( @@ -573,6 +576,7 @@ impl RolloutStore { } } merged_generations.insert(flushed.generation); + merged_paths.push(flushed.path.clone()); } // Append the merged rows to the base table. @@ -617,6 +621,40 @@ impl RolloutStore { }) .await?; + // Delete the merged generations' blob directories now that no manifest + // references them. Ordering matters: the drain above already removed + // these ids from `flushed_generations`, so a reader can no longer resolve + // them — deleting the data second (never before) keeps the sequence + // crash-safe. If the process dies between the drain and here, the rows + // are already in the base table and the manifest no longer lists these + // generations, so nothing reads them; they simply become storage that a + // sweep can reclaim later. + // + // Best-effort: a delete failure must NOT fail the merge — the merge has + // logically succeeded (data appended, manifest drained). A failed delete + // only leaks one directory, which the same reclamation path handles. + // Skipping this deletion is exactly the historical storage leak: every + // merged generation left its `_mem_wal/{shard}/{gen}/` directory behind + // forever. + let object_store = self.dataset.object_store(None).await?; + let branch_path = self.dataset.branch_location().path.clone(); + for path in &merged_paths { + let gen_dir = branch_path + .clone() + .join("_mem_wal") + .join(self.write_shard.to_string().as_str()) + .join(path.as_str()); + if let Err(err) = object_store.remove_dir_all(gen_dir.clone()).await { + tracing::warn!( + shard = %self.write_shard, + generation_path = %path, + error = %err, + "failed to delete merged MemWAL generation directory; \ + it will remain until reclaimed" + ); + } + } + Ok(()) } diff --git a/crates/lance-context-core/tests/wal_merge_generation_cleanup.rs b/crates/lance-context-core/tests/wal_merge_generation_cleanup.rs new file mode 100644 index 0000000..f6410c0 --- /dev/null +++ b/crates/lance-context-core/tests/wal_merge_generation_cleanup.rs @@ -0,0 +1,121 @@ +//! Regression test: WAL merge must delete the merged generation's blob dir. +//! +//! One store, one shard, fully serial — no concurrency, no fence. This isolates +//! a single property: `merge_own_shard` appends merged generations into the base +//! table and drains them from the shard manifest, and must ALSO delete the +//! `_mem_wal/{shard}/{gen}/` blob directory. Historically it did not, leaking one +//! zombie directory per merged generation (a pure storage leak: data was correct +//! but `_mem_wal/` grew without bound). After the fix, the count of on-disk +//! `_gen_` directories tracks the manifest's pending generations exactly. + +use std::path::Path; + +use lance_context_core::{RolloutRecord, RolloutStore, RolloutStoreOptions, ROLE_ASSISTANT}; + +fn rec(id: &str) -> RolloutRecord { + RolloutRecord { + id: id.to_string(), + rollout_id: "r".to_string(), + problem_id: "p".to_string(), + dataset: Some("d".to_string()), + sequence_order: 0, + role: ROLE_ASSISTANT.to_string(), + created_at: chrono::Utc::now(), + content: Some("x".to_string()), + content_type: "text/plain".to_string(), + input_tokens: None, + output_tokens: None, + num_input_tokens: None, + num_output_tokens: None, + output_logprobs: None, + input_logprobs: None, + ref_logprobs: None, + loss_mask: None, + advantage: None, + reward: None, + raw_reward: None, + grader_id: None, + score: None, + include_in_training: None, + exclude_reason: None, + policy_version: None, + relationships: vec![], + binary_payload: None, + payload_size: None, + payload_checksum: None, + artifact_type: None, + metadata: None, + } +} + +fn count_gen_dirs_on_disk(dataset_dir: &Path) -> usize { + let mem_wal = dataset_dir.join("_mem_wal"); + let mut count = 0; + if let Ok(shards) = std::fs::read_dir(&mem_wal) { + for shard in shards.flatten() { + if let Ok(entries) = std::fs::read_dir(shard.path()) { + for e in entries.flatten() { + let name = e.file_name().to_string_lossy().to_string(); + if name.contains("_gen_") && e.path().is_dir() { + count += 1; + } + } + } + } + } + count +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn serial_merge_deletes_merged_generation_dirs() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + + // Merge every 10 accumulated generations into the base table. + let opts = RolloutStoreOptions { + shard_id: Some("solo".to_string()), + merge_after_generations: Some(10), + ..Default::default() + }; + + let mut store = RolloutStore::open_with_options(&uri, opts.clone()) + .await + .unwrap(); + + // 30 serial single-row appends. Each append flushes one generation; every + // 10th triggers a count-merge that folds the accumulated generations into + // the base table and drains them from the manifest. + let n = 30; + for i in 0..n { + store.add(&[rec(&format!("row-{i}"))]).await.unwrap(); + } + + let obs = store.observe().await.unwrap(); + let manifest_pending = obs.pending_wal_generations as usize; + let row_count = obs.row_count; + + let on_disk = count_gen_dirs_on_disk(Path::new(&uri)); + + eprintln!( + "appends={n} row_count={row_count} \ + manifest_pending_generations={manifest_pending} \ + gen_dirs_on_disk={on_disk} leaked={}", + on_disk.saturating_sub(manifest_pending) + ); + + // Data correctness: all rows present exactly once. + assert_eq!( + row_count as usize, n, + "all rows must be readable exactly once" + ); + + // The fix: merged generations are drained from the manifest AND their blob + // dirs are deleted, so on-disk gen dirs never exceed the manifest's pending + // count. (Before the fix this was 30 dirs on disk vs 0 pending = 30 leaks.) + assert_eq!( + on_disk, manifest_pending, + "on-disk generation dirs ({on_disk}) must match manifest pending \ + generations ({manifest_pending}); a surplus means merged generations \ + leaked their blob directories" + ); +}