fix(l2): use checkpoints to persist previous batch state

cmd/ethrex/Cargo.toml

@@ -83,7 +83,7 @@ c-kzg = [
   "ethrex-crypto/c-kzg",
 ]
 metrics = ["ethrex-blockchain/metrics", "ethrex-l2?/metrics"]
-rocksdb = ["ethrex-storage/rocksdb", "ethrex-p2p/rocksdb"]
+rocksdb = ["ethrex-storage/rocksdb", "ethrex-p2p/rocksdb", "ethrex-l2?/rocksdb"]
 jemalloc = ["dep:tikv-jemallocator"]
 jemalloc_profiling = [
   "jemalloc",

cmd/ethrex/l2/initializers.rs

@@ -7,7 +7,8 @@ use crate::l2::{L2Options, SequencerOptions};
 use crate::utils::{
     NodeConfigFile, get_client_version, init_datadir, read_jwtsecret_file, store_node_config_file,
 };
-use ethrex_blockchain::{Blockchain, BlockchainType};
+use ethrex_blockchain::{Blockchain, BlockchainOptions, BlockchainType};
+use ethrex_common::types::Genesis;
 use ethrex_common::types::fee_config::{FeeConfig, OperatorFeeConfig};
 use ethrex_common::{Address, types::DEFAULT_BUILDER_GAS_CEIL};
 use ethrex_l2::SequencerConfig;
@@ -18,6 +19,8 @@ use ethrex_p2p::{
     sync_manager::SyncManager,
     types::{Node, NodeRecord},
 };
+#[cfg(feature = "rocksdb")]
+use ethrex_storage::EngineType;
 use ethrex_storage::Store;
 use ethrex_storage_rollup::{EngineTypeRollup, StoreRollup};
 use secp256k1::SecretKey;
@@ -150,10 +153,13 @@ pub async fn init_l2(
     init_datadir(&opts.node_opts.datadir);
     let rollup_store_dir = datadir.join("rollup_store");
 
+    // Checkpoints are stored in the main datadir
+    let checkpoints_dir = datadir.clone();
+
     let network = get_network(&opts.node_opts);
 
     let genesis = network.get_genesis()?;
-    let store = init_store(&datadir, genesis).await;
+    let store = init_store(&datadir, genesis.clone()).await;
     let rollup_store = init_rollup_store(&rollup_store_dir).await;
 
     let operator_fee_config = get_operator_fee_config(&opts.sequencer_opts).await?;
@@ -172,10 +178,18 @@ pub async fn init_l2(
         perf_logs_enabled: true,
     };
 
-    let blockchain = init_blockchain(store.clone(), blockchain_opts);
+    let blockchain = init_blockchain(store.clone(), blockchain_opts.clone());
 
     regenerate_head_state(&store, &blockchain).await?;
 
+    let (initial_checkpoint_store, initial_checkpoint_blockchain) = initialize_checkpoint(
+        &store,
+        &checkpoints_dir.join("initial_checkpoint"),
+        genesis.clone(),
+        blockchain_opts,
+    )
+    .await?;
+
     let signer = get_signer(&datadir);
 
     let local_p2p_node = get_local_p2p_node(&opts.node_opts, &signer);
@@ -269,6 +283,10 @@ pub async fn init_l2(
             "http://{}:{}",
             opts.node_opts.http_addr, opts.node_opts.http_port
         ),
+        initial_checkpoint_store,
+        initial_checkpoint_blockchain,
+        genesis,
+        checkpoints_dir,
     )
     .into_future();
 
@@ -321,3 +339,92 @@ pub async fn get_operator_fee_config(
         };
     Ok(operator_fee_config)
 }
