perf(core): snapshots

cmd/ethrex/cli.rs

@@ -376,7 +376,7 @@ pub async fn import_blocks(
         utils::read_chain_file(path)
     };
     let size = blocks.len();
-    for block in &blocks {
+    for block in blocks.iter() {
         let hash = block.hash();
         let number = block.header.number;
         info!("Adding block {number} with hash {hash:#x}.");
@@ -405,6 +405,11 @@ pub async fn import_blocks(
 
     // Make head canonical and label all special blocks correctly.
     if let Some(block) = blocks.last() {
+        _ = store
+            .reconstruct_snapshots_for_new_canonical_chain(block.hash())
+            .await
+            .inspect_err(|error| warn!("Failed to reconstruct snapshot: {}", error));
+
         store
             .update_finalized_block_number(block.header.number)
             .await?;

cmd/ethrex/l2/command.rs

@@ -341,7 +341,7 @@ impl Command {
                                 .await
                                 .expect("Error applying account updates");
 
-                            let (new_state_root, state_updates, accounts_updates) =
+                            let (new_state_root, state_updates, storage_updates) =
                                 (
                                     account_updates_list.state_trie_hash,
                                     account_updates_list.state_updates,
@@ -350,13 +350,12 @@ impl Command {
 
                             let pseudo_update_batch = UpdateBatch {
                                 account_updates: state_updates,
-                                storage_updates: accounts_updates,
-                                blocks: vec![],
-                                receipts: vec![],
-                                code_updates: vec![],
+                                storage_updates,
+                                ..Default::default()
                             };
 
-                            store.store_block_updates(pseudo_update_batch).await.expect("Error storing trie updates");
+                            let account_updates: Vec<_> = account_updates.values().cloned().collect();
+                            store.store_block_updates(pseudo_update_batch, &[&account_updates]).await.expect("Error storing trie updates");
 
                             new_trie = store.open_state_trie(new_state_root).expect("Error opening new state trie");
 

crates/blockchain/blockchain.rs

@@ -345,6 +345,14 @@ impl Blockchain {
         let state_updates = apply_updates_list.state_updates;
         let accounts_updates = apply_updates_list.storage_updates;
         let code_updates = apply_updates_list.code_updates;
+        let account_info_log_updates = self
+            .storage
+            .build_account_info_logs(std::iter::once((block, &account_updates.to_vec())))
+            .map_err(ChainError::StoreError)?;
+        let storage_log_updates = self
+            .storage
+            .build_account_storage_logs(std::iter::once((block, &account_updates.to_vec())))
+            .map_err(ChainError::StoreError)?;
 
         // Check state root matches the one in block header
         validate_state_root(&block.header, new_state_root)?;
@@ -355,10 +363,11 @@ impl Blockchain {
             blocks: vec![block.clone()],
             receipts: vec![(block.hash(), execution_result.receipts)],
             code_updates,
+            account_info_log_updates,
+            storage_log_updates,
         };
 
         self.storage
-            .clone()
             .store_block_updates(update_batch)
             .await
             .map_err(|e| e.into())
@@ -449,6 +458,8 @@ impl Blockchain {
         let mut total_gas_used = 0;
         let mut transactions_count = 0;
 
+        let mut account_updates_acc = Vec::with_capacity(blocks_len);
+
         let interval = Instant::now();
         for (i, block) in blocks.iter().enumerate() {
             // for the first block, we need to query the store
@@ -484,11 +495,12 @@ impl Blockchain {
             total_gas_used += block.header.gas_used;
             transactions_count += block.body.transactions.len();
             all_receipts.push((block.hash(), receipts));
-        }
 
-        let account_updates = vm
-            .get_state_transitions()
-            .map_err(|err| (ChainError::EvmError(err), None))?;
+            let account_updates = vm
+                .get_state_transitions()
+                .map_err(|err| (ChainError::EvmError(err), None))?;
+            account_updates_acc.push(account_updates);
+        }
 
         let last_block = blocks
             .last()
@@ -497,6 +509,11 @@ impl Blockchain {
         let last_block_number = last_block.header.number;
         let last_block_gas_limit = last_block.header.gas_limit;
 
+        let account_updates = account_updates_acc
+            .iter()
+            .flatten()
+            .cloned()
+            .collect::<Vec<_>>();
         // Apply the account updates over all blocks and compute the new state root
         let account_updates_list = self
             .storage
@@ -505,6 +522,8 @@ impl Blockchain {
             .map_err(|e| (e.into(), None))?
             .ok_or((ChainError::ParentStateNotFound, None))?;
 
+        let account_updates_per_block = blocks.iter().zip(account_updates_acc.iter());
+
         let new_state_root = account_updates_list.state_trie_hash;
         let state_updates = account_updates_list.state_updates;
         let accounts_updates = account_updates_list.storage_updates;
@@ -513,12 +532,25 @@ impl Blockchain {
         // Check state root matches the one in block header
         validate_state_root(&last_block.header, new_state_root).map_err(|e| (e, None))?;
 
+        // call helper function to build write logs
+        let account_info_log_updates = self
+            .storage
+            .build_account_info_logs(account_updates_per_block.clone())
+            .map_err(|e| (e.into(), None))?;
+
+        let storage_log_updates = self
+            .storage
+            .build_account_storage_logs(account_updates_per_block)
+            .map_err(|e| (e.into(), None))?;
+
         let update_batch = UpdateBatch {
             account_updates: state_updates,
             storage_updates: accounts_updates,
             blocks,
             receipts: all_receipts,
             code_updates,
+            account_info_log_updates,
+            storage_log_updates,
         };
 
         self.storage

crates/blockchain/fork_choice.rs

@@ -66,10 +66,10 @@ pub async fn apply_fork_choice(
         return Err(InvalidForkChoice::UnlinkedHead);
     };
 
-    let link_block_number = match new_canonical_blocks.last() {
-        Some((number, _)) => *number,
-        None => head.number,
-    };
+    let link_block_number = new_canonical_blocks
+        .last()
+        .map(|(number, _)| *number)
+        .unwrap_or(head.number);
 
     // Check that finalized and safe blocks are part of the new canonical chain.
     if let Some(ref finalized) = finalized_res {
@@ -99,7 +99,6 @@ pub async fn apply_fork_choice(
     }
 
     // Finished all validations.
-
     // Make all ancestors to head canonical.
     for (number, hash) in new_canonical_blocks {
         store.set_canonical_block(number, hash).await?;
@@ -112,6 +111,11 @@ pub async fn apply_fork_choice(
 
     // Make head canonical and label all special blocks correctly.
     store.set_canonical_block(head.number, head_hash).await?;
+
+    store
+        .reconstruct_snapshots_for_new_canonical_chain(head_hash)
+        .await?;
+
     if let Some(finalized) = finalized_res {
         store
             .update_finalized_block_number(finalized.number)
@@ -176,10 +180,8 @@ async fn find_link_with_canonical_chain(
         let parent_hash = header.parent_hash;
 
         // Check that the parent exists.
-        let parent_header = match store.get_block_header_by_hash(parent_hash) {
-            Ok(Some(header)) => header,
-            Ok(None) => return Ok(None),
-            Err(error) => return Err(error),
+        let Some(parent_header) = store.get_block_header_by_hash(parent_hash)? else {
+            return Ok(None);
         };
 
         if is_canonical(store, block_number, parent_hash).await? {

crates/blockchain/vm.rs

@@ -42,15 +42,43 @@ impl StoreVmDatabase {
 
 impl VmDatabase for StoreVmDatabase {
     fn get_account_info(&self, address: Address) -> Result<Option<AccountInfo>, EvmError> {
-        self.store
-            .get_account_info_by_hash(self.block_hash, address)
-            .map_err(|e| EvmError::DB(e.to_string()))
+        let block_for_current_snapshot = self
+            .store
+            .get_block_for_current_snapshot()
+            .map_err(|e| EvmError::DB(e.to_string()))?;
+
+        if Some(self.block_hash) == block_for_current_snapshot {
+            self.store.get_current_account_info(address)
+        } else {
+            tracing::warn!(
+                "account_info snapshot miss: expected: {:?} got: {:?}",
+                block_for_current_snapshot,
+                self.block_hash
+            );
+            self.store
+                .get_account_info_by_hash(self.block_hash, address)
+        }
+        .map_err(|e| EvmError::DB(e.to_string()))
     }
 
     fn get_storage_slot(&self, address: Address, key: H256) -> Result<Option<U256>, EvmError> {
-        self.store
-            .get_storage_at_hash(self.block_hash, address, key)
-            .map_err(|e| EvmError::DB(e.to_string()))
+        let block_for_current_snapshot = self
+            .store
+            .get_block_for_current_snapshot()
+            .map_err(|e| EvmError::DB(e.to_string()))?;
+
+        if Some(self.block_hash) == block_for_current_snapshot {
+            self.store.get_current_storage(address, key)
+        } else {
+            tracing::warn!(
+                "account_storage snapshot miss: expected: {:?} got: {:?}",
+                block_for_current_snapshot,
+                self.block_hash
+            );
+            self.store
+                .get_storage_at_hash(self.block_hash, address, key)
+        }
+        .map_err(|e| EvmError::DB(e.to_string()))
     }
 
     fn get_block_hash(&self, block_number: u64) -> Result<H256, EvmError> {

crates/common/types/account.rs

@@ -25,7 +25,7 @@ pub struct Account {
     pub storage: HashMap<H256, U256>,
 }
 
-#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Eq)]
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Eq, Hash)]
 pub struct AccountInfo {
     pub code_hash: H256,
     pub balance: U256,

crates/networking/p2p/sync.rs

@@ -457,13 +457,24 @@ impl Syncer {
                     .set_latest_valid_ancestor(failed_block_hash, last_valid_hash)
                     .await?;
 
+                // We also need to use the correct snapshot.
+                store
+                    .reconstruct_snapshots_for_new_canonical_chain(last_valid_hash)
+                    .await
+                    .map_err(SyncError::Store)?;
                 // TODO(#2127): Just marking the failing ancestor is enough for the the Missing Ancestors hive test,
                 // we want to look at a more robust solution in the future if needed.
             }
 
             return Err(error.into());
         }
 
+        // We also need to use the correct snapshot.
+        store
+            .reconstruct_snapshots_for_new_canonical_chain(last_block.hash())
+            .await
+            .map_err(SyncError::Store)?;
+
         store
             .update_latest_block_number(last_block.header.number)
             .await?;

crates/storage/api.rs

@@ -1,8 +1,9 @@
 use bytes::Bytes;
 use ethereum_types::{H256, U256};
+use ethrex_common::Address;
 use ethrex_common::types::{
-    AccountState, Block, BlockBody, BlockHash, BlockHeader, BlockNumber, ChainConfig, Index,
-    Receipt, Transaction, payload::PayloadBundle,
+    AccountInfo, AccountState, Block, BlockBody, BlockHash, BlockHeader, BlockNumber, ChainConfig,
+    Index, Receipt, Transaction, payload::PayloadBundle,
 };
 use std::{fmt::Debug, panic::RefUnwindSafe};
 
@@ -17,6 +18,10 @@ pub trait StoreEngine: Debug + Send + Sync + RefUnwindSafe {
     /// Store changes in a batch from a vec of blocks
     async fn apply_updates(&self, update_batch: UpdateBatch) -> Result<(), StoreError>;
 
+    async fn undo_writes_until_canonical(&self) -> Result<(), StoreError>;
+
+    async fn replay_writes_until_head(&self, head: H256) -> Result<(), StoreError>;
+
     /// Add a batch of blocks in a single transaction.
     /// This will store -> BlockHeader, BlockBody, BlockTransactions, BlockNumber.
     async fn add_blocks(&self, blocks: Vec<Block>) -> Result<(), StoreError>;
@@ -304,6 +309,27 @@ pub trait StoreEngine: Debug + Send + Sync + RefUnwindSafe {
         &self,
     ) -> Result<Option<[H256; STATE_TRIE_SEGMENTS]>, StoreError>;
 
+    async fn setup_genesis_flat_account_storage(
+        &self,
+        genesis_block_number: u64,
+        genesis_block_hash: H256,
+        genesis_accounts: &[(Address, H256, U256)],
+    ) -> Result<(), StoreError>;
+
+    async fn setup_genesis_flat_account_info(
+        &self,
+        genesis_block_number: u64,
+        genesis_block_hash: H256,
+        genesis_accounts: &[(Address, u64, U256, H256, bool)],
+    ) -> Result<(), StoreError>;
+
+    fn get_block_for_current_snapshot(&self) -> Result<Option<BlockHash>, StoreError>;
+
+    fn get_current_storage(&self, address: Address, key: H256) -> Result<Option<U256>, StoreError>;
+
+    fn get_current_account_info(&self, address: Address)
+    -> Result<Option<AccountInfo>, StoreError>;
+
     /// Sets storage trie paths in need of healing, grouped by hashed address
     /// This will overwite previously stored paths for the received storages but will not remove other storage's paths
     async fn set_storage_heal_paths(