DB: prune non-persisted diffs from prev block on every new block commit

crates/apps/src/lib/node/ledger/storage/rocksdb.rs

@@ -1556,6 +1556,39 @@ impl DB for RocksDB {
         Ok(())
     }
 
+    fn prune_non_persisted_diffs(
+        &mut self,
+        batch: &mut Self::WriteBatch,
+        height: BlockHeight,
+    ) -> Result<()> {
+        let rollback_cf = self.get_column_family(ROLLBACK_CF)?;
+
+        let diff_old_key_prefix = Key {
+            segments: vec![
+                height.to_db_key(),
+                OLD_DIFF_PREFIX.to_string().to_db_key(),
+            ],
+        };
+        for (key_str, _val, _) in
+            iter_prefix(self, rollback_cf, None, Some(&diff_old_key_prefix))
+        {
+            batch.0.delete_cf(rollback_cf, key_str)
+        }
+
+        let diff_new_key_prefix = Key {
+            segments: vec![
+                height.to_db_key(),
+                NEW_DIFF_PREFIX.to_string().to_db_key(),
+            ],
+        };
+        for (key_str, _val, _) in
+            iter_prefix(self, rollback_cf, None, Some(&diff_new_key_prefix))
+        {
+            batch.0.delete_cf(rollback_cf, key_str)
+        }
+        Ok(())
+    }
+
     #[inline]
     fn overwrite_entry(
         &self,

crates/state/src/wl_state.rs

@@ -598,6 +598,10 @@ where
             // prune old merkle tree stores
             self.prune_merkle_tree_stores(&mut batch)?;
         }
+        // If there's a previous block, prune non-persisted diffs from it
+        if let Some(height) = self.in_mem.block.height.checked_prev() {
+            self.db.prune_non_persisted_diffs(&mut batch, height)?;
+        }
         self.db.exec_batch(batch)?;
         Ok(())
     }

crates/storage/src/db.rs

@@ -268,6 +268,13 @@ pub trait DB: Debug {
         batch: &mut Self::WriteBatch,
     ) -> Result<()>;
 
+    /// Prune non-persisted diffs that are only kept for one block for rollback
+    fn prune_non_persisted_diffs(
+        &mut self,
+        batch: &mut Self::WriteBatch,
+        height: BlockHeight,
+    ) -> Result<()>;
+
     /// Overwrite a new value in storage, taking into
     /// account values stored at a previous height
     fn overwrite_entry(

crates/storage/src/mockdb.rs

@@ -526,7 +526,8 @@ impl DB for MockDB {
         let diff_prefix = Key::from(height.to_db_key());
         let mut db = self.0.borrow_mut();
 
-        // Diffs
+        // Diffs - Note that this is different from RocksDB that has a separate
+        // CF for non-persisted diffs (ROLLBACK_CF)
         let size_diff =
             match db.insert(subspace_key.to_string(), value.to_owned()) {
                 Some(prev_value) => {
@@ -585,6 +586,8 @@ impl DB for MockDB {
         let diff_prefix = Key::from(height.to_db_key());
         let mut db = self.0.borrow_mut();
 
+        // Diffs - Note that this is different from RocksDB that has a separate
+        // CF for non-persisted diffs (ROLLBACK_CF)
         let size_diff = match db.remove(&subspace_key.to_string()) {
             Some(value) => {
                 let old_key = diff_prefix
@@ -691,6 +694,16 @@ impl DB for MockDB {
         Ok(())
     }
 
+    fn prune_non_persisted_diffs(
+        &mut self,
+        _batch: &mut Self::WriteBatch,
+        _height: BlockHeight,
+    ) -> Result<()> {
+        // No-op - Note that this is different from RocksDB that has a separate
+        // CF for non-persisted diffs (ROLLBACK_CF)
+        Ok(())
+    }
+
     fn overwrite_entry(
         &self,
         _batch: &mut Self::WriteBatch,