Merge branch 'tomas/fix-rollback' (#1616)

* origin/tomas/fix-rollback:
  ledger/storage/rocksdb: fix rollback to restore deleted keys
  test/ledger/storage/rocksdb: add a test for block rollback

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

@@ -435,6 +435,29 @@ impl RocksDB {
             },
         )?;
 
+        // Look for diffs in this block to find what has been deleted
+        let diff_new_key_prefix = Key {
+            segments: vec![
+                last_block.height.to_db_key(),
+                "new".to_string().to_db_key(),
+            ],
+        };
+        {
+            let mut batch_guard = batch.lock().unwrap();
+            let subspace_cf = self.get_column_family(SUBSPACE_CF)?;
+            for (key, val, _) in
+                iter_diffs_prefix(self, last_block.height, true)
+            {
+                let key = Key::parse(key).unwrap();
+                let diff_new_key = diff_new_key_prefix.join(&key);
+                if self.read_subspace_val(&diff_new_key)?.is_none() {
+                    // If there is no new value, it has been deleted in this
+                    // block and we have to restore it
+                    batch_guard.put_cf(subspace_cf, key.to_string(), val)
+                }
+            }
+        }
+
         tracing::info!("Deleting keys prepended with the last height");
         let mut batch = batch.into_inner().unwrap();
         let prefix = last_block.height.to_string();
@@ -1443,6 +1466,7 @@ mod test {
     use namada::types::address::EstablishedAddressGen;
     use namada::types::storage::{BlockHash, Epoch, Epochs};
     use tempfile::tempdir;
+    use test_log::test;
 
     use super::*;
 
@@ -1462,36 +1486,7 @@ mod test {
         )
         .unwrap();
 
-        let merkle_tree = MerkleTree::<Sha256Hasher>::default();
-        let merkle_tree_stores = merkle_tree.stores();
-        let hash = BlockHash::default();
-        let time = DateTimeUtc::now();
-        let epoch = Epoch::default();
-        let pred_epochs = Epochs::default();
-        let height = BlockHeight::default();
-        let next_epoch_min_start_height = BlockHeight::default();
-        let next_epoch_min_start_time = DateTimeUtc::now();
-        let update_epoch_blocks_delay = None;
-        let address_gen = EstablishedAddressGen::new("whatever");
-        let tx_queue = TxQueue::default();
-        let results = BlockResults::default();
-        let block = BlockStateWrite {
-            merkle_tree_stores,
-            header: None,
-            hash: &hash,
-            height,
-            time,
-            epoch,
-            results: &results,
-            pred_epochs: &pred_epochs,
-            next_epoch_min_start_height,
-            next_epoch_min_start_time,
-            update_epoch_blocks_delay,
-            address_gen: &address_gen,
-            tx_queue: &tx_queue,
-        };
-
-        db.write_block(block, &mut batch, true).unwrap();
+        write_block(&mut db, &mut batch, BlockHeight::default()).unwrap();
         db.exec_batch(batch.0).unwrap();
 
         let _state = db
@@ -1653,4 +1648,116 @@ mod test {
             .collect();
         itertools::assert_equal(all_keys, itered_keys);
     }
+
+    #[test]
+    fn test_rollback() {
+        let dir = tempdir().unwrap();
+        let mut db = open(dir.path(), None).unwrap();
+
+        // A key that's gonna be added on a second block
+        let add_key = Key::parse("add").unwrap();
+        // A key that's gonna be deleted on a second block
+        let delete_key = Key::parse("delete").unwrap();
+        // A key that's gonna be overwritten on a second block
+        let overwrite_key = Key::parse("overwrite").unwrap();
+
+        // Write first block
+        let mut batch = RocksDB::batch();
+        let height_0 = BlockHeight(100);
+        let to_delete_val = vec![1_u8, 1, 0, 0];
+        let to_overwrite_val = vec![1_u8, 1, 1, 0];
+        db.batch_write_subspace_val(
+            &mut batch,
+            height_0,
+            &delete_key,
+            &to_delete_val,
+        )
+        .unwrap();
+        db.batch_write_subspace_val(
+            &mut batch,
+            height_0,
+            &overwrite_key,
+            &to_overwrite_val,
+        )
+        .unwrap();
+
+        write_block(&mut db, &mut batch, height_0).unwrap();
+        db.exec_batch(batch.0).unwrap();
+
+        // Write second block
+        let mut batch = RocksDB::batch();
+        let height_1 = BlockHeight(101);
+        let add_val = vec![1_u8, 0, 0, 0];
+        let overwrite_val = vec![1_u8, 1, 1, 1];
+        db.batch_write_subspace_val(&mut batch, height_1, &add_key, &add_val)
+            .unwrap();
+        db.batch_write_subspace_val(
+            &mut batch,
+            height_1,
+            &overwrite_key,
+            &overwrite_val,
+        )
+        .unwrap();
+        db.batch_delete_subspace_val(&mut batch, height_1, &delete_key)
+            .unwrap();
+
+        write_block(&mut db, &mut batch, height_1).unwrap();
+        db.exec_batch(batch.0).unwrap();
+
+        // Check that the values are as expected from second block
+        let added = db.read_subspace_val(&add_key).unwrap();
+        assert_eq!(added, Some(add_val));
+        let overwritten = db.read_subspace_val(&overwrite_key).unwrap();
+        assert_eq!(overwritten, Some(overwrite_val));
+        let deleted = db.read_subspace_val(&delete_key).unwrap();
+        assert_eq!(deleted, None);
+
+        // Rollback to the first block height
+        db.rollback(height_0).unwrap();
+
+        // Check that the values are back to the state at the first block
+        let added = db.read_subspace_val(&add_key).unwrap();
+        assert_eq!(added, None);
+        let overwritten = db.read_subspace_val(&overwrite_key).unwrap();
+        assert_eq!(overwritten, Some(to_overwrite_val));
+        let deleted = db.read_subspace_val(&delete_key).unwrap();
+        assert_eq!(deleted, Some(to_delete_val));
+    }
+
+    /// A test helper to write a block
+    fn write_block(
+        db: &mut RocksDB,
+        batch: &mut RocksDBWriteBatch,
+        height: BlockHeight,
+    ) -> Result<()> {
+        let merkle_tree = MerkleTree::<Sha256Hasher>::default();
+        let merkle_tree_stores = merkle_tree.stores();
+        let hash = BlockHash::default();
+        let time = DateTimeUtc::now();
+        let epoch = Epoch::default();
+        let pred_epochs = Epochs::default();
+        let next_epoch_min_start_height = BlockHeight::default();
+        let next_epoch_min_start_time = DateTimeUtc::now();
+        let update_epoch_blocks_delay = None;
+        let address_gen = EstablishedAddressGen::new("whatever");
+        let tx_queue = TxQueue::default();
+        let results = BlockResults::default();
+        let block = BlockStateWrite {
+            merkle_tree_stores,
+            header: None,
+            hash: &hash,
+            height,
+            time,
+            epoch,
+            results: &results,
+            pred_epochs: &pred_epochs,
+            next_epoch_min_start_height,
+            next_epoch_min_start_time,
+            update_epoch_blocks_delay,
+            address_gen: &address_gen,
+            tx_queue: &tx_queue,
+        };
+
+        db.write_block(block, batch, true)
+    }
 }