refactor codes

crates/blockchain-tree/src/blockchain_tree.rs

@@ -20,16 +20,18 @@ use reth_primitives::{
     SealedBlock, SealedBlockWithSenders, SealedHeader, StaticFileSegment, B256, U256,
 };
 use reth_provider::{
-    BlockExecutionWriter, BlockNumReader, BlockWriter, CanonStateNotification,
-    CanonStateNotificationSender, CanonStateNotifications, ChainSpecProvider, ChainSplit,
-    ChainSplitTarget, DisplayBlocksChain, HeaderProvider, ProviderError, StaticFileProviderFactory,
+    BlockExecutionReader, BlockExecutionWriter, BlockNumReader, BlockWriter,
+    CanonStateNotification, CanonStateNotificationSender, CanonStateNotifications,
+    ChainSpecProvider, ChainSplit, ChainSplitTarget, DisplayBlocksChain, HeaderProvider,
+    ProviderError, StaticFileProviderFactory,
 };
 use reth_prune_types::PruneModes;
 use reth_stages_api::{MetricEvent, MetricEventsSender};
 use reth_storage_errors::provider::{ProviderResult, RootMismatch};
 use reth_trie::{hashed_cursor::HashedPostStateCursorFactory, StateRoot};
 use std::{
     collections::{btree_map::Entry, BTreeMap, HashSet},
+    ops::Deref,
     sync::Arc,
 };
 use tracing::{debug, error, info, instrument, trace, warn};
@@ -437,7 +439,50 @@ where
             block_validation_kind,
         )?;
 
-        self.insert_chain(chain);
+        let Some(chain_id) = self.block_indices().get_block_chain_id(&block_num_hash.hash) else {
+            unreachable!("Should not happen as the canonical chain should keep the canonical block")
+        };
+        let Some(prepend) = self.state.chains.get(&chain_id) else {
+            unreachable!("Should not happen as the canonical chain should be kept")
+        };
+
+        let mut chains = vec![chain.into_inner()];
+        let in_memory =
+            match prepend.deref().clone().split(ChainSplitTarget::Hash(block_num_hash.hash)) {
+                ChainSplit::Split { canonical: lower, pending: _higher } => lower,
+                ChainSplit::NoSplitCanonical(lower) => lower,
+                ChainSplit::NoSplitPending(_) => {
+                    unreachable!("Should not happen as block indices guarantee structure of blocks")
+                }
+            };
+        let memory_height = in_memory.first().number;
+        chains.push(in_memory);
+
+        // need to get states which are canonical but not in memory (i.e, during restart)
+        if memory_height > self.block_indices().last_finalized_block() {
+            let provider_ro = self.externals.provider_factory.provider()?;
+
+            let range = (self.block_indices().last_finalized_block() + 1)..=memory_height - 1;
+            // Read block and execution result from database.
+            let in_db = provider_ro.get_block_and_execution_range(range).map_err(|_| {
+                InsertBlockErrorKind::Tree(BlockchainTreeError::CanonicalChain {
+                    block_hash: block_num_hash.hash,
+                })
+            })?;
+            chains.push(in_db);
+        }
+
+        if let Some(mut new_chain) = chains.pop() {
+            for chain in chains.into_iter().rev() {
+                new_chain.append_chain(chain).map_err(|_| {
+                    InsertBlockErrorKind::Tree(BlockchainTreeError::CanonicalChain {
+                        block_hash: block_num_hash.hash,
+                    })
+                })?;
+            }
+            self.insert_chain(AppendableChain::new(new_chain));
+        }
+
         self.try_connect_buffered_blocks(block_num_hash);
 
         Ok(BlockStatus::Valid(block_attachment))
