fix(core): delete merged MemWAL generation dirs after WAL merge

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<u64> = 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<String> = Vec::new();
         let mut batches: Vec<RecordBatch> = 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(())
     }
 

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"
+    );
+}