perf: better state retrieval (#3221)

crates/rpc/rpc/src/eth/api/mod.rs

@@ -177,21 +177,14 @@ impl<Provider, Pool, Network> EthApi<Provider, Pool, Network>
 where
     Provider: BlockProviderIdExt + StateProviderFactory + EvmEnvProvider + 'static,
 {
-    fn convert_block_number(&self, num: BlockNumberOrTag) -> Result<Option<u64>> {
-        self.provider().convert_block_number(num)
-    }
-
     /// Returns the state at the given [BlockId] enum.
     pub fn state_at_block_id(&self, at: BlockId) -> EthResult<StateProviderBox<'_>> {
-        match at {
+        Ok(self.provider().state_by_block_id(at)?)
-            BlockId::Hash(hash) => Ok(self.state_at_hash(hash.into())?),
-            BlockId::Number(num) => {
-                self.state_at_block_number(num)?.ok_or(EthApiError::UnknownBlockNumber)
-            }
-        }
     }
 
     /// Returns the state at the given [BlockId] enum or the latest.
+    ///
+    /// Convenience function to interprets `None` as `BlockId::Number(BlockNumberOrTag::Latest)`
     pub fn state_at_block_id_or_latest(
         &self,
         block_id: Option<BlockId>,
@@ -203,33 +196,11 @@ where
         }
     }
 
-    /// Returns the state at the given [BlockNumberOrTag] enum
-    ///
-    /// Returns `None` if no state available.
-    pub fn state_at_block_number(
-        &self,
-        num: BlockNumberOrTag,
-    ) -> Result<Option<StateProviderBox<'_>>> {
-        if let Some(number) = self.convert_block_number(num)? {
-            self.state_at_number(number).map(Some)
-        } else {
-            Ok(None)
-        }
-    }
-
     /// Returns the state at the given block number
     pub fn state_at_hash(&self, block_hash: H256) -> Result<StateProviderBox<'_>> {
         self.provider().history_by_block_hash(block_hash)
     }
 
-    /// Returns the state at the given block number
-    pub fn state_at_number(&self, block_number: u64) -> Result<StateProviderBox<'_>> {
-        match self.convert_block_number(BlockNumberOrTag::Latest)? {
-            Some(num) if num == block_number => self.latest_state(),
-            _ => self.provider().history_by_block_number(block_number),
-        }
-    }
-
     /// Returns the _latest_ state
     pub fn latest_state(&self) -> Result<StateProviderBox<'_>> {
         self.provider().latest()

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

@@ -409,6 +409,65 @@ where
         self.database.latest()
     }
 
+    /// Returns a [StateProviderBox] indexed by the given [BlockId].
+    fn state_by_block_id(&self, block_id: BlockId) -> Result<StateProviderBox<'_>> {
+        match block_id {
+            BlockId::Number(block_number) => self.state_by_block_number_or_tag(block_number),
+            BlockId::Hash(rpc_block_hash) => {
+                let block_hash = rpc_block_hash.into();
+                let mut state = self.history_by_block_hash(block_hash);
+
+                // we failed to get the state by hash, from disk, hash block be the pending block
+                if state.is_err() && !rpc_block_hash.require_canonical.unwrap_or(false) {
+                    if let Ok(Some(pending)) = self.pending_state_by_hash(block_hash) {
+                        // we found pending block by hash
+                        state = Ok(pending)
+                    }
+                }
+
+                state
+            }
+        }
+    }
+
+    /// Returns a [StateProviderBox] indexed by the given block number or tag.
+    fn state_by_block_number_or_tag(
+        &self,
+        number_or_tag: BlockNumberOrTag,
+    ) -> Result<StateProviderBox<'_>> {
+        match number_or_tag {
+            BlockNumberOrTag::Latest => self.latest(),
+            BlockNumberOrTag::Finalized => {
+                // we can only get the finalized state by hash, not by num
+                let hash = match self.finalized_block_hash()? {
+                    Some(hash) => hash,
+                    None => return Err(ProviderError::FinalizedBlockNotFound.into()),
+                };
+
+                self.state_by_block_hash(hash)
+            }
+            BlockNumberOrTag::Safe => {
+                // we can only get the safe state by hash, not by num
+                let hash = match self.safe_block_hash()? {
+                    Some(hash) => hash,
+                    None => return Err(ProviderError::SafeBlockNotFound.into()),
+                };
+
+                self.state_by_block_hash(hash)
+            }
+            BlockNumberOrTag::Earliest => self.history_by_block_number(0),
+            BlockNumberOrTag::Pending => self.pending(),
+            BlockNumberOrTag::Number(num) => {
+                let mut state = self.history_by_block_number(num);
+                if state.is_err() && num == self.chain_info.get_canonical_block_number() + 1 {
+                    // we don't have the block on disk yet but the number is the pending block
+                    state = self.pending();
+                }
+                state
+            }
+        }
+    }
+
     fn history_by_block_number(&self, block_number: BlockNumber) -> Result<StateProviderBox<'_>> {
         trace!(target: "providers::blockchain", ?block_number, "Getting history by block number");
         self.ensure_canonical_block(block_number)?;