@@ -797,6 +842,27 @@ where
 
     /// Finalize blocks up until and including `finalized_block`, and remove them from the tree.
     pub fn finalize_block(&mut self, finalized_block: BlockNumber) -> ProviderResult<()> {
+        let block_hash = self.externals.find_hash_by_number(finalized_block).unwrap();
+        let Some(chain_id) = self.block_indices().get_block_chain_id(&block_hash) else {
+            unreachable!("Should not happen as the canonical chain should keep the canonical block")
+        };
+
+        let canonical = self.state.chains.remove(&chain_id);
+        let mut outcome = ExecutionOutcome::default();
+        match canonical.unwrap().into_inner().split(ChainSplitTarget::Number(finalized_block)) {
+            ChainSplit::Split { canonical: lower, pending: higher } => {
+                outcome.clone_from(lower.execution_outcome());
+
+                self.state.block_indices.insert_chain(chain_id, &higher);
+                self.state.chains.insert(chain_id, AppendableChain::new(higher));
+            }
+            ChainSplit::NoSplitCanonical(lower) => {
+                outcome.clone_from(lower.execution_outcome());
+            }
+            ChainSplit::NoSplitPending(_) => {}
+        }
+        apply_bundle_state(outcome.bundle);
+
         // remove blocks
         let mut remove_chains = self.state.block_indices.finalize_canonical_blocks(
             finalized_block,
@@ -1067,7 +1133,7 @@ where
         };
 
         // we are splitting chain at the block hash that we want to make canonical
-        let Some(canonical) = self.remove_and_split_chain(chain_id, block_hash.into()) else {
+        let Some(mut canonical) = self.remove_and_split_chain(chain_id, block_hash.into()) else {
             debug!(target: "blockchain_tree", ?block_hash, ?chain_id, "Chain not present");
             return Err(CanonicalError::from(BlockchainTreeError::BlockSideChainIdConsistency {
                 chain_id: chain_id.into(),
@@ -1077,7 +1143,7 @@ where
         durations_recorder.record_relative(MakeCanonicalAction::SplitChain);
 
         let mut fork_block = canonical.fork_block();
-        let mut chains_to_promote = vec![canonical];
+        let mut chains_to_promote = vec![canonical.clone()];
 
         // loop while fork blocks are found in Tree.
         while let Some(chain_id) = self.block_indices().get_block_chain_id(&fork_block.hash) {
@@ -1117,6 +1183,11 @@ where
 
         if chain_appended {
             trace!(target: "blockchain_tree", ?new_canon_chain, "Canonical chain appended");
+        } else {
+            canonical.append_chain(new_canon_chain).map_err(|_| {
+                CanonicalError::from(BlockchainTreeError::BlockHashNotFoundInChain { block_hash })
+            })?;
+            new_canon_chain = canonical;
         }
         // update canonical index
         self.state.block_indices.canonicalize_blocks(new_canon_chain.blocks());
@@ -1262,7 +1333,6 @@ where
         };
         recorder.record_relative(MakeCanonicalAction::RetrieveStateTrieUpdates);
 
-        let cloned_bundle = state.bundle.clone();
         let provider_rw = self.externals.provider_factory.provider_rw()?;
         provider_rw
             .append_blocks_with_state(
@@ -1276,9 +1346,6 @@ where
         provider_rw.commit()?;
         recorder.record_relative(MakeCanonicalAction::CommitCanonicalChainToDatabase);
 
-        // update global canonical cache
-        apply_bundle_state(cloned_bundle);
-
         Ok(())
     }
 

crates/blockchain-tree/src/externals.rs

@@ -3,10 +3,10 @@
 use reth_consensus::Consensus;
 use reth_db::{static_file::HeaderMask, tables};
 use reth_db_api::{cursor::DbCursorRO, database::Database, transaction::DbTx};
-use reth_primitives::{BlockHash, BlockNumber, StaticFileSegment};
+use reth_primitives::{BlockHash, BlockNumber, StaticFileSegment, B256};
 use reth_provider::{
-    FinalizedBlockReader, FinalizedBlockWriter, ProviderFactory, StaticFileProviderFactory,
-    StatsReader,
+    BlockHashReader, FinalizedBlockReader, FinalizedBlockWriter, ProviderFactory,
+    StaticFileProviderFactory, StatsReader,
 };
 use reth_storage_errors::provider::ProviderResult;
 use std::{collections::BTreeMap, sync::Arc};
@@ -98,4 +98,10 @@ impl<DB: Database, E> TreeExternals<DB, E> {
         provider_rw.commit()?;
         Ok(())
     }
+
+    pub(crate) fn find_hash_by_number(&self, block_number: BlockNumber) -> ProviderResult<B256> {
+        let provider_ro = self.provider_factory.provider()?;
+        let hash = provider_ro.block_hash(block_number)?;
+        Ok(hash.unwrap_or_default())
+    }
 }