perf(core): snapshots

cmd/ethrex/cli.rs

@@ -419,6 +419,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

@@ -384,7 +384,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,
@@ -393,10 +393,8 @@ 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");

crates/blockchain/blockchain.rs

@@ -362,16 +362,25 @@ impl Blockchain {
         account_updates: &[AccountUpdate],
     ) -> Result<(), ChainError> {
         // Apply the account updates over the last block's state and compute the new state root
+        let parent_hash = block.header.parent_hash;
         let apply_updates_list = self
             .storage
-            .apply_account_updates_batch(block.header.parent_hash, account_updates)
+            .apply_account_updates_batch(parent_hash, account_updates)
             .await?
             .ok_or(ChainError::ParentStateNotFound)?;
 
         let new_state_root = apply_updates_list.state_trie_hash;
         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(parent_hash, account_updates.iter())
+            .map_err(ChainError::StoreError)?;
+        let storage_log_updates = self
+            .storage
+            .build_account_storage_logs(parent_hash, account_updates.iter())
+            .map_err(ChainError::StoreError)?;
 
         // Check state root matches the one in block header
         validate_state_root(&block.header, new_state_root)?;
@@ -382,10 +391,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())
@@ -517,21 +527,22 @@ impl Blockchain {
             log_batch_progress(blocks_len, i);
         }
 
-        let account_updates = vm
-            .get_state_transitions()
-            .map_err(|err| (ChainError::EvmError(err), None))?;
-
         let last_block = blocks
             .last()
             .ok_or_else(|| (ChainError::Custom("Last block not found".into()), None))?;
 
         let last_block_number = last_block.header.number;
         let last_block_gas_limit = last_block.header.gas_limit;
 
+        let account_updates = vm
+            .get_state_transitions()
+            .map_err(|err| (ChainError::EvmError(err), None))?;
+
+        let parent_hash = first_block_header.parent_hash;
         // Apply the account updates over all blocks and compute the new state root
         let account_updates_list = self
             .storage
-            .apply_account_updates_batch(first_block_header.parent_hash, &account_updates)
+            .apply_account_updates_batch(parent_hash, &account_updates)
             .await
             .map_err(|e| (e.into(), None))?
             .ok_or((ChainError::ParentStateNotFound, None))?;
@@ -544,12 +555,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(parent_hash, account_updates.iter())
+            .map_err(|e| (e.into(), None))?;
+
+        let storage_log_updates = self
+            .storage
+            .build_account_storage_logs(parent_hash, account_updates.iter())
+            .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,33 @@ 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 {
+            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 {
+            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

@@ -463,13 +463,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,17 @@ 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>;
 
+    /// Reverts the changes in the state using write logs until reaching a canonical block.
+    /// Used to restore state after a reorg by walking backwards through logs.
+    async fn undo_writes_until_canonical(&self) -> Result<(), StoreError>;
+
+    /// Replays writes from the write logs until the head block is reached.
+    /// This is used to restore the flat tables from the write logs after a reorg.
+    /// Assumes that the current flat representation corresponds to a block in the canonical chain.
+    /// Must be called after [`undo_writes_until_canonical`](StoreEngine::undo_writes_until_canonical)
+    /// to complete reorg recovery.
+    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>;
@@ -307,6 +319,32 @@ pub trait StoreEngine: Debug + Send + Sync + RefUnwindSafe {
         &self,
     ) -> Result<Option<[H256; STATE_TRIE_SEGMENTS]>, StoreError>;
 
+    /// Sets up the flat tables of account storage for the genesis block
+    async fn setup_genesis_flat_account_storage(
+        &self,
+        genesis_block_number: u64,
+        genesis_block_hash: H256,
+        genesis_accounts: &[(Address, H256, U256)],
+    ) -> Result<(), StoreError>;
+
+    /// Sets up the flat tables of account info for the genesis block
+    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>;
+
+    /// Gets the block hash for the current snapshot
+    fn get_block_for_current_snapshot(&self) -> Result<Option<BlockHash>, StoreError>;
+
+    /// Gets the account storage for an account
+    fn get_current_storage(&self, address: Address, key: H256) -> Result<Option<U256>, StoreError>;
+
+    /// Gets the account info for an account
+    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(

crates/storage/error.rs

@@ -53,3 +53,10 @@ pub enum StoreError {
     #[error("Failed to lock database for writing")]
     LockError,
 }
+
+#[cfg(feature = "libmdbx")]
+impl From<anyhow::Error> for StoreError {
+    fn from(err: anyhow::Error) -> Self {
+        StoreError::LibmdbxError(err)
+    }
+}

crates/storage/rlp.rs

@@ -1,7 +1,14 @@
 use std::fmt::Debug;
 use std::marker::PhantomData;
 
+#[cfg(feature = "redb")]
+use crate::store_db::codec::{
+    account_info_log_entry::AccountInfoLogEntry, account_storage_log_entry::AccountStorageLogEntry,
+    block_num_hash::BlockNumHash,
+};
 use bytes::Bytes;
+#[cfg(feature = "redb")]
+use ethrex_common::{Address, U256, types::AccountInfo};
 use ethrex_common::{
     H256,
     types::{
@@ -24,6 +31,22 @@ pub type AccountHashRLP = Rlp<H256>;
 pub type AccountStateRLP = Rlp<AccountState>;
 pub type TriePathsRLP = Rlp<Vec<Nibbles>>;
 
+// Snapshot state RLP types
+#[cfg(feature = "redb")]
+pub type AccountInfoLogEntryRLP = Rlp<AccountInfoLogEntry>;
+#[cfg(feature = "redb")]
+pub type AccountStorageLogEntryRLP = Rlp<AccountStorageLogEntry>;
+#[cfg(feature = "redb")]
+pub type BlockNumHashRLP = Rlp<BlockNumHash>;
+#[cfg(feature = "redb")]
+pub type AccountAddressRLP = Rlp<Address>;
+#[cfg(feature = "redb")]
+pub type AccountInfoRLP = Rlp<AccountInfo>;
+#[cfg(feature = "redb")]
+pub type AccountStorageKeyRLP = Rlp<H256>;
+#[cfg(feature = "redb")]
+pub type AccountStorageValueRLP = Rlp<U256>;
+
 // Block types
 pub type BlockHashRLP = Rlp<BlockHash>;
 pub type BlockHeaderRLP = Rlp<BlockHeader>;