feat(stages): add table checkpoint to AccountHashing and StorageHashing

bin/reth/src/dump_stage/hashing_account.rs

@@ -36,8 +36,7 @@ pub(crate) async fn dump_hashing_account_stage<DB: Database>(
     unwind_and_copy::<DB>(db_tool, from, tip_block_number, &output_db).await?;
 
     if should_run {
-        println!("\n# AccountHashing stage does not support dry run, so it will actually be committing changes.");
-        run(output_db, to, from).await?;
+        dry_run(output_db, to, from).await?;
     }
 
     Ok(())
@@ -69,7 +68,7 @@ async fn unwind_and_copy<DB: Database>(
 }
 
 /// Try to re-execute the stage straightaway
-async fn run(
+async fn dry_run(
     output_db: reth_db::mdbx::Env<reth_db::mdbx::WriteMap>,
     to: u64,
     from: u64,
@@ -82,15 +81,21 @@ async fn run(
         ..Default::default()
     };
 
-    exec_stage
-        .execute(
-            &mut tx,
-            reth_stages::ExecInput {
-                previous_stage: Some((StageId("Another"), to)),
-                stage_progress: Some(from),
-            },
-        )
-        .await?;
+    let mut exec_output = false;
+    while !exec_output {
+        exec_output = exec_stage
+            .execute(
+                &mut tx,
+                reth_stages::ExecInput {
+                    previous_stage: Some((StageId("Another"), to)),
+                    stage_progress: Some(from),
+                },
+            )
+            .await?
+            .done;
+    }
+
+    tx.drop()?;
 
     info!(target: "reth::cli", "Success.");
 

bin/reth/src/dump_stage/hashing_storage.rs

@@ -6,8 +6,9 @@ use crate::{
 use eyre::Result;
 use reth_db::{database::Database, table::TableImporter, tables};
 use reth_provider::Transaction;
-use reth_stages::{stages::StorageHashingStage, Stage, UnwindInput};
+use reth_stages::{stages::StorageHashingStage, Stage, StageId, UnwindInput};
 use std::ops::DerefMut;
+use tracing::info;
 
 pub(crate) async fn dump_hashing_storage_stage<DB: Database>(
     db_tool: &mut DbTool<'_, DB>,
@@ -16,14 +17,14 @@ pub(crate) async fn dump_hashing_storage_stage<DB: Database>(
     output_db: &PlatformPath<DbPath>,
     should_run: bool,
 ) -> Result<()> {
-    if should_run {
-        eyre::bail!("StorageHashing stage does not support dry run.")
-    }
-
     let (output_db, tip_block_number) = setup::<DB>(from, to, output_db, db_tool)?;
 
     unwind_and_copy::<DB>(db_tool, from, tip_block_number, &output_db).await?;
 
+    if should_run {
+        dry_run(output_db, to, from).await?;
+    }
+
     Ok(())
 }
 
@@ -53,3 +54,38 @@ async fn unwind_and_copy<DB: Database>(
 
     Ok(())
 }
+
+/// Try to re-execute the stage straightaway
+async fn dry_run(
+    output_db: reth_db::mdbx::Env<reth_db::mdbx::WriteMap>,
+    to: u64,
+    from: u64,
+) -> eyre::Result<()> {
+    info!(target: "reth::cli", "Executing stage.");
+
+    let mut tx = Transaction::new(&output_db)?;
+    let mut exec_stage = StorageHashingStage {
+        clean_threshold: 1, // Forces hashing from scratch
+        ..Default::default()
+    };
+
+    let mut exec_output = false;
+    while !exec_output {
+        exec_output = exec_stage
+            .execute(
+                &mut tx,
+                reth_stages::ExecInput {
+                    previous_stage: Some((StageId("Another"), to)),
+                    stage_progress: Some(from),
+                },
+            )
+            .await?
+            .done;
+    }
+
+    tx.drop()?;
+
+    info!(target: "reth::cli", "Success.");
+
+    Ok(())
+}

crates/primitives/src/checkpoints.rs

@@ -0,0 +1,42 @@
+use crate::{Address, H256};
+use reth_codecs::{main_codec, Compact};
+
+/// Saves the progress of MerkleStage
+#[main_codec]
+#[derive(Default, Debug, Copy, Clone, PartialEq)]
+pub struct ProofCheckpoint {
+    /// The next hashed account to insert into the trie.
+    pub hashed_address: Option<H256>,
+    /// The next storage entry to insert into the trie.
+    pub storage_key: Option<H256>,
+    /// Current intermediate root for `AccountsTrie`.
+    pub account_root: Option<H256>,
+    /// Current intermediate storage root from an account.
+    pub storage_root: Option<H256>,
+}
+
+/// Saves the progress of AccountHashing
+#[main_codec]
+#[derive(Default, Debug, Copy, Clone, PartialEq)]
+pub struct AccountHashingCheckpoint {
+    /// The next account to start hashing from
+    pub address: Option<Address>,
+    /// Start transition id
+    pub from: u64,
+    /// Last transition id
+    pub to: u64,
+}
+
+/// Saves the progress of StorageHashing
+#[main_codec]
+#[derive(Default, Debug, Copy, Clone, PartialEq)]
+pub struct StorageHashingCheckpoint {
+    /// The next account to start hashing from
+    pub address: Option<Address>,
+    /// The next storage slot to start hashing from
+    pub storage: Option<H256>,
+    /// Start transition id
+    pub from: u64,
+    /// Last transition id
+    pub to: u64,
+}

crates/primitives/src/lib.rs

@@ -14,6 +14,7 @@ mod bits;
 mod block;
 pub mod bloom;
 mod chain;
+mod checkpoints;
 pub mod constants;
 pub mod contract;
 mod error;
@@ -34,7 +35,6 @@ mod withdrawal;
 
 /// Helper function for calculating Merkle proofs and hashes
 pub mod proofs;
-pub use proofs::ProofCheckpoint;
 
 pub use account::{Account, Bytecode};
 pub use bits::H512;
@@ -46,6 +46,7 @@ pub use chain::{
     AllGenesisFormats, Chain, ChainInfo, ChainSpec, ChainSpecBuilder, ForkCondition, GOERLI,
     MAINNET, SEPOLIA,
 };
+pub use checkpoints::{AccountHashingCheckpoint, ProofCheckpoint, StorageHashingCheckpoint};
 pub use constants::{
     EMPTY_OMMER_ROOT, GOERLI_GENESIS, KECCAK_EMPTY, MAINNET_GENESIS, SEPOLIA_GENESIS,
 };

crates/primitives/src/proofs.rs

@@ -1,5 +1,3 @@
-use std::collections::HashMap;
-
 use crate::{
     keccak256, Address, Bytes, GenesisAccount, Header, Log, Receipt, TransactionSigned, Withdrawal,
     H256,
@@ -8,8 +6,8 @@ use bytes::BytesMut;
 use hash_db::Hasher;
 use hex_literal::hex;
 use plain_hasher::PlainHasher;
-use reth_codecs::{main_codec, Compact};
 use reth_rlp::Encodable;
+use std::collections::HashMap;
 use triehash::{ordered_trie_root, sec_trie_root};
 
 /// Keccak-256 hash of the RLP of an empty list, KEC("\xc0").
@@ -35,23 +33,8 @@ impl Hasher for KeccakHasher {
     }
 }
 
-/// Saves the progress of MerkleStage
-#[main_codec]
-#[derive(Default, Debug, Copy, Clone, PartialEq)]
-pub struct ProofCheckpoint {
-    /// The next hashed account to insert into the trie.
-    pub hashed_address: Option<H256>,
-    /// The next storage entry to insert into the trie.
-    pub storage_key: Option<H256>,
-    /// Current intermediate root for `AccountsTrie`.
-    pub account_root: Option<H256>,
-    /// Current intermediate storage root from an account.
-    pub storage_root: Option<H256>,
-}
-
 /// Calculate a transaction root.
 ///
-/// Iterates over the given transactions and the merkle merkle trie root of
 /// `(rlp(index), encoded(tx))` pairs.
 pub fn calculate_transaction_root<'a>(
     transactions: impl IntoIterator<Item = &'a TransactionSigned>,

crates/stages/Cargo.toml

@@ -18,6 +18,7 @@ normal = [
 reth-primitives = { path = "../primitives" }
 reth-interfaces = { path = "../interfaces" }
 reth-db = { path = "../storage/db" }
+reth-codecs = { path = "../storage/codecs" }
 reth-provider = { path = "../storage/provider" }
 reth-metrics-derive = { path = "../metrics/metrics-derive" }
 

crates/stages/src/stages/hashing_account.rs

@@ -1,11 +1,12 @@
 use crate::{ExecInput, ExecOutput, Stage, StageError, StageId, UnwindInput, UnwindOutput};
+use reth_codecs::Compact;
 use reth_db::{
     cursor::{DbCursorRO, DbCursorRW},
     database::Database,
     tables,
     transaction::{DbTx, DbTxMut},
 };
-use reth_primitives::keccak256;
+use reth_primitives::{keccak256, AccountHashingCheckpoint};
 use reth_provider::Transaction;
 use std::{collections::BTreeMap, fmt::Debug, ops::Range};
 use tracing::*;
@@ -30,6 +31,43 @@ impl Default for AccountHashingStage {
     }
 }
 
+impl AccountHashingStage {
+    /// Saves the hashing progress
+    pub fn save_checkpoint<DB: Database>(
+        &mut self,
+        tx: &Transaction<'_, DB>,
+        checkpoint: AccountHashingCheckpoint,
+    ) -> Result<(), StageError> {
+        debug!(target: "sync::stages::account_hashing::exec", checkpoint = ?checkpoint, "Saving inner account hashing checkpoint");
+
+        let mut buf = vec![];
+        checkpoint.to_compact(&mut buf);
+
+        Ok(tx.put::<tables::SyncStageProgress>(ACCOUNT_HASHING.0.into(), buf)?)
+    }
+
+    /// Gets the hashing progress
+    pub fn get_checkpoint<DB: Database>(
+        &self,
+        tx: &Transaction<'_, DB>,
+    ) -> Result<AccountHashingCheckpoint, StageError> {
+        let buf =
+            tx.get::<tables::SyncStageProgress>(ACCOUNT_HASHING.0.into())?.unwrap_or_default();
+
+        if buf.is_empty() {
+            return Ok(AccountHashingCheckpoint::default())
+        }
+
+        let (checkpoint, _) = AccountHashingCheckpoint::from_compact(&buf, buf.len());
+
+        if checkpoint.address.is_some() {
+            debug!(target: "sync::stages::account_hashing::exec", checkpoint = ?checkpoint, "Continuing inner account hashing checkpoint");
+        }
+
+        Ok(checkpoint)
+    }
+}
+
 #[derive(Clone, Debug)]
 /// `SeedOpts` provides configuration parameters for calling `AccountHashingStage::seed`
 /// in unit tests or benchmarks to generate an initial database state for running the
@@ -137,43 +175,58 @@ impl<DB: Database> Stage<DB> for AccountHashingStage {
         // AccountHashing table. Also, if we start from genesis, we need to hash from scratch, as
         // genesis accounts are not in changeset.
         if to_transition - from_transition > self.clean_threshold || stage_progress == 0 {
-            // clear table, load all accounts and hash it
-            tx.clear::<tables::HashedAccount>()?;
-            tx.commit()?;
-
-            let mut first_key = None;
-            loop {
-                let next_key = {
-                    let mut accounts = tx.cursor_read::<tables::PlainAccountState>()?;
-
-                    let hashed_batch = accounts
-                        .walk(first_key)?
-                        .take(self.commit_threshold as usize)
-                        .map(|res| res.map(|(address, account)| (keccak256(address), account)))
-                        .collect::<Result<BTreeMap<_, _>, _>>()?;
-
-                    let mut hashed_account_cursor = tx.cursor_write::<tables::HashedAccount>()?;
-
-                    // iterate and put presorted hashed accounts
-                    if first_key.is_none() {
-                        hashed_batch
-                            .into_iter()
-                            .try_for_each(|(k, v)| hashed_account_cursor.append(k, v))?;
-                    } else {
-                        hashed_batch
-                            .into_iter()
-                            .try_for_each(|(k, v)| hashed_account_cursor.insert(k, v))?;
-                    }
+            let mut checkpoint = self.get_checkpoint(tx)?;
+
+            if checkpoint.address.is_none() ||
+                // Checkpoint is no longer valid if the range of transitions changed. 
+                // An already hashed account may have been changed with the new range, and therefore should be hashed again. 
+                checkpoint.to != to_transition ||
+                checkpoint.from != from_transition
+            {
+                // clear table, load all accounts and hash it
+                tx.clear::<tables::HashedAccount>()?;
+
+                checkpoint = AccountHashingCheckpoint::default();
+                self.save_checkpoint(tx, checkpoint)?;
+            }
 
-                    // next key of iterator
-                    accounts.next()?
-                };
-                tx.commit()?;
-                if let Some((next_key, _)) = next_key {
-                    first_key = Some(next_key);
-                    continue
+            let start_address = checkpoint.address.take();
+            let next_address = {
+                let mut accounts = tx.cursor_read::<tables::PlainAccountState>()?;
+
+                let hashed_batch = accounts
+                    .walk(start_address)?
+                    .take(self.commit_threshold as usize)
+                    .map(|res| res.map(|(address, account)| (keccak256(address), account)))
+                    .collect::<Result<BTreeMap<_, _>, _>>()?;
+
+                let mut hashed_account_cursor = tx.cursor_write::<tables::HashedAccount>()?;
+
+                // iterate and put presorted hashed accounts
+                if start_address.is_none() {
+                    hashed_batch
+                        .into_iter()
+                        .try_for_each(|(k, v)| hashed_account_cursor.append(k, v))?;
+                } else {
+                    hashed_batch
+                        .into_iter()
+                        .try_for_each(|(k, v)| hashed_account_cursor.insert(k, v))?;
                 }
-                break
+
+                // next key of iterator
+                accounts.next()?
+            };
+
+            if let Some((next_address, _)) = &next_address {
+                checkpoint.address = Some(*next_address);
+                checkpoint.from = from_transition;
+                checkpoint.to = to_transition;
+            }
+
+            self.save_checkpoint(tx, checkpoint)?;
+
+            if next_address.is_some() {
+                return Ok(ExecOutput { stage_progress, done: false })
             }
         } else {
             // Aggregate all transition changesets and and make list of account that have been