From 87ff44b32e0a9800450c5271989843c110700fb0 Mon Sep 17 00:00:00 2001
From: Harper <harper.jme@gmail.com>
Date: Fri, 8 Mar 2024 23:31:31 +0000
Subject: [PATCH] feat(rolldb): allow optionally overlap of WAL over immutable
 chain (#419)

---
 pallas-rolldb/src/wal/store.rs  | 10 +++++++--
 pallas-rolldb/src/wal/stream.rs |  2 +-
 pallas-rolldb/src/wal/tests.rs  | 36 ++++++++++++++++++++++++++++++++-
 3 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/pallas-rolldb/src/wal/store.rs b/pallas-rolldb/src/wal/store.rs
index 1275a741..715e738b 100644
--- a/pallas-rolldb/src/wal/store.rs
+++ b/pallas-rolldb/src/wal/store.rs
@@ -167,10 +167,15 @@ pub struct Store {
     pub tip_change: Arc<tokio::sync::Notify>,
     wal_seq: u64,
     k_param: u64,
+    immutable_overlap: u64,
 }
 
 impl Store {
-    pub fn open(path: impl AsRef<Path>, k_param: u64) -> Result<Self, Error> {
+    pub fn open(
+        path: impl AsRef<Path>,
+        k_param: u64,
+        immutable_overlap: Option<u64>,
+    ) -> Result<Self, Error> {
         let mut opts = Options::default();
         opts.create_if_missing(true);
         opts.create_missing_column_families(true);
@@ -184,6 +189,7 @@ impl Store {
             tip_change: Arc::new(tokio::sync::Notify::new()),
             wal_seq,
             k_param,
+            immutable_overlap: immutable_overlap.unwrap_or(0),
         };
 
         Ok(out)
@@ -378,7 +384,7 @@ impl Store {
             // get the number of slots that have passed since the wal point
             let slot_delta = tip - value.slot().unwrap_or(0);
 
-            if slot_delta <= self.k_param {
+            if slot_delta <= self.k_param + self.immutable_overlap {
                 break;
             } else {
                 WalKV::stage_delete(&self.db, wal_key, &mut batch);
diff --git a/pallas-rolldb/src/wal/stream.rs b/pallas-rolldb/src/wal/stream.rs
index 69c867d0..9b95a026 100644
--- a/pallas-rolldb/src/wal/stream.rs
+++ b/pallas-rolldb/src/wal/stream.rs
@@ -46,7 +46,7 @@ mod tests {
     #[tokio::test]
     async fn test_stream_waiting() {
         let path = tempfile::tempdir().unwrap().into_path();
-        let mut db = Store::open(path.clone(), 30).unwrap();
+        let mut db = Store::open(path.clone(), 30, None).unwrap();
 
         for i in 0..=100 {
             let (slot, hash, body) = dummy_block(i * 10);
diff --git a/pallas-rolldb/src/wal/tests.rs b/pallas-rolldb/src/wal/tests.rs
index df64d8c1..f10637c4 100644
--- a/pallas-rolldb/src/wal/tests.rs
+++ b/pallas-rolldb/src/wal/tests.rs
@@ -2,7 +2,16 @@ use super::{BlockBody, BlockHash, BlockSlot, Store};
 
 fn with_tmp_db<T>(k_param: u64, op: fn(store: Store) -> T) {
     let path = tempfile::tempdir().unwrap().into_path();
-    let store = Store::open(path.clone(), k_param).unwrap();
+    let store = Store::open(path.clone(), k_param, None).unwrap();
+
+    op(store);
+
+    Store::destroy(path).unwrap();
+}
+
+fn with_tmp_db_overlap<T>(k_param: u64, overlap: u64, op: fn(store: Store) -> T) {
+    let path = tempfile::tempdir().unwrap().into_path();
+    let store = Store::open(path.clone(), k_param, Some(overlap)).unwrap();
 
     op(store);
 
@@ -109,6 +118,8 @@ fn test_prune_linear() {
             db.roll_forward(slot, hash, body).unwrap();
         }
 
+        // db contains slots: 0, 10, ..., 980, 990
+        // this should prune slots less than (990 - 30) = 960
         db.prune_wal().unwrap();
 
         let mut wal = db.crawl_after(None);
@@ -122,6 +133,29 @@ fn test_prune_linear() {
     });
 }
 
+#[test]
+fn test_prune_linear_with_overlap() {
+    with_tmp_db_overlap(30, 20, |mut db| {
+        for i in 0..100 {
+            let (slot, hash, body) = dummy_block(i * 10);
+            db.roll_forward(slot, hash, body).unwrap();
+        }
+
+        // db contains slots: 0, 10, ..., 980, 990
+        // this should prune slots less than (990 - 30 - 20) = 940
+        db.prune_wal().unwrap();
+
+        let mut wal = db.crawl_after(None);
+
+        for i in 94..100 {
+            let (_, val) = wal.next().unwrap().unwrap();
+            assert_eq!(val.slot().unwrap(), i * 10);
+        }
+
+        assert!(wal.next().is_none());
+    });
+}
+
 #[test]
 fn test_prune_with_rollback() {
     with_tmp_db(30, |mut db| {