feat(state): Implements reconsider_block method

zebra-state/src/constants.rs

@@ -117,6 +117,12 @@ pub const MAX_FIND_BLOCK_HEADERS_RESULTS: u32 = 160;
 /// These database versions can be recreated from their directly preceding versions.
 pub const RESTORABLE_DB_VERSIONS: [u64; 1] = [26];
 
+/// The maximum number of invalidated blocks per entry.
+///
+/// This limits the memory for each entry to around:
+/// `256 blocks * 2 MB per block = 0.5 GB`
+pub const MAX_INVALIDATED_BLOCKS: usize = 256;
+
 lazy_static! {
     /// Regex that matches the RocksDB error when its lock file is already open.
     pub static ref LOCK_FILE_ERROR: Regex = Regex::new("(lock file).*(temporarily unavailable)|(in use)|(being used by another process)").expect("regex is valid");

zebra-state/src/error.rs

@@ -51,9 +51,17 @@ pub struct CommitSemanticallyVerifiedError(#[from] ValidateContextError);
 #[derive(Debug, Error)]
 pub enum ReconsiderError {
     #[error("Block with hash {0} was not previously invalidated")]
-    NonPreviouslyInvalidatedBlock(block::Hash),
+    MissingInvalidatedBlock(block::Hash),
+
     #[error("Parent chain not found for block {0}")]
     ParentChainNotFound(block::Hash),
+
+    #[error("Invalidated blocks list is empty when it should contain at least one block")]
+    InvalidatedBlocksEmpty,
+
+    #[error("Invalid height {0:?}: parent height would be negative")]
+    InvalidHeight(block::Height),
+
     #[error("{0}")]
     ValidationError(#[from] ValidateContextError),
 }
@@ -63,6 +71,10 @@ pub enum ReconsiderError {
 #[non_exhaustive]
 #[allow(missing_docs)]
 pub enum ValidateContextError {
+    #[error("block hash {block_hash} was previously invalidated")]
+    #[non_exhaustive]
+    BlockPreviouslyInvalidated { block_hash: block::Hash },
+
     #[error("block parent not found in any chain, or not enough blocks in chain")]
     #[non_exhaustive]
     NotReadyToBeCommitted,

zebra-state/src/service/non_finalized_state.rs

@@ -8,15 +8,17 @@ use std::{
     sync::Arc,
 };
 
+use indexmap::IndexMap;
 use zebra_chain::{
-    block::{self, Block, Hash},
+    block::{self, Block, Hash, Height},
     parameters::Network,
     sprout::{self},
     transparent,
+    value_balance::ValueBalance,
 };
 
 use crate::{
-    constants::MAX_NON_FINALIZED_CHAIN_FORKS,
+    constants::{MAX_INVALIDATED_BLOCKS, MAX_NON_FINALIZED_CHAIN_FORKS},
     error::ReconsiderError,
     request::{ContextuallyVerifiedBlock, FinalizableBlock},
     service::{check, finalized_state::ZebraDb},
@@ -49,7 +51,7 @@ pub struct NonFinalizedState {
 
     /// Blocks that have been invalidated in, and removed from, the non finalized
     /// state.
-    invalidated_blocks: HashMap<Hash, Arc<Vec<ContextuallyVerifiedBlock>>>,
+    invalidated_blocks: IndexMap<Height, Arc<Vec<ContextuallyVerifiedBlock>>>,
 
     // Configuration
     //
@@ -235,6 +237,10 @@ impl NonFinalizedState {
             self.insert(side_chain);
         }
 
+        // Remove all invalidated_blocks at or below the finalized height
+        self.invalidated_blocks
+            .retain(|height, _blocks| *height > best_chain_root.height);
+
         self.update_metrics_for_chains();
 
         // Add the treestate to the finalized block.
@@ -296,8 +302,14 @@ impl NonFinalizedState {
             invalidated_blocks
         };
 
-        self.invalidated_blocks
-            .insert(block_hash, Arc::new(invalidated_blocks));
+        self.invalidated_blocks.insert(
+            invalidated_blocks.first().unwrap().clone().height,
+            Arc::new(invalidated_blocks),
+        );
+
+        while self.invalidated_blocks.len() > MAX_INVALIDATED_BLOCKS {
+            self.invalidated_blocks.shift_remove_index(0);
+        }
 
         self.update_metrics_for_chains();
         self.update_metrics_bars();
@@ -308,48 +320,76 @@ impl NonFinalizedState {
     /// into the chain_set.
     pub fn reconsider_block(&mut self, block_hash: block::Hash) -> Result<(), ReconsiderError> {
         // Get the invalidated blocks that were invalidated by the given block_hash
-        let Some(invalidated_blocks) = self.invalidated_blocks.get(&block_hash) else {
-            return Err(ReconsiderError::NonPreviouslyInvalidatedBlock(block_hash));
-        };
+        let height = self
+            .invalidated_blocks
+            .iter()
+            .find_map(|(height, blocks)| {
+                if blocks.first()?.hash == block_hash {
+                    Some(height)
+                } else {
+                    None
+                }
+            })
+            .ok_or(ReconsiderError::MissingInvalidatedBlock(block_hash))?;
 
-        let mut blocks = invalidated_blocks.clone();
-        let mut_blocks = Arc::make_mut(&mut blocks);
-
-        let Some(invalidated_root) = mut_blocks.first() else {
-            return Err(ReconsiderError::NonPreviouslyInvalidatedBlock(block_hash));
-        };
+        let mut invalidated_blocks = self
+            .invalidated_blocks
+            .clone()
+            .shift_remove(height)
+            .ok_or(ReconsiderError::MissingInvalidatedBlock(block_hash))?;
+        let mut_blocks = Arc::make_mut(&mut invalidated_blocks);
 
         // Find and fork the parent chain of the invalidated_root. Update the parent chain
         // with the invalidated_descendants
-        let prev_block_hash = invalidated_root.block.header.previous_block_hash;
-        let Ok(parent_chain) = self.parent_chain(prev_block_hash) else {
-            return Err(ReconsiderError::ParentChainNotFound(block_hash));
-        };
+        let invalidated_root = mut_blocks
+            .first()
+            .ok_or(ReconsiderError::InvalidatedBlocksEmpty)?;
+
+        let root_parent_hash = invalidated_root.block.header.previous_block_hash;
+        let parent_chain = self
+            .parent_chain(root_parent_hash)
+            .map_err(|_| ReconsiderError::ParentChainNotFound(block_hash));
+
+        let modified_chain = match parent_chain {
+            // 1. If a parent chain exist use the parent chain
+            Ok(parent_chain) => {
+                let mut chain = Arc::unwrap_or_clone(parent_chain);
+
+                for block in Arc::unwrap_or_clone(invalidated_blocks) {
+                    chain = chain.push(block)?;
+                }
 
-        let Some(new_chain) = parent_chain.fork(parent_chain.non_finalized_tip_hash()) else {
-            return Err(ReconsiderError::ValidationError(
-                ValidateContextError::NonSequentialBlock {
-                    candidate_height: invalidated_root.height,
-                    parent_height: parent_chain.non_finalized_tip_height(),
-                },
-            ));
+                Arc::new(chain)
+            }
+            // 2. If a parent chain does not exist create a new chain from the non finalized state tip
+            Err(_) => {
+                let mut chain = Chain::new(
+                    &self.network,
+                    invalidated_root.height.previous().map_err(|_| {
+                        ReconsiderError::InvalidHeight(Height(invalidated_root.height.0 - 1))
+                    })?,
+                    Default::default(),
+                    Default::default(),
+                    Default::default(),
+                    Default::default(),
+                    ValueBalance::zero(),
+                );
+
+                for block in Arc::unwrap_or_clone(invalidated_blocks) {
+                    chain = chain.push(block)?;
+                }
+
+                Arc::new(chain)
+            }
         };
 
-        let mut chain = new_chain
-            .push(mut_blocks.remove(0))
-            .map_err(ReconsiderError::ValidationError)?;
-        for block in mut_blocks {
-            chain = chain
-                .push(block.clone())
-                .map_err(ReconsiderError::ValidationError)?;
-        }
+        let (height, hash) = modified_chain.non_finalized_tip();
 
-        self.insert_with(Arc::new(chain), |chain_set| {
-            chain_set.retain(|c| c != &parent_chain)
+        self.insert_with(modified_chain, |chain_set| {
+            chain_set.retain(|chain| chain.non_finalized_tip_hash() != root_parent_hash)
         });
 
-        self.update_metrics_for_chains();
-        self.update_metrics_bars();
+        self.update_metrics_for_committed_block(height, hash);
 
         Ok(())
     }
@@ -405,6 +445,12 @@ impl NonFinalizedState {
         prepared: SemanticallyVerifiedBlock,
         finalized_state: &ZebraDb,
     ) -> Result<Arc<Chain>, ValidateContextError> {
+        if self.invalidated_blocks.contains_key(&prepared.height) {
+            return Err(ValidateContextError::BlockPreviouslyInvalidated {
+                block_hash: prepared.hash,
+            });
+        }
+
         // Reads from disk
         //
         // TODO: if these disk reads show up in profiles, run them in parallel, using std::thread::spawn()
@@ -677,7 +723,7 @@ impl NonFinalizedState {
     }
 
     /// Return the invalidated blocks.
-    pub fn invalidated_blocks(&self) -> HashMap<block::Hash, Arc<Vec<ContextuallyVerifiedBlock>>> {
+    pub fn invalidated_blocks(&self) -> IndexMap<Height, Arc<Vec<ContextuallyVerifiedBlock>>> {
         self.invalidated_blocks.clone()
     }
 

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

@@ -278,26 +278,36 @@ fn invalidate_block_removes_block_and_descendants_from_chain_for_network(
     );
 
     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();
+    // Find an entry in the IndexMap that contains block2 hash
+    let (_, invalidated_blocks_state_descendants) = invalidated_blocks_state
+        .iter()
+        .find_map(|(height, blocks)| {
+            assert!(
+                blocks.iter().any(|block| block.hash == block2.hash()),
+                "invalidated_blocks should reference the hash of block2"
+            );
+
+            if blocks.iter().any(|block| block.hash == block2.hash()) {
+                Some((height, blocks))
+            } else {
+                None
+            }
+        })
+        .unwrap();
 
     match network {
         Network::Mainnet => assert!(
             invalidated_blocks_state_descendants
                 .iter()
                 .any(|block| block.height == block::Height(653601)),
-            "invalidated descendants vec should contain block3"
+            "invalidated descendants should contain block3"
         ),
         Network::Testnet(_parameters) => assert!(
             invalidated_blocks_state_descendants
                 .iter()
                 .any(|block| block.height == block::Height(584001)),
-            "invalidated descendants vec should contain block3"
+            "invalidated descendants should contain block3"
         ),
     }