bench: total difficulty (#1215)

crates/stages/benches/criterion.rs

@@ -5,13 +5,13 @@ use criterion::{
 use reth_db::mdbx::{Env, WriteMap};
 use reth_primitives::H256;
 use reth_stages::{
-    stages::{SenderRecoveryStage, TransactionLookupStage},
+    stages::{SenderRecoveryStage, TotalDifficultyStage, TransactionLookupStage},
     test_utils::TestTransaction,
     ExecInput, Stage, StageId, UnwindInput,
 };
 use std::path::{Path, PathBuf};
 
-criterion_group!(benches, tx_lookup, senders);
+criterion_group!(benches, tx_lookup, senders, total_difficulty);
 criterion_main!(benches);
 
 fn senders(c: &mut Criterion) {
@@ -27,7 +27,7 @@ fn senders(c: &mut Criterion) {
         stage.batch_size = batch;
         stage.commit_threshold = num_blocks;
         let label = format!("SendersRecovery-batch-{}", batch);
-        measure_stage(&mut group, stage, num_blocks - 1 /* why do we need - 1 here? */, label);
+        measure_stage(&mut group, stage, num_blocks, label);
     }
 }
 
@@ -43,6 +43,18 @@ fn tx_lookup(c: &mut Criterion) {
     measure_stage(&mut group, stage, num_blocks, "TransactionLookup".to_string());
 }
 
+fn total_difficulty(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Stages");
+    group.measurement_time(std::time::Duration::from_millis(2000));
+    group.warm_up_time(std::time::Duration::from_millis(2000));
+    // don't need to run each stage for that many times
+    group.sample_size(10);
+
+    let num_blocks = 10_000;
+    let stage = TotalDifficultyStage::default();
+    measure_stage(&mut group, stage, num_blocks, "TotalDifficulty".to_string());
+}
+
 fn measure_stage<S: Clone + Default + Stage<Env<WriteMap>>>(
     group: &mut BenchmarkGroup<WallTime>,
     stage: S,
@@ -94,9 +106,23 @@ fn txs_testdata(num_blocks: usize) -> PathBuf {
         let tx = TestTransaction::new(&path);
 
         // This takes a while because it does sig recovery internally
-        let blocks = random_block_range(0..num_blocks as u64, H256::zero(), txs_range);
+        let blocks = random_block_range(0..num_blocks as u64 + 1, H256::zero(), txs_range);
+
+        // insert all blocks
         tx.insert_blocks(blocks.iter(), None).unwrap();
-        tx.inner().commit().unwrap();
+
+        // // initialize TD
+        use reth_db::{
+            cursor::DbCursorRO,
+            tables,
+            transaction::{DbTx, DbTxMut},
+        };
+        tx.commit(|tx| {
+            let (head, _) =
+                tx.cursor_read::<tables::Headers>()?.first()?.unwrap_or_default().into();
+            tx.put::<tables::HeaderTD>(head, reth_primitives::U256::from(0).into())
+        })
+        .unwrap();
     }
 
     path

crates/stages/src/stages/total_difficulty.rs

@@ -19,7 +19,7 @@ const TOTAL_DIFFICULTY: StageId = StageId("TotalDifficulty");
 /// This stage walks over inserted headers and computes total difficulty
 /// at each block. The entries are inserted into [`HeaderTD`][reth_db::tables::HeaderTD]
 /// table.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct TotalDifficultyStage {
     /// The chain specification.
     pub chain_spec: ChainSpec,