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,9 @@ 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, L2Config};
+use ethrex_blockchain::{Blockchain, BlockchainOptions, BlockchainType, L2Config};
+use ethrex_common::fd_limit::raise_fd_limit;
+use ethrex_common::types::Genesis;
 use ethrex_common::types::fee_config::{FeeConfig, L1FeeConfig, OperatorFeeConfig};
 use ethrex_common::{Address, types::DEFAULT_BUILDER_GAS_CEIL};
 use ethrex_l2::SequencerConfig;
@@ -18,6 +20,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;
@@ -146,14 +150,19 @@ pub async fn init_l2(
     opts: L2Options,
     log_filter_handler: Option<reload::Handle<EnvFilter, Registry>>,
 ) -> eyre::Result<()> {
+    raise_fd_limit()?;
+
     let datadir = opts.node_opts.datadir.clone();
     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?;
@@ -180,10 +189,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);
@@ -277,6 +294,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();
 
@@ -344,3 +365,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_header = checkpoint_store
+        .get_block_header(checkpoint_latest_block_number)?
+        .ok_or(eyre::eyre!(
+            "latest block header ({checkpoint_latest_block_number}) not found in checkpoint store"
+        ))?;
+
+    let db_latest_block_header = store
+        .get_block_header(db_latest_block_number)?
+        .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_header.state_root != db_latest_block_header.state_root {
+        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

@@ -201,13 +201,25 @@ impl Blockchain {
         blocks: &[Block],
         fee_configs: Option<&[FeeConfig]>,
     ) -> 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;
+
+        // Later on, we need to access block hashes by number. To avoid needing
+        // to apply fork choice for each block, we cache them here.
+        let mut block_hashes_map: BTreeMap<u64, H256> = blocks
+            .iter()
+            .cloned()
+            .map(|block| (block.header.number, block.hash()))
+            .collect();
+
+        block_hashes_map.insert(
+            first_block_header.number.saturating_sub(1),
+            first_block_header.parent_hash,
+        );
 
         // Get state at previous block
         let trie = self
@@ -216,7 +228,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
@@ -232,10 +254,18 @@ impl Blockchain {
 
         for (i, block) in blocks.iter().enumerate() {
             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 {
+
+            let mut vm = match self.options.r#type {
                 BlockchainType::L1 => Evm::new_from_db_for_l1(logger.clone()),
                 BlockchainType::L2(_) => {
                     let l2_config = match fee_configs {
@@ -253,7 +283,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()?;
 
@@ -276,7 +307,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 {
@@ -320,6 +353,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
@@ -342,14 +376,21 @@ 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?;
+
             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())
@@ -359,15 +400,40 @@ 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_hashes_map
+                            .get(&block.header.number)
+                            .ok_or(ChainError::WitnessGeneration(
+                                "Block hash not found for witness generation".to_string(),
+                            ))?
+                            .to_owned(),
+                    )
+                    .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
@@ -376,6 +442,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
@@ -385,17 +452,26 @@ 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()),
+            let hash = block_hashes_map
+                .get(&block_number)
+                .ok_or(ChainError::WitnessGeneration(format!(
+                    "Failed to get block {block_number} hash"
+                )))?;
+            let header = self.storage.get_block_header_by_hash(*hash)?.ok_or(
+                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"]