Add invalidate block method and invalidated_blocks field to NonFinalizedState

zebra-state/src/request.rs

@@ -1,9 +1,10 @@
 //! State [`tower::Service`] request types.
 
 use std::{
-    collections::{HashMap, HashSet},
+    collections::{BTreeMap, HashMap, HashSet},
     ops::{Deref, DerefMut, RangeInclusive},
     sync::Arc,
+    time::SystemTime,
 };
 
 use zebra_chain::{
@@ -234,6 +235,18 @@ pub struct ContextuallyVerifiedBlock {
     pub(crate) chain_value_pool_change: ValueBalance<NegativeAllowed>,
 }
 
+/// Data related to an invalidated block including the [`ContextuallyVerifiedBlock`] and
+/// descendants.
+#[derive(Clone, Debug)]
+pub struct InvalidatedBlockData {
+    /// The block that was invalidated.
+    pub block: ContextuallyVerifiedBlock,
+    /// All descendant blocks that were also removed in order by [`block::Height`].
+    pub descendants: BTreeMap<block::Height, ContextuallyVerifiedBlock>,
+    /// Timestamp in which the block was invalidated.
+    pub timestamp: SystemTime,
+}
+
 /// Wraps note commitment trees and the history tree together.
 ///
 /// The default instance represents the treestate that corresponds to the genesis block.

zebra-state/src/service/non_finalized_state.rs

@@ -6,17 +6,19 @@ use std::{
     collections::{BTreeSet, HashMap},
     mem,
     sync::Arc,
+    time::SystemTime,
 };
 
+use chain::UpdateWith;
 use zebra_chain::{
-    block::{self, Block},
+    block::{self, Block, Hash},
     parameters::Network,
     sprout, transparent,
 };
 
 use crate::{
     constants::MAX_NON_FINALIZED_CHAIN_FORKS,
-    request::{ContextuallyVerifiedBlock, FinalizableBlock},
+    request::{ContextuallyVerifiedBlock, FinalizableBlock, InvalidatedBlockData},
     service::{check, finalized_state::ZebraDb},
     SemanticallyVerifiedBlock, ValidateContextError,
 };
@@ -45,6 +47,9 @@ pub struct NonFinalizedState {
     /// callers should migrate to `chain_iter().next()`.
     chain_set: BTreeSet<Arc<Chain>>,
 
+    /// Invalidated blocks and descendants
+    invalidated_blocks: HashMap<block::Hash, InvalidatedBlockData>,
+
     // Configuration
     //
     /// The configured Zcash network.
@@ -92,6 +97,7 @@ impl Clone for NonFinalizedState {
         Self {
             chain_set: self.chain_set.clone(),
             network: self.network.clone(),
+            invalidated_blocks: self.invalidated_blocks.clone(),
 
             #[cfg(feature = "getblocktemplate-rpcs")]
             should_count_metrics: self.should_count_metrics,
@@ -112,6 +118,7 @@ impl NonFinalizedState {
         NonFinalizedState {
             chain_set: Default::default(),
             network: network.clone(),
+            invalidated_blocks: Default::default(),
             #[cfg(feature = "getblocktemplate-rpcs")]
             should_count_metrics: true,
             #[cfg(feature = "progress-bar")]
@@ -264,6 +271,49 @@ impl NonFinalizedState {
         Ok(())
     }
 
+    /// Invalidate block with has `block_hash` and all descendants from the non-finalized state. Insert
+    /// the new chain into the chain_set and discard the previous.
+    pub fn invalidate_block(&mut self, block_hash: Hash) {
+        if !self.any_chain_contains(&block_hash) {
+            return;
+        }
+
+        let mut chain = self
+            .find_chain(|chain| chain.contains_block_hash(block_hash))
+            .expect("block hash exist in a chain");
+
+        let new_chain = Arc::make_mut(&mut chain);
+        let block_height = new_chain.height_by_hash(block_hash).unwrap();
+
+        // Split the new chain at the the `block_hash` and invalidate the block with the
+        // block_hash along with the block's descendants
+        let mut invalidated_blocks = new_chain.blocks.split_off(&block_height);
+        for (_, ctx_block) in invalidated_blocks.iter().rev() {
+            new_chain.revert_chain_with(ctx_block, chain::RevertPosition::Tip);
+        }
+
+        self.invalidated_blocks.insert(
+            block_hash,
+            InvalidatedBlockData {
+                block: invalidated_blocks.remove(&block_height).unwrap(),
+                descendants: invalidated_blocks,
+                timestamp: SystemTime::now(),
+            },
+        );
+
+        // If the new chain still contains blocks:
+        // - add the new chain fork or updated chain to the set of recent chains
+        // - remove the chain containing the hash of the block from the chain set
+        if !new_chain.is_empty() {
+            self.insert_with(Arc::new(new_chain.clone()), |chain_set| {
+                chain_set.retain(|c| !c.contains_block_hash(block_hash))
+            });
+
+            self.update_metrics_for_chains();
+            self.update_metrics_bars();
+        }
+    }
+
     /// Commit block to the non-finalized state as a new chain where its parent
     /// is the finalized tip.
     #[tracing::instrument(level = "debug", skip(self, finalized_state, prepared))]
@@ -578,6 +628,11 @@ impl NonFinalizedState {
         self.chain_set.len()
     }
 
+    /// Return the invalidated blocks.
+    pub fn invalidated_blocks(&self) -> HashMap<block::Hash, InvalidatedBlockData> {
+        self.invalidated_blocks.clone()
+    }
+
     /// Return the chain whose tip block hash is `parent_hash`.
     ///
     /// The chain can be an existing chain in the non-finalized state, or a freshly

zebra-state/src/service/non_finalized_state/chain.rs

@@ -1565,7 +1565,7 @@ impl DerefMut for Chain {
 
 /// The revert position being performed on a chain.
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
-enum RevertPosition {
+pub(crate) enum RevertPosition {
     /// The chain root is being reverted via [`Chain::pop_root`], when a block
     /// is finalized.
     Root,
@@ -1584,7 +1584,7 @@ enum RevertPosition {
 /// and [`Chain::pop_tip`] functions, and fear that it would be easy to
 /// introduce bugs when updating them, unless the code was reorganized to keep
 /// related operations adjacent to each other.
-trait UpdateWith<T> {
+pub(crate) trait UpdateWith<T> {
     /// When `T` is added to the chain tip,
     /// update [`Chain`] cumulative data members to add data that are derived from `T`.
     fn update_chain_tip_with(&mut self, _: &T) -> Result<(), ValidateContextError>;

zebra-state/src/service/non_finalized_state/tests/vectors.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 
 use zebra_chain::{
     amount::NonNegative,
-    block::{Block, Height},
+    block::{self, Block, Height},
     history_tree::NonEmptyHistoryTree,
     parameters::{Network, NetworkUpgrade},
     serialization::ZcashDeserializeInto,
@@ -216,6 +216,91 @@ fn finalize_pops_from_best_chain_for_network(network: Network) -> Result<()> {
     Ok(())
 }
 
+fn invalidate_block_removes_block_and_descendants_from_chain_for_network(
+    network: Network,
+) -> Result<()> {
+    let block1: Arc<Block> = Arc::new(network.test_block(653599, 583999).unwrap());
+    let block2 = block1.make_fake_child().set_work(10);
+    let block3 = block2.make_fake_child().set_work(1);
+
+    let mut state = NonFinalizedState::new(&network);
+    let finalized_state = FinalizedState::new(
+        &Config::ephemeral(),
+        &network,
+        #[cfg(feature = "elasticsearch")]
+        false,
+    );
+
+    let fake_value_pool = ValueBalance::<NonNegative>::fake_populated_pool();
+    finalized_state.set_finalized_value_pool(fake_value_pool);
+
+    state.commit_new_chain(block1.clone().prepare(), &finalized_state)?;
+    state.commit_block(block2.clone().prepare(), &finalized_state)?;
+    state.commit_block(block3.clone().prepare(), &finalized_state)?;
+
+    assert_eq!(
+        state
+            .best_chain()
+            .unwrap_or(&Arc::new(Chain::default()))
+            .blocks
+            .len(),
+        3
+    );
+
+    state.invalidate_block(block2.hash());
+
+    let post_invalidated_chain = state.best_chain().unwrap();
+
+    assert_eq!(post_invalidated_chain.blocks.len(), 1);
+    assert!(
+        post_invalidated_chain.contains_block_hash(block1.hash()),
+        "the new modified chain should contain block1"
+    );
+
+    assert!(
+        !post_invalidated_chain.contains_block_hash(block2.hash()),
+        "the new modified chain should not contain block2"
+    );
+    assert!(
+        !post_invalidated_chain.contains_block_hash(block3.hash()),
+        "the new modified chain should not contain block3"
+    );
+
+    let invalidated_blocks_state = &state.invalidated_blocks;
+    assert!(
+        invalidated_blocks_state.contains_key(&block2.hash()),
+        "invalidated blocks map should reference the hash of block2"
+    );
+
+    let invalidated_blocks_state_descendants = &invalidated_blocks_state
+        .get(&block2.hash())
+        .unwrap()
+        .descendants;
+    match network {
+        Network::Mainnet => assert!(
+            invalidated_blocks_state_descendants.contains_key(&block::Height(653601)),
+            "invalidated blocks map should reference block3's height in descendants"
+        ),
+        Network::Testnet(_parameters) => assert!(
+            invalidated_blocks_state_descendants.contains_key(&block::Height(584001)),
+            "invalidated blocks map should reference block3's height in descendants"
+        ),
+    }
+
+    Ok(())
+}
+
+#[test]
+fn invalidate_block_removes_block_and_descendants_from_chain() -> Result<()> {
+    let _init_guard = zebra_test::init();
+
+    for network in Network::iter() {
+        invalidate_block_removes_block_and_descendants_from_chain_for_network(network)?;
+    }
+
+    Ok(())
+}
+
 #[test]
 // This test gives full coverage for `take_chain_if`
 fn commit_block_extending_best_chain_doesnt_drop_worst_chains() -> Result<()> {