feat(storage): account for pruned account/storage history (#4092)

crates/interfaces/src/provider.rs

@@ -94,4 +94,6 @@ pub enum ProviderError {
         /// Block hash
         block_hash: BlockHash,
     },
+    #[error("State at block #{0} is pruned")]
+    StateAtBlockPruned(BlockNumber),
 }

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

@@ -87,7 +87,7 @@ impl<DB: Database> ProviderFactory<DB> {
     }
 
     /// Storage provider for state at that given block
-    pub fn history_by_block_number(
+    fn state_provider_by_block_number(
         &self,
         mut block_number: BlockNumber,
     ) -> Result<StateProviderBox<'_>> {
@@ -102,30 +102,49 @@ impl<DB: Database> ProviderFactory<DB> {
         // +1 as the changeset that we want is the one that was applied after this block.
         block_number += 1;
 
+        let account_history_prune_checkpoint =
+            provider.get_prune_checkpoint(PrunePart::AccountHistory)?;
+        let storage_history_prune_checkpoint =
+            provider.get_prune_checkpoint(PrunePart::StorageHistory)?;
+
+        let mut state_provider = HistoricalStateProvider::new(provider.into_tx(), block_number);
+
+        // If we pruned account or storage history, we can't return state on every historical block.
+        // Instead, we should cap it at the latest prune checkpoint for corresponding prune part.
+        if let Some(prune_checkpoint) = account_history_prune_checkpoint {
+            state_provider = state_provider.with_lowest_available_account_history_block_number(
+                prune_checkpoint.block_number + 1,
+            );
+        }
+        if let Some(prune_checkpoint) = storage_history_prune_checkpoint {
+            state_provider = state_provider.with_lowest_available_storage_history_block_number(
+                prune_checkpoint.block_number + 1,
+            );
+        }
+
+        Ok(Box::new(state_provider))
+    }
+
+    /// Storage provider for state at that given block
+    pub fn history_by_block_number(
+        &self,
+        block_number: BlockNumber,
+    ) -> Result<StateProviderBox<'_>> {
+        let state_provider = self.state_provider_by_block_number(block_number)?;
         trace!(target: "providers::db", ?block_number, "Returning historical state provider for block number");
-        Ok(Box::new(HistoricalStateProvider::new(provider.into_tx(), block_number)))
+        Ok(state_provider)
     }
 
     /// Storage provider for state at that given block hash
     pub fn history_by_block_hash(&self, block_hash: BlockHash) -> Result<StateProviderBox<'_>> {
-        let provider = self.provider()?;
-
-        let mut block_number = provider
+        let block_number = self
+            .provider()?
             .block_number(block_hash)?
             .ok_or(ProviderError::BlockHashNotFound(block_hash))?;
 
-        if block_number == provider.best_block_number().unwrap_or_default() &&
-            block_number == provider.last_block_number().unwrap_or_default()
-        {
-            return Ok(Box::new(LatestStateProvider::new(provider.into_tx())))
-        }
-
-        // +1 as the changeset that we want is the one that was applied after this block.
-        // as the  changeset contains old values.
-        block_number += 1;
-
-        trace!(target: "providers::db", ?block_hash, "Returning historical state provider for block hash");
-        Ok(Box::new(HistoricalStateProvider::new(provider.into_tx(), block_number)))
+        let state_provider = self.state_provider_by_block_number(block_number)?;
+        trace!(target: "providers::db", ?block_number, "Returning historical state provider for block hash");
+        Ok(state_provider)
     }
 }
 

crates/storage/provider/src/providers/state/historical.rs

@@ -29,24 +29,46 @@ pub struct HistoricalStateProviderRef<'a, 'b, TX: DbTx<'a>> {
     tx: &'b TX,
     /// Block number is main index for the history state of accounts and storages.
     block_number: BlockNumber,
+    /// Lowest blocks at which different parts of the state are available.
+    lowest_available_blocks: LowestAvailableBlocks,
     /// Phantom lifetime `'a`
     _phantom: PhantomData<&'a TX>,
 }
 
+#[derive(Debug, Eq, PartialEq)]
 pub enum HistoryInfo {
     NotYetWritten,
     InChangeset(u64),
     InPlainState,
 }
 
 impl<'a, 'b, TX: DbTx<'a>> HistoricalStateProviderRef<'a, 'b, TX> {
