feat: add reth db snapshot transactions | receipts commands

bin/reth/src/db/snapshots/bench.rs

@@ -4,7 +4,7 @@ use reth_primitives::{
     ChainSpec, SnapshotSegment,
 };
 use reth_provider::{DatabaseProviderRO, ProviderFactory};
-use std::{sync::Arc, time::Instant};
+use std::{fmt::Debug, sync::Arc, time::Instant};
 
 #[derive(Debug)]
 pub(crate) enum BenchKind {
@@ -14,7 +14,7 @@ pub(crate) enum BenchKind {
     RandomHash,
 }
 
-pub(crate) fn bench<F1, F2>(
+pub(crate) fn bench<F1, F2, R>(
     bench_kind: BenchKind,
     db: (DatabaseEnvRO, Arc<ChainSpec>),
     segment: SnapshotSegment,
@@ -24,28 +24,34 @@ pub(crate) fn bench<F1, F2>(
     database_method: F2,
 ) -> eyre::Result<()>
 where
-    F1: FnMut() -> eyre::Result<()>,
-    F2: Fn(DatabaseProviderRO<'_, DatabaseEnvRO>) -> eyre::Result<()>,
+    F1: FnMut() -> eyre::Result<R>,
+    F2: Fn(DatabaseProviderRO<'_, DatabaseEnvRO>) -> eyre::Result<R>,
+    R: Debug + PartialEq,
 {
     let (db, chain) = db;
 
     println!();
     println!("############");
     println!("## [{segment:?}] [{compression:?}] [{filters:?}] [{bench_kind:?}]");
-    {
+    let snap_result = {
         let start = Instant::now();
-        snapshot_method()?;
+        let result = snapshot_method()?;
         let end = start.elapsed().as_micros();
         println!("# snapshot {bench_kind:?} | {end} μs");
-    }
-    {
+        result
+    };
+
+    let db_result = {
         let factory = ProviderFactory::new(db, chain);
         let provider = factory.provider()?;
         let start = Instant::now();
-        database_method(provider)?;
+        let result = database_method(provider)?;
         let end = start.elapsed().as_micros();
         println!("# database {bench_kind:?} | {end} μs");
-    }
+        result
+    };
+
+    assert_eq!(snap_result, db_result);
 
     Ok(())
 }

bin/reth/src/db/snapshots/headers.rs

@@ -2,23 +2,22 @@ use super::{
     bench::{bench, BenchKind},
     Command,
 };
-use crate::utils::DbTool;
 use rand::{seq::SliceRandom, Rng};
-use reth_db::{database::Database, open_db_read_only, table::Decompress, DatabaseEnvRO};
+use reth_db::{database::Database, open_db_read_only, table::Decompress};
 use reth_interfaces::db::LogLevel;
 use reth_nippy_jar::NippyJar;
 use reth_primitives::{
     snapshot::{Compression, Filters, InclusionFilter, PerfectHashingFunction},
     ChainSpec, Header, SnapshotSegment,
 };
-use reth_provider::{HeaderProvider, ProviderError, ProviderFactory};
+use reth_provider::{DatabaseProviderRO, HeaderProvider, ProviderError, ProviderFactory};
 use reth_snapshot::segments::{get_snapshot_segment_file_name, Headers, Segment};
 use std::{path::Path, sync::Arc};
 
 impl Command {
-    pub(crate) fn generate_headers_snapshot(
+    pub(crate) fn generate_headers_snapshot<DB: Database>(
         &self,
-        tool: &DbTool<'_, DatabaseEnvRO>,
+        provider: &DatabaseProviderRO<'_, DB>,
         compression: Compression,
         inclusion_filter: InclusionFilter,
         phf: PerfectHashingFunction,
@@ -31,7 +30,7 @@ impl Command {
                 Filters::WithoutFilters
             },
         );
-        segment.snapshot(&tool.db.tx()?, self.from..=(self.from + self.block_interval - 1))?;
+        segment.snapshot::<DB>(provider, self.from..=(self.from + self.block_interval - 1))?;
 
         Ok(())
     }
@@ -56,7 +55,7 @@ impl Command {
         let mut row_indexes = range.clone().collect::<Vec<_>>();
         let mut rng = rand::thread_rng();
         let mut dictionaries = None;
-        let mut jar = NippyJar::load_without_header(&get_snapshot_segment_file_name(
+        let mut jar = NippyJar::load(&get_snapshot_segment_file_name(
             SnapshotSegment::Headers,
             filters,
             compression,
@@ -114,18 +113,16 @@ impl Command {
                 filters,
                 compression,
                 || {
-                    Header::decompress(
+                    Ok(Header::decompress(
                         cursor
                             .row_by_number_with_cols::<0b01, 2>((num - self.from) as usize)?
                             .ok_or(ProviderError::HeaderNotFound((num as u64).into()))?[0],
-                    )?;
-                    Ok(())
+                    )?)
                 },
                 |provider| {
-                    provider
+                    Ok(provider
                         .header_by_number(num as u64)?
-                        .ok_or(ProviderError::HeaderNotFound((num as u64).into()))?;
-                    Ok(())
+                        .ok_or(ProviderError::HeaderNotFound((num as u64).into()))?)
                 },
             )?;
         }
@@ -154,13 +151,12 @@ impl Command {
 
                     // Might be a false positive, so in the real world we have to validate it
                     assert_eq!(header.hash_slow(), header_hash);
-                    Ok(())
+                    Ok(header)
                 },
                 |provider| {
-                    provider
+                    Ok(provider
                         .header(&header_hash)?
-                        .ok_or(ProviderError::HeaderNotFound(header_hash.into()))?;
-                    Ok(())
+                        .ok_or(ProviderError::HeaderNotFound(header_hash.into()))?)
                 },
             )?;
         }

bin/reth/src/db/snapshots/mod.rs

@@ -1,44 +1,26 @@
-use crate::{db::genesis_value_parser, utils::DbTool};
 use clap::Parser;
 use itertools::Itertools;
-use reth_db::open_db_read_only;
+use reth_db::{open_db_read_only, DatabaseEnvRO};
 use reth_interfaces::db::LogLevel;
 use reth_nippy_jar::{
     compression::{DecoderDictionary, Decompressor},
     NippyJar,
 };
 use reth_primitives::{
-    snapshot::{Compression, InclusionFilter, PerfectHashingFunction},
+    snapshot::{Compression, InclusionFilter, PerfectHashingFunction, SegmentHeader},
     BlockNumber, ChainSpec, SnapshotSegment,
 };
-use reth_provider::providers::SnapshotProvider;
+use reth_provider::{providers::SnapshotProvider, ProviderFactory};
 use std::{path::Path, sync::Arc};
 
 mod bench;
 mod headers;
+mod receipts;
+mod transactions;
 
 #[derive(Parser, Debug)]
 /// Arguments for the `reth db snapshot` command.
 pub struct Command {
-    /// The chain this node is running.
-    ///
-    /// Possible values are either a built-in chain or the path to a chain specification file.
-    ///
-    /// Built-in chains:
-    /// - mainnet
-    /// - goerli
-    /// - sepolia
-    /// - holesky
-    #[arg(
-        long,
-        value_name = "CHAIN_OR_PATH",
-        verbatim_doc_comment,
-        default_value = "mainnet",
-        value_parser = genesis_value_parser,
-        global = true,
-    )]
-    chain: Arc<ChainSpec>,
-
     /// Snapshot segments to generate.
     segments: Vec<SnapshotSegment>,
 
@@ -87,19 +69,33 @@ impl Command {
 
         {
             let db = open_db_read_only(db_path, None)?;
-            let tool = DbTool::new(&db, chain.clone())?;
+            let factory = ProviderFactory::new(db, chain.clone());
+            let provider = factory.provider()?;
 
             if !self.only_bench {
                 for ((mode, compression), phf) in all_combinations.clone() {
                     match mode {
-                        SnapshotSegment::Headers => self.generate_headers_snapshot(
-                            &tool,
-                            *compression,
-                            InclusionFilter::Cuckoo,
-                            *phf,
-                        )?,
-                        SnapshotSegment::Transactions => todo!(),
-                        SnapshotSegment::Receipts => todo!(),
+                        SnapshotSegment::Headers => self
+                            .generate_headers_snapshot::<DatabaseEnvRO>(
+                                &provider,
+                                *compression,
+                                InclusionFilter::Cuckoo,
+                                *phf,
+                            )?,
+                        SnapshotSegment::Transactions => self
+                            .generate_transactions_snapshot::<DatabaseEnvRO>(
+                                &provider,
+                                *compression,
+                                InclusionFilter::Cuckoo,
+                                *phf,
+                            )?,
+                        SnapshotSegment::Receipts => self
+                            .generate_receipts_snapshot::<DatabaseEnvRO>(
+                                &provider,
+                                *compression,
+                                InclusionFilter::Cuckoo,
+                                *phf,
+                            )?,
                     }
                 }
             }
@@ -116,8 +112,22 @@ impl Command {
                         InclusionFilter::Cuckoo,
                         *phf,
                     )?,
-                    SnapshotSegment::Transactions => todo!(),
-                    SnapshotSegment::Receipts => todo!(),
+                    SnapshotSegment::Transactions => self.bench_transactions_snapshot(
+                        db_path,
+                        log_level,
+                        chain.clone(),
+                        *compression,
+                        InclusionFilter::Cuckoo,
+                        *phf,
+                    )?,
+                    SnapshotSegment::Receipts => self.bench_receipts_snapshot(
+                        db_path,
+                        log_level,
+                        chain.clone(),
+                        *compression,
+                        InclusionFilter::Cuckoo,
+                        *phf,
+                    )?,
                 }
             }
         }
@@ -129,7 +139,7 @@ impl Command {
     /// [`DecoderDictionary`] and [`Decompressor`] if necessary.
     fn prepare_jar_provider<'a>(
         &self,
-        jar: &'a mut NippyJar,
+        jar: &'a mut NippyJar<SegmentHeader>,
         dictionaries: &'a mut Option<Vec<DecoderDictionary<'_>>>,
     ) -> eyre::Result<(SnapshotProvider<'a>, Vec<Decompressor<'a>>)> {
         let mut decompressors: Vec<Decompressor<'_>> = vec![];
@@ -140,6 +150,6 @@ impl Command {
             }
         }
 
-        Ok((SnapshotProvider { jar: &*jar, jar_start_block: self.from }, decompressors))
+        Ok((SnapshotProvider { jar: &*jar }, decompressors))
     }
 }