+
+/// Initializes a checkpoint of the given store at the specified path.
+///
+/// If there's no previous checkpoint, it creates one from the current store state.
+///
+/// This function performs the following steps:
+/// 1. Creates a checkpoint of the provided store at the specified path.
+/// 2. Initializes a new store and blockchain for the checkpoint.
+/// 3. Regenerates the head state in the checkpoint store.
+/// 4. Validates that the checkpoint store's head block number and latest block match those of the original store.
+async fn initialize_checkpoint(
+    store: &Store,
+    path: &Path,
+    genesis: Genesis,
+    blockchain_opts: BlockchainOptions,
+) -> eyre::Result<(Store, Arc<Blockchain>)> {
+    // If the checkpoint is not present, create it
+    if !path.exists() {
+        store.create_checkpoint(path).await?;
+    }
+
+    // We now load the checkpoint, validate it, and regenerate its state.
+
+    #[cfg(feature = "rocksdb")]
+    let engine_type = EngineType::RocksDB;
+    #[cfg(not(feature = "rocksdb"))]
+    let engine_type = EngineType::InMemory;
+
+    let checkpoint_store = {
+        let checkpoint_store_inner = Store::new(path, engine_type)?;
+
+        checkpoint_store_inner
+            .add_initial_state(genesis.clone())
+            .await?;
+
+        checkpoint_store_inner
+    };
+
+    let checkpoint_blockchain =
+        Arc::new(Blockchain::new(checkpoint_store.clone(), blockchain_opts));
+
+    let checkpoint_head_block_number = checkpoint_store.get_latest_block_number().await?;
+
+    let db_head_block_number = store.get_latest_block_number().await?;
+
+    if checkpoint_head_block_number != db_head_block_number {
+        return Err(eyre::eyre!(
+            "checkpoint store head block number does not match main store head block number before regeneration"
+        ));
+    }
+
+    regenerate_head_state(&checkpoint_store, &checkpoint_blockchain).await?;
+
+    let checkpoint_latest_block_number = checkpoint_store.get_latest_block_number().await?;
+
+    let db_latest_block_number = store.get_latest_block_number().await?;
+
+    let checkpoint_latest_block = checkpoint_store
+        .get_block_by_number(checkpoint_latest_block_number)
+        .await?
+        .ok_or(eyre::eyre!("latest block not found in checkpoint store"))?;
+
+    let db_latest_block = store
+        .get_block_by_number(db_latest_block_number)
+        .await?
+        .ok_or(eyre::eyre!("latest block not found in main store"))?;
+
+    // Final sanity check
+
+    if !checkpoint_store.has_state_root(checkpoint_latest_block.header.state_root)? {
+        return Err(eyre::eyre!(
+            "checkpoint store state is not regenerated properly"
+        ));
+    }
+
+    if checkpoint_latest_block_number != db_head_block_number {
+        return Err(eyre::eyre!(
+            "checkpoint store latest block number does not match main store head block number after regeneration"
+        ));
+    }
+
+    if checkpoint_latest_block.hash() != db_latest_block.hash() {
+        return Err(eyre::eyre!(
+            "checkpoint store latest block hash does not match main store latest block hash after regeneration"
+        ));
+    }
+
+    Ok((checkpoint_store, checkpoint_blockchain))
+}

crates/blockchain/blockchain.rs

@@ -48,6 +48,8 @@ use ethrex_metrics::metrics_blocks::METRICS_BLOCKS;
 #[cfg(feature = "c-kzg")]
 use ethrex_common::types::BlobsBundle;
 
+use crate::fork_choice::apply_fork_choice;
+
 const MAX_PAYLOADS: usize = 10;
 const MAX_MEMPOOL_SIZE_DEFAULT: usize = 10_000;
 