-    /// Create new StateProvider from history transaction number
+    /// Create new StateProvider for historical block number
     pub fn new(tx: &'b TX, block_number: BlockNumber) -> Self {
-        Self { tx, block_number, _phantom: PhantomData {} }
+        Self {
+            tx,
+            block_number,
+            lowest_available_blocks: Default::default(),
+            _phantom: PhantomData {},
+        }
+    }
+
+    /// Create new StateProvider for historical block number and lowest block numbers at which
+    /// account & storage histories are available.
+    pub fn new_with_lowest_available_blocks(
+        tx: &'b TX,
+        block_number: BlockNumber,
+        lowest_available_blocks: LowestAvailableBlocks,
+    ) -> Self {
+        Self { tx, block_number, lowest_available_blocks, _phantom: PhantomData {} }
     }
 
     /// Lookup an account in the AccountHistory table
     pub fn account_history_lookup(&self, address: Address) -> Result<HistoryInfo> {
+        if !self.lowest_available_blocks.is_account_history_available(self.block_number) {
+            return Err(ProviderError::StateAtBlockPruned(self.block_number).into())
+        }
+
         // history key to search IntegerList of block number changesets.
         let history_key = ShardedKey::new(address, self.block_number);
         self.history_info::<tables::AccountHistory, _>(history_key, |key| key.key == address)
@@ -58,6 +80,10 @@ impl<'a, 'b, TX: DbTx<'a>> HistoricalStateProviderRef<'a, 'b, TX> {
         address: Address,
         storage_key: StorageKey,
     ) -> Result<HistoryInfo> {
+        if !self.lowest_available_blocks.is_storage_history_available(self.block_number) {
+            return Err(ProviderError::StateAtBlockPruned(self.block_number).into())
+        }
+
         // history key to search IntegerList of block number changesets.
         let history_key = StorageShardedKey::new(address, storage_key, self.block_number);
         self.history_info::<tables::StorageHistory, _>(history_key, |key| {
@@ -199,29 +225,85 @@ pub struct HistoricalStateProvider<'a, TX: DbTx<'a>> {
     tx: TX,
     /// State at the block number is the main indexer of the state.
     block_number: BlockNumber,
+    /// Lowest blocks at which different parts of the state are available.
+    lowest_available_blocks: LowestAvailableBlocks,
     /// Phantom lifetime `'a`
     _phantom: PhantomData<&'a TX>,
 }
 
 impl<'a, TX: DbTx<'a>> HistoricalStateProvider<'a, TX> {
-    /// Create new StateProvider from history transaction number
+    /// Create new StateProvider for historical block number
     pub fn new(tx: TX, block_number: BlockNumber) -> Self {
-        Self { tx, block_number, _phantom: PhantomData {} }
+        Self {
+            tx,
+            block_number,
+            lowest_available_blocks: Default::default(),
+            _phantom: PhantomData {},
+        }
+    }
+
+    /// Set the lowest block number at which the account history is available.
+    pub fn with_lowest_available_account_history_block_number(
+        mut self,
+        block_number: BlockNumber,
+    ) -> Self {
+        self.lowest_available_blocks.account_history_block_number = Some(block_number);
+        self
+    }
+
+    /// Set the lowest block number at which the storage history is available.
+    pub fn with_lowest_available_storage_history_block_number(
+        mut self,
+        block_number: BlockNumber,
+    ) -> Self {
+        self.lowest_available_blocks.storage_history_block_number = Some(block_number);
+        self
     }
 
     /// Returns a new provider that takes the `TX` as reference
     #[inline(always)]
     fn as_ref<'b>(&'b self) -> HistoricalStateProviderRef<'a, 'b, TX> {
-        HistoricalStateProviderRef::new(&self.tx, self.block_number)
+        HistoricalStateProviderRef::new_with_lowest_available_blocks(
+            &self.tx,
+            self.block_number,
+            self.lowest_available_blocks,
+        )
     }
 }
 
 // Delegates all provider impls to [HistoricalStateProviderRef]
 delegate_provider_impls!(HistoricalStateProvider<'a, TX> where [TX: DbTx<'a>]);
 