@@ -443,6 +502,13 @@ where
         self.latest()
     }
 
+    fn pending_state_by_hash(&self, block_hash: H256) -> Result<Option<StateProviderBox<'_>>> {
+        if let Some(state) = self.tree.find_pending_state_provider(block_hash) {
+            return Ok(Some(self.pending_with_provider(state)?))
+        }
+        Ok(None)
+    }
+
     fn pending_with_provider(
         &self,
         post_state_data: Box<dyn PostStateDataProvider>,

crates/storage/provider/src/test_utils/mock.rs

@@ -462,6 +462,10 @@ impl StateProviderFactory for MockEthProvider {
         todo!()
     }
 
+    fn pending_state_by_hash(&self, _block_hash: H256) -> Result<Option<StateProviderBox<'_>>> {
+        todo!()
+    }
+
     fn pending_with_provider<'a>(
         &'a self,
         _post_state_data: Box<dyn PostStateDataProvider + 'a>,
@@ -491,6 +495,10 @@ impl StateProviderFactory for Arc<MockEthProvider> {
         todo!()
     }
 
+    fn pending_state_by_hash(&self, _block_hash: H256) -> Result<Option<StateProviderBox<'_>>> {
+        todo!()
+    }
+
     fn pending_with_provider<'a>(
         &'a self,
         _post_state_data: Box<dyn PostStateDataProvider + 'a>,

crates/storage/provider/src/test_utils/noop.rs

@@ -303,6 +303,10 @@ impl StateProviderFactory for NoopProvider {
         Ok(Box::new(*self))
     }
 
+    fn pending_state_by_hash(&self, _block_hash: H256) -> Result<Option<StateProviderBox<'_>>> {
+        Ok(Some(Box::new(*self)))
+    }
+
     fn pending_with_provider<'a>(
         &'a self,
         _post_state_data: Box<dyn crate::PostStateDataProvider + 'a>,

crates/storage/provider/src/traits/state.rs

@@ -103,13 +103,13 @@ pub trait StateProviderFactory: BlockIdProvider + Send + Sync {
     /// Returns a [StateProvider] indexed by the given [BlockId].
     fn state_by_block_id(&self, block_id: BlockId) -> Result<StateProviderBox<'_>> {
         match block_id {
-            BlockId::Number(block_number) => self.history_by_block_number_or_tag(block_number),
+            BlockId::Number(block_number) => self.state_by_block_number_or_tag(block_number),
             BlockId::Hash(block_hash) => self.history_by_block_hash(block_hash.into()),
         }
     }
 
     /// Returns a [StateProvider] indexed by the given block number or tag.
-    fn history_by_block_number_or_tag(
+    fn state_by_block_number_or_tag(
         &self,
         number_or_tag: BlockNumberOrTag,
     ) -> Result<StateProviderBox<'_>> {
@@ -161,6 +161,13 @@ pub trait StateProviderFactory: BlockIdProvider + Send + Sync {
     /// If there's no `pending` block, then this is equal to [StateProviderFactory::latest]
     fn pending(&self) -> Result<StateProviderBox<'_>>;
 
+    /// Storage provider for pending state for the given block hash.
+    ///
+    /// Represents the state at the block that extends the canonical chain.
+    ///
+    /// If the block couldn't be found, returns `None`.
+    fn pending_state_by_hash(&self, block_hash: H256) -> Result<Option<StateProviderBox<'_>>>;
+
     /// Return a [StateProvider] that contains post state data provider.
     /// Used to inspect or execute transaction on the pending state.
     fn pending_with_provider(