@@ -186,13 +188,12 @@ impl Blockchain {
         &self,
         blocks: &[Block],
     ) -> Result<ExecutionWitness, ChainError> {
-        let first_block_header = blocks
+        let first_block_header = &blocks
             .first()
             .ok_or(ChainError::WitnessGeneration(
                 "Empty block batch".to_string(),
             ))?
-            .header
-            .clone();
+            .header;
 
         // Get state at previous block
         let trie = self
@@ -201,7 +202,17 @@ impl Blockchain {
             .map_err(|_| ChainError::ParentStateNotFound)?
             .ok_or(ChainError::ParentStateNotFound)?;
 
-        let (state_trie_witness, mut trie) = TrieLogger::open_trie(trie);
+        let (mut current_trie_witness, mut trie) = TrieLogger::open_trie(trie);
+
+        // For each block, a new TrieLogger will be opened, each containing the
+        // witness accessed during the block execution. We need to accumulate
+        // all the nodes accessed during the entire batch execution.
+        let mut accumulated_state_trie_witness = current_trie_witness
+            .lock()
+            .map_err(|_| {
+                ChainError::WitnessGeneration("Failed to lock state trie witness".to_string())
+            })?
+            .clone();
 
         let mut touched_account_storage_slots = BTreeMap::new();
         // This will become the state trie + storage trie
@@ -214,12 +225,20 @@ impl Blockchain {
 
         let mut block_hashes = HashMap::new();
         let mut codes = Vec::new();
-
+                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    
         for block in blocks {
             let parent_hash = block.header.parent_hash;
+
+            // This assumes that the user has the necessary state stored already,
+            // so if the user only has the state previous to the first block, it
+            // will fail in the second iteration of this for loop. To ensure this,
+            // doesn't fail, later in this function we store the new state after
+            // re-execution.
             let vm_db: DynVmDatabase =
                 Box::new(StoreVmDatabase::new(self.storage.clone(), parent_hash));
+
             let logger = Arc::new(DatabaseLogger::new(Arc::new(Mutex::new(Box::new(vm_db)))));
+
             let mut vm = match self.options.r#type {
                 BlockchainType::L1 => Evm::new_from_db_for_l1(logger.clone()),
                 BlockchainType::L2(fee_config) => {
@@ -228,7 +247,8 @@ impl Blockchain {
             };
 
             // Re-execute block with logger
-            vm.execute_block(block)?;
+            let execution_result = vm.execute_block(block)?;
+
             // Gather account updates
             let account_updates = vm.get_state_transitions()?;
 
@@ -251,7 +271,9 @@ impl Blockchain {
                     ChainError::WitnessGeneration("Failed to get block hashes".to_string())
                 })?
                 .clone();
+
             block_hashes.extend(logger_block_hashes);
+
             // Access all the accounts needed for withdrawals
             if let Some(withdrawals) = block.body.withdrawals.as_ref() {
                 for withdrawal in withdrawals {
@@ -295,6 +317,7 @@ impl Blockchain {
                     used_storage_tries.insert(*account, (storage_trie_witness, storage_trie));
                 }
             }
+
             // Store all the accessed evm bytecodes
             for code_hash in logger
                 .code_accessed
@@ -317,14 +340,39 @@ impl Blockchain {
             }
 
             // Apply account updates to the trie recording all the necessary nodes to do so
-            let (updated_trie, storage_tries_after_update) = self
+            let (storage_tries_after_update, account_updates_list) = self
                 .storage
                 .apply_account_updates_from_trie_with_witness(
                     trie,
                     &account_updates,
                     used_storage_tries,
                 )
                 .await?;
+
+            // We cannot ensure that the users of this function have the necessary
+            // state stored, so in order for it to not assume anything, we update
+            // the storage with the new state after re-execution
+            self.store_block(block.clone(), account_updates_list, execution_result)
+                .await?;
+
+            // Later in this function we query the batch block headers by number,
+            // to do so, they need to be marked as canonical.
+            // Only users that have the state previous to the first batch block
+            // need this step.
+            if self
+                .storage
+                .get_block_header(block.header.number)?
+                .is_none()
+            {
+                apply_fork_choice(&self.storage, block.hash(), block.hash(), block.hash())
+                    .await
+                    .map_err(|e| {
+                        ChainError::WitnessGeneration(format!(
+                            "Failed to apply fork choice after witness generation: {e}"
+                        ))
+                    })?;
+            }
+
             for (address, (witness, _storage_trie)) in storage_tries_after_update {
                 let mut witness = witness.lock().map_err(|_| {
                     ChainError::WitnessGeneration("Failed to lock storage trie witness".to_string())
@@ -334,15 +382,33 @@ impl Blockchain {
                 used_trie_nodes.extend_from_slice(&witness);
                 touched_account_storage_slots.entry(address).or_default();
             }
+
+            let (new_state_trie_witness, updated_trie) = TrieLogger::open_trie(
+                self.storage
+                    .state_trie(block.hash())
+                    .map_err(|_| ChainError::ParentStateNotFound)?
+                    .ok_or(ChainError::ParentStateNotFound)?,
+            );
+
+            // Use the updated state trie for the next block
             trie = updated_trie;
+
+            for state_trie_witness in current_trie_witness
+                .lock()
+                .map_err(|_| {
+                    ChainError::WitnessGeneration("Failed to lock state trie witness".to_string())
+                })?
+                .iter()
+            {
+                accumulated_state_trie_witness.insert(state_trie_witness.clone());
+            }
+
+            current_trie_witness = new_state_trie_witness;
         }
 
-        // Get the witness for the state trie
-        let mut state_trie_witness = state_trie_witness.lock().map_err(|_| {
-            ChainError::WitnessGeneration("Failed to lock state trie witness".to_string())
-        })?;
-        let state_trie_witness = std::mem::take(&mut *state_trie_witness);
-        used_trie_nodes.extend_from_slice(&Vec::from_iter(state_trie_witness.into_iter()));
+        used_trie_nodes
+            .extend_from_slice(&Vec::from_iter(accumulated_state_trie_witness.into_iter()));
+
         // If the witness is empty at least try to store the root
         if used_trie_nodes.is_empty()
             && let Some(root) = root_node
@@ -351,6 +417,7 @@ impl Blockchain {
         }
 
         let mut needed_block_numbers = block_hashes.keys().collect::<Vec<_>>();
+
         needed_block_numbers.sort();
 
         // Last needed block header for the witness is the parent of the last block we need to execute
@@ -360,17 +427,21 @@ impl Blockchain {
             .header
             .number
             .saturating_sub(1);
+
         // The first block number we need is either the parent of the first block number or the earliest block number used by BLOCKHASH
         let mut first_needed_block_number = first_block_header.number.saturating_sub(1);
+
         if let Some(block_number_from_logger) = needed_block_numbers.first()
             && **block_number_from_logger < first_needed_block_number
         {
             first_needed_block_number = **block_number_from_logger;
         }
+
         let mut block_headers_bytes = Vec::new();
+
         for block_number in first_needed_block_number..=last_needed_block_number {
             let header = self.storage.get_block_header(block_number)?.ok_or(
-                ChainError::WitnessGeneration("Failed to get block header".to_string()),
+                ChainError::WitnessGeneration(format!("Failed to get block {block_number} header")),
             )?;
             block_headers_bytes.push(header.encode_to_vec());
         }

crates/l2/Cargo.toml

@@ -80,5 +80,6 @@ panic = "deny"
 
 [features]
 default = ["l2"]
+rocksdb = []
 metrics = ["ethrex-blockchain/metrics"]
 l2 = ["guest_program/l2"]