feat: search for a snapshot that fulfills a queried BlockHash or TxHash (#5373)

Co-authored-by: Alexey Shekhirin <a.shekhirin@gmail.com>
Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>

crates/storage/provider/src/providers/snapshot/manager.rs

@@ -1,23 +1,45 @@
 use super::{LoadedJar, SnapshotJarProvider};
 use crate::{BlockHashReader, BlockNumReader, HeaderProvider, TransactionsProvider};
 use dashmap::DashMap;
-use reth_interfaces::RethResult;
+use parking_lot::RwLock;
+use reth_db::{
+    codecs::CompactU256,
+    snapshot::{HeaderMask, TransactionMask},
+};
+use reth_interfaces::{provider::ProviderError, RethResult};
 use reth_nippy_jar::NippyJar;
 use reth_primitives::{
-    snapshot::{HighestSnapshots, BLOCKS_PER_SNAPSHOT},
-    Address, BlockHash, BlockHashOrNumber, BlockNumber, ChainInfo, Header, SealedHeader,
-    SnapshotSegment, TransactionMeta, TransactionSigned, TransactionSignedNoHash, TxHash, TxNumber,
-    B256, U256,
+    snapshot::HighestSnapshots, Address, BlockHash, BlockHashOrNumber, BlockNumber, ChainInfo,
+    Header, SealedHeader, SnapshotSegment, TransactionMeta, TransactionSigned,
+    TransactionSignedNoHash, TxHash, TxNumber, B256, U256,
+};
+use revm::primitives::HashMap;
+use std::{
+    collections::BTreeMap,
+    ops::{RangeBounds, RangeInclusive},
+    path::{Path, PathBuf},
 };
-use std::{ops::RangeBounds, path::PathBuf};
 use tokio::sync::watch;
 
-/// SnapshotProvider
+/// Alias type for a map that can be queried for transaction/block ranges from a block/transaction
+/// segment respectively. It uses `BlockNumber` to represent the block end of a snapshot range or
+/// `TxNumber` to represent the transaction end of a snapshot range.
+///
+/// Can be in one of the two formats:
+/// - `HashMap<SnapshotSegment, BTreeMap<BlockNumber, RangeInclusive<TxNumber>>>`
+/// - `HashMap<SnapshotSegment, BTreeMap<TxNumber, RangeInclusive<BlockNumber>>>`
+type SegmentRanges = HashMap<SnapshotSegment, BTreeMap<u64, RangeInclusive<u64>>>;
+
+/// [`SnapshotProvider`] manages all existing [`SnapshotJarProvider`].
 #[derive(Debug, Default)]
 pub struct SnapshotProvider {
     /// Maintains a map which allows for concurrent access to different `NippyJars`, over different
     /// segments and ranges.
     map: DashMap<(BlockNumber, SnapshotSegment), LoadedJar>,
+    /// Available snapshot ranges on disk indexed by max blocks.
+    snapshots_block_index: RwLock<SegmentRanges>,
+    /// Available snapshot ranges on disk indexed by max transactions.
+    snapshots_tx_index: RwLock<SegmentRanges>,
     /// Tracks the highest snapshot of every segment.
     highest_tracker: Option<watch::Receiver<Option<HighestSnapshots>>>,
     /// Directory where snapshots are located
@@ -26,8 +48,14 @@ pub struct SnapshotProvider {
 
 impl SnapshotProvider {
     /// Creates a new [`SnapshotProvider`].
-    pub fn new(path: PathBuf) -> Self {
-        Self { map: Default::default(), highest_tracker: None, path }
+    pub fn new(path: impl AsRef<Path>) -> Self {
+        Self {
+            map: Default::default(),
+            snapshots_block_index: Default::default(),
+            snapshots_tx_index: Default::default(),
+            highest_tracker: None,
+            path: path.as_ref().to_path_buf(),
+        }
     }
 
     /// Adds a highest snapshot tracker to the provider
@@ -39,30 +67,137 @@ impl SnapshotProvider {
         self
     }
 
-    /// Gets the provider of the requested segment and range.
-    pub fn get_segment_provider(
+    /// Gets the [`SnapshotJarProvider`] of the requested segment and block.
+    pub fn get_segment_provider_from_block(
         &self,
         segment: SnapshotSegment,
         block: BlockNumber,
-        mut path: Option<PathBuf>,
+        path: Option<&Path>,
     ) -> RethResult<SnapshotJarProvider<'_>> {
-        // TODO this invalidates custom length snapshots.
-        let snapshot = block / BLOCKS_PER_SNAPSHOT;
-        let key = (snapshot, segment);
+        self.get_segment_provider(
+            segment,
+            || self.get_segment_ranges_from_block(segment, block),
+            path,
+        )?
+        .ok_or_else(|| ProviderError::MissingSnapshotBlock(segment, block).into())
+    }
 
+    /// Gets the [`SnapshotJarProvider`] of the requested segment and transaction.
+    pub fn get_segment_provider_from_transaction(
+        &self,
+        segment: SnapshotSegment,
+        tx: TxNumber,
+        path: Option<&Path>,
+    ) -> RethResult<SnapshotJarProvider<'_>> {
+        self.get_segment_provider(
+            segment,
+            || self.get_segment_ranges_from_transaction(segment, tx),
+            path,
+        )?
+        .ok_or_else(|| ProviderError::MissingSnapshotTx(segment, tx).into())
+    }
+
+    /// Gets the [`SnapshotJarProvider`] of the requested segment and block or transaction.
+    pub fn get_segment_provider(
+        &self,
+        segment: SnapshotSegment,
+        fn_ranges: impl Fn() -> Option<(RangeInclusive<BlockNumber>, RangeInclusive<TxNumber>)>,
+        path: Option<&Path>,
+    ) -> RethResult<Option<SnapshotJarProvider<'_>>> {
+        // If we have a path, then get the block range and transaction range from its name.
+        // Otherwise, check `self.available_snapshots`
+        let snapshot_ranges = match path {
+            Some(path) => {
+                SnapshotSegment::parse_filename(path.file_name().ok_or_else(|| {
+                    ProviderError::MissingSnapshotPath(segment, path.to_path_buf())
+                })?)
+                .and_then(|(parsed_segment, block_range, tx_range)| {
+                    if parsed_segment == segment {
+                        return Some((block_range, tx_range));
+                    }
+                    None
+                })
+            }
+            None => fn_ranges(),
+        };
+
+        // Return cached `LoadedJar` or insert it for the first time, and then, return it.
+        if let Some((block_range, tx_range)) = snapshot_ranges {
+            return Ok(Some(self.get_or_create_jar_provider(segment, &block_range, &tx_range)?));
+        }
+
+        Ok(None)
+    }
+
+    /// Given a segment, block range and transaction range it returns a cached
+    /// [`SnapshotJarProvider`]. TODO: we should check the size and pop N if there's too many.
+    fn get_or_create_jar_provider(
+        &self,
+        segment: SnapshotSegment,
+        block_range: &RangeInclusive<u64>,
+        tx_range: &RangeInclusive<u64>,
+    ) -> Result<SnapshotJarProvider<'_>, reth_interfaces::RethError> {
+        let key = (*block_range.end(), segment);
         if let Some(jar) = self.map.get(&key) {
-            return Ok(jar.into())
-        }
-
-        if let Some(path) = &path {
-            self.map.insert(key, LoadedJar::new(NippyJar::load(path)?)?);
+            Ok(jar.into())
         } else {
-            path = Some(self.path.join(segment.filename(
-                &((snapshot * BLOCKS_PER_SNAPSHOT)..=((snapshot + 1) * BLOCKS_PER_SNAPSHOT - 1)),
-            )));
+            self.map.insert(
+                key,
+                LoadedJar::new(NippyJar::load(
+                    &self.path.join(segment.filename(block_range, tx_range)),
+                )?)?,
+            );
+            Ok(self.map.get(&key).expect("qed").into())
         }
+    }
 
-        self.get_segment_provider(segment, block, path)
+    /// Gets a snapshot segment's block range and transaction range from the provider inner block
+    /// index.
+    fn get_segment_ranges_from_block(
+        &self,
+        segment: SnapshotSegment,
+        block: u64,
+    ) -> Option<(RangeInclusive<BlockNumber>, RangeInclusive<TxNumber>)> {
+        let snapshots = self.snapshots_block_index.read();
+        let segment_snapshots = snapshots.get(&segment)?;
+
+        // It's more probable that the request comes from a newer block height, so we iterate
+        // the snapshots in reverse.
+        let mut snapshots_rev_iter = segment_snapshots.iter().rev().peekable();
+
+        while let Some((block_end, tx_range)) = snapshots_rev_iter.next() {
+            // `unwrap_or(0) is safe here as it sets block_start to 0 if the iterator is empty,
+            // indicating the lowest height snapshot has been reached.
+            let block_start =
+                snapshots_rev_iter.peek().map(|(block_end, _)| *block_end + 1).unwrap_or(0);
+            if block_start <= block {
+                return Some((block_start..=*block_end, tx_range.clone()));
+            }
+        }
+        None
+    }
+
+    /// Gets a snapshot segment's block range and transaction range from the provider inner
+    /// transaction index.
+    fn get_segment_ranges_from_transaction(
+        &self,
+        segment: SnapshotSegment,
+        tx: u64,
+    ) -> Option<(RangeInclusive<BlockNumber>, RangeInclusive<TxNumber>)> {
+        let snapshots = self.snapshots_tx_index.read();
+        let segment_snapshots = snapshots.get(&segment)?;
+
+        // It's more probable that the request comes from a newer tx height, so we iterate
+        // the snapshots in reverse.
+        let mut snapshots_rev_iter = segment_snapshots.iter().rev().peekable();
+
+        while let Some((tx_end, block_range)) = snapshots_rev_iter.next() {
+            let tx_start = snapshots_rev_iter.peek().map(|(tx_end, _)| *tx_end + 1).unwrap_or(0);
+            if tx_start <= tx {
+                return Some((block_range.clone(), tx_start..=*tx_end));
+            }
+        }
+        None
     }
 
     /// Gets the highest snapshot if it exists for a snapshot segment.
@@ -71,23 +206,72 @@ impl SnapshotProvider {
             .as_ref()
             .and_then(|tracker| tracker.borrow().and_then(|highest| highest.highest(segment)))
     }
+
+    /// Iterates through segment snapshots in reverse order, executing a function until it returns
+    /// some object. Useful for finding objects by [`TxHash`] or [`BlockHash`].
+    pub fn find_snapshot<T>(
+        &self,
+        segment: SnapshotSegment,
+        func: impl Fn(SnapshotJarProvider<'_>) -> RethResult<Option<T>>,
+    ) -> RethResult<Option<T>> {
+        let snapshots = self.snapshots_block_index.read();
+        if let Some(segment_snapshots) = snapshots.get(&segment) {
+            // It's more probable that the request comes from a newer block height, so we iterate
+            // the snapshots in reverse.
+            let mut snapshots_rev_iter = segment_snapshots.iter().rev().peekable();
+
+            while let Some((block_end, tx_range)) = snapshots_rev_iter.next() {
+                // `unwrap_or(0) is safe here as it sets block_start to 0 if the iterator
+                // is empty, indicating the lowest height snapshot has been reached.
+                let block_start =
+                    snapshots_rev_iter.peek().map(|(block_end, _)| *block_end + 1).unwrap_or(0);
+
+                if let Some(res) = func(self.get_or_create_jar_provider(
+                    segment,
+                    &(block_start..=*block_end),
+                    tx_range,
+                )?)? {
+                    return Ok(Some(res))
+                }
+            }
+        }
+
+        Ok(None)
+    }
 }
 
 impl HeaderProvider for SnapshotProvider {
-    fn header(&self, _block_hash: &BlockHash) -> RethResult<Option<Header>> {
-        todo!()
+    fn header(&self, block_hash: &BlockHash) -> RethResult<Option<Header>> {
+        self.find_snapshot(SnapshotSegment::Headers, |jar_provider| {
+            Ok(jar_provider
+                .cursor()?
+                .get_two::<HeaderMask<Header, BlockHash>>(block_hash.into())?
+                .and_then(|(header, hash)| {
+                    if &hash == block_hash {
+                        return Some(header)
+                    }
+                    None
+                }))
+        })
     }
 
     fn header_by_number(&self, num: BlockNumber) -> RethResult<Option<Header>> {
-        self.get_segment_provider(SnapshotSegment::Headers, num, None)?.header_by_number(num)
+        self.get_segment_provider_from_block(SnapshotSegment::Headers, num, None)?
+            .header_by_number(num)
     }
 
-    fn header_td(&self, _block_hash: &BlockHash) -> RethResult<Option<U256>> {
-        todo!()
+    fn header_td(&self, block_hash: &BlockHash) -> RethResult<Option<U256>> {
+        self.find_snapshot(SnapshotSegment::Headers, |jar_provider| {
+            Ok(jar_provider
+                .cursor()?
+                .get_two::<HeaderMask<CompactU256, BlockHash>>(block_hash.into())?
+                .and_then(|(td, hash)| (&hash == block_hash).then_some(td.0)))
+        })
     }
 
-    fn header_td_by_number(&self, _number: BlockNumber) -> RethResult<Option<U256>> {
-        todo!();
+    fn header_td_by_number(&self, num: BlockNumber) -> RethResult<Option<U256>> {
+        self.get_segment_provider_from_block(SnapshotSegment::Headers, num, None)?
+            .header_td_by_number(num)
     }
 
     fn headers_range(&self, _range: impl RangeBounds<BlockNumber>) -> RethResult<Vec<Header>> {
@@ -101,14 +285,15 @@ impl HeaderProvider for SnapshotProvider {
         todo!();
     }
 
-    fn sealed_header(&self, _number: BlockNumber) -> RethResult<Option<SealedHeader>> {
-        todo!();
+    fn sealed_header(&self, num: BlockNumber) -> RethResult<Option<SealedHeader>> {
+        self.get_segment_provider_from_block(SnapshotSegment::Headers, num, None)?
+            .sealed_header(num)
     }
 }
 
 impl BlockHashReader for SnapshotProvider {
-    fn block_hash(&self, _number: u64) -> RethResult<Option<B256>> {
-        todo!()
+    fn block_hash(&self, num: u64) -> RethResult<Option<B256>> {
+        self.get_segment_provider_from_block(SnapshotSegment::Headers, num, None)?.block_hash(num)
     }
 
     fn canonical_hashes_range(
@@ -139,26 +324,42 @@ impl BlockNumReader for SnapshotProvider {
 }
 
 impl TransactionsProvider for SnapshotProvider {
-    fn transaction_id(&self, _tx_hash: TxHash) -> RethResult<Option<TxNumber>> {
-        todo!()
+    fn transaction_id(&self, tx_hash: TxHash) -> RethResult<Option<TxNumber>> {
+        self.find_snapshot(SnapshotSegment::Transactions, |jar_provider| {
+            let mut cursor = jar_provider.cursor()?;
+            if cursor
+                .get_one::<TransactionMask<TransactionSignedNoHash>>((&tx_hash).into())?
+                .and_then(|tx| (tx.hash() == tx_hash).then_some(tx))
+                .is_some()
+            {
+                Ok(Some(cursor.number()))
+            } else {
+                Ok(None)
+            }
+        })
     }
 
     fn transaction_by_id(&self, num: TxNumber) -> RethResult<Option<TransactionSigned>> {
-        // TODO `num` is provided after checking the index
-        let block_num = num;
-        self.get_segment_provider(SnapshotSegment::Transactions, block_num, None)?
+        self.get_segment_provider_from_transaction(SnapshotSegment::Transactions, num, None)?
             .transaction_by_id(num)
     }
 
     fn transaction_by_id_no_hash(
         &self,
-        _id: TxNumber,
+        num: TxNumber,
     ) -> RethResult<Option<TransactionSignedNoHash>> {
-        todo!()
+        self.get_segment_provider_from_transaction(SnapshotSegment::Transactions, num, None)?
+            .transaction_by_id_no_hash(num)
     }
 
-    fn transaction_by_hash(&self, _hash: TxHash) -> RethResult<Option<TransactionSigned>> {
-        todo!()
+    fn transaction_by_hash(&self, hash: TxHash) -> RethResult<Option<TransactionSigned>> {
+        self.find_snapshot(SnapshotSegment::Transactions, |jar_provider| {
+            Ok(jar_provider
+                .cursor()?
+                .get_one::<TransactionMask<TransactionSignedNoHash>>((&hash).into())?
+                .map(|tx| tx.with_hash())
+                .and_then(|tx| (tx.hash_ref() == &hash).then_some(tx)))
+        })
     }
 
     fn transaction_by_hash_with_meta(
@@ -197,7 +398,7 @@ impl TransactionsProvider for SnapshotProvider {
         todo!()
     }
 
-    fn transaction_sender(&self, _id: TxNumber) -> RethResult<Option<Address>> {
-        todo!()
+    fn transaction_sender(&self, id: TxNumber) -> RethResult<Option<Address>> {
+        Ok(self.transaction_by_id_no_hash(id)?.and_then(|tx| tx.recover_signer()))
     }
 }

crates/storage/provider/src/providers/snapshot/mod.rs

@@ -66,7 +66,8 @@ mod test {
         // Data sources
         let db = create_test_rw_db();
         let factory = ProviderFactory::new(&db, MAINNET.clone());
-        let snap_file = tempfile::NamedTempFile::new().unwrap();
+        let snap_path = tempfile::tempdir().unwrap();
+        let snap_file = snap_path.path().join(SnapshotSegment::Headers.filename(&range, &range));
 
         // Setup data
         let mut headers = random_header_range(
@@ -96,7 +97,7 @@ mod test {
             let with_compression = true;
             let with_filter = true;
 
-            let mut nippy_jar = NippyJar::new(3, snap_file.path(), segment_header);
+            let mut nippy_jar = NippyJar::new(3, snap_file.as_path(), segment_header);
 
             if with_compression {
                 nippy_jar = nippy_jar.with_zstd(false, 0);
@@ -134,9 +135,9 @@ mod test {
         // Use providers to query Header data and compare if it matches
         {
             let db_provider = factory.provider().unwrap();
-            let manager = SnapshotProvider::default();
+            let manager = SnapshotProvider::new(snap_path.path());
             let jar_provider = manager
-                .get_segment_provider(SnapshotSegment::Headers, 0, Some(snap_file.path().into()))
+                .get_segment_provider_from_block(SnapshotSegment::Headers, 0, Some(&snap_file))
                 .unwrap();
 
             assert!(!headers.is_empty());