+/// Lowest blocks at which different parts of the state are available.
+/// They may be [Some] if pruning is enabled.
+#[derive(Default, Copy, Clone)]
+pub struct LowestAvailableBlocks {
+    /// Lowest block number at which the account history is available. It may not be available if
+    /// [reth_primitives::PrunePart::AccountHistory] was pruned.
+    pub account_history_block_number: Option<BlockNumber>,
+    /// Lowest block number at which the storage history is available. It may not be available if
+    /// [reth_primitives::PrunePart::StorageHistory] was pruned.
+    pub storage_history_block_number: Option<BlockNumber>,
+}
+
+impl LowestAvailableBlocks {
+    /// Check if account history is available at the provided block number, i.e. lowest available
+    /// block number for account history is less than or equal to the provided block number.
+    pub fn is_account_history_available(&self, at: BlockNumber) -> bool {
+        self.account_history_block_number.map(|block_number| block_number <= at).unwrap_or(true)
+    }
+
+    /// Check if storage history is available at the provided block number, i.e. lowest available
+    /// block number for storage history is less than or equal to the provided block number.
+    pub fn is_storage_history_available(&self, at: BlockNumber) -> bool {
+        self.storage_history_block_number.map(|block_number| block_number <= at).unwrap_or(true)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{
+        providers::state::historical::{HistoryInfo, LowestAvailableBlocks},
         AccountReader, HistoricalStateProvider, HistoricalStateProviderRef, StateProvider,
     };
     use reth_db::{
@@ -232,6 +314,7 @@ mod tests {
         transaction::{DbTx, DbTxMut},
         BlockNumberList,
     };
+    use reth_interfaces::provider::ProviderError;
     use reth_primitives::{hex_literal::hex, Account, StorageEntry, H160, H256, U256};
 
     const ADDRESS: H160 = H160(hex!("0000000000000000000000000000000000000001"));
@@ -440,4 +523,61 @@ mod tests {
             Ok(Some(higher_entry_plain.value))
         );
     }
+
+    #[test]
+    fn history_provider_unavailable() {
+        let db = create_test_rw_db();
+        let tx = db.tx().unwrap();
+
+        // provider block_number < lowest available block number,
+        // i.e. state at provider block is pruned
+        let provider = HistoricalStateProviderRef::new_with_lowest_available_blocks(
+            &tx,
+            2,
+            LowestAvailableBlocks {
+                account_history_block_number: Some(3),
+                storage_history_block_number: Some(3),
+            },
+        );
+        assert_eq!(
+            provider.account_history_lookup(ADDRESS),
+            Err(ProviderError::StateAtBlockPruned(provider.block_number).into())
+        );
+        assert_eq!(
+            provider.storage_history_lookup(ADDRESS, STORAGE),
+            Err(ProviderError::StateAtBlockPruned(provider.block_number).into())
+        );
+
+        // provider block_number == lowest available block number,
+        // i.e. state at provider block is available
+        let provider = HistoricalStateProviderRef::new_with_lowest_available_blocks(
+            &tx,
+            2,
+            LowestAvailableBlocks {
+                account_history_block_number: Some(2),
+                storage_history_block_number: Some(2),
+            },
+        );
+        assert_eq!(provider.account_history_lookup(ADDRESS), Ok(HistoryInfo::NotYetWritten));
+        assert_eq!(
+            provider.storage_history_lookup(ADDRESS, STORAGE),
+            Ok(HistoryInfo::NotYetWritten)
+        );
+
+        // provider block_number == lowest available block number,
+        // i.e. state at provider block is available
+        let provider = HistoricalStateProviderRef::new_with_lowest_available_blocks(
+            &tx,
+            2,
+            LowestAvailableBlocks {
+                account_history_block_number: Some(1),
+                storage_history_block_number: Some(1),
+            },
+        );
+        assert_eq!(provider.account_history_lookup(ADDRESS), Ok(HistoryInfo::NotYetWritten));
+        assert_eq!(
+            provider.storage_history_lookup(ADDRESS, STORAGE),
+            Ok(HistoryInfo::NotYetWritten)
+        );
+    }
 }