chore: add StorageLocation to BlockBodyWriter trait (#13266)

paradigmxyz · Dec 10, 2024 · 4c39b98 · 4c39b98
1 parent d97449d
commit 4c39b98
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 41 deletions.
diff --git a/crates/optimism/node/src/node.rs b/crates/optimism/node/src/node.rs
@@ -36,7 +36,7 @@ use reth_payload_builder::{PayloadBuilderHandle, PayloadBuilderService};
 use reth_primitives::{BlockBody, PooledTransactionsElement, TransactionSigned};
 use reth_provider::{
     providers::ChainStorage, BlockBodyReader, BlockBodyWriter, CanonStateSubscriptions,
-    ChainSpecProvider, DBProvider, EthStorage, ProviderResult, ReadBodyInput,
+    ChainSpecProvider, DBProvider, EthStorage, ProviderResult, ReadBodyInput, StorageLocation,
 };
 use reth_rpc_server_types::RethRpcModule;
 use reth_tracing::tracing::{debug, info};
@@ -56,16 +56,18 @@ impl<Provider: DBProvider<Tx: DbTxMut>> BlockBodyWriter<Provider, BlockBody> for
         &self,
         provider: &Provider,
         bodies: Vec<(u64, Option<BlockBody>)>,
+        write_to: StorageLocation,
     ) -> ProviderResult<()> {
-        self.0.write_block_bodies(provider, bodies)
+        self.0.write_block_bodies(provider, bodies, write_to)
     }
 
     fn remove_block_bodies_above(
         &self,
         provider: &Provider,
         block: alloy_primitives::BlockNumber,
+        remove_from: StorageLocation,
     ) -> ProviderResult<()> {
-        self.0.remove_block_bodies_above(provider, block)
+        self.0.remove_block_bodies_above(provider, block, remove_from)
     }
 }
 

diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs
@@ -2893,12 +2893,12 @@ impl<TX: DbTxMut + DbTx + 'static, N: NodeTypesForProvider + 'static> BlockWrite
     fn append_block_bodies(
         &self,
         bodies: Vec<(BlockNumber, Option<BodyTy<N>>)>,
-        write_transactions_to: StorageLocation,
+        write_to: StorageLocation,
     ) -> ProviderResult<()> {
         let Some(from_block) = bodies.first().map(|(block, _)| *block) else { return Ok(()) };
 
         // Initialize writer if we will be writing transactions to staticfiles
-        let mut tx_static_writer = write_transactions_to
+        let mut tx_static_writer = write_to
             .static_files()
             .then(|| {
                 self.static_file_provider.get_writer(from_block, StaticFileSegment::Transactions)
@@ -2909,7 +2909,7 @@ impl<TX: DbTxMut + DbTx + 'static, N: NodeTypesForProvider + 'static> BlockWrite
         let mut tx_block_cursor = self.tx.cursor_write::<tables::TransactionBlocks>()?;
 
         // Initialize cursor if we will be writing transactions to database
-        let mut tx_cursor = write_transactions_to
+        let mut tx_cursor = write_to
             .database()
             .then(|| self.tx.cursor_write::<tables::Transactions<TxTy<N>>>())
             .transpose()?;
@@ -2962,15 +2962,15 @@ impl<TX: DbTxMut + DbTx + 'static, N: NodeTypesForProvider + 'static> BlockWrite
             );
         }
 
-        self.storage.writer().write_block_bodies(self, bodies)?;
+        self.storage.writer().write_block_bodies(self, bodies, write_to)?;
 
         Ok(())
     }
 
     fn remove_blocks_above(
         &self,
         block: BlockNumber,
-        remove_transactions_from: StorageLocation,
+        remove_from: StorageLocation,
     ) -> ProviderResult<()> {
         let mut canonical_headers_cursor = self.tx.cursor_write::<tables::CanonicalHeaders>()?;
         let mut rev_headers = canonical_headers_cursor.walk_back(None)?;
@@ -3010,17 +3010,17 @@ impl<TX: DbTxMut + DbTx + 'static, N: NodeTypesForProvider + 'static> BlockWrite
 
         self.remove::<tables::TransactionSenders>(unwind_tx_from..)?;
 
-        self.remove_bodies_above(block, remove_transactions_from)?;
+        self.remove_bodies_above(block, remove_from)?;
 
         Ok(())
     }
 
     fn remove_bodies_above(
         &self,
         block: BlockNumber,
-        remove_transactions_from: StorageLocation,
+        remove_from: StorageLocation,
     ) -> ProviderResult<()> {
-        self.storage.writer().remove_block_bodies_above(self, block)?;
+        self.storage.writer().remove_block_bodies_above(self, block, remove_from)?;
 
         // First transaction to be removed
         let unwind_tx_from = self
@@ -3032,11 +3032,11 @@ impl<TX: DbTxMut + DbTx + 'static, N: NodeTypesForProvider + 'static> BlockWrite
         self.remove::<tables::BlockBodyIndices>(block + 1..)?;
         self.remove::<tables::TransactionBlocks>(unwind_tx_from..)?;
 
-        if remove_transactions_from.database() {
+        if remove_from.database() {
             self.remove::<tables::Transactions<TxTy<N>>>(unwind_tx_from..)?;
         }
 
-        if remove_transactions_from.static_files() {
+        if remove_from.static_files() {
             let static_file_tx_num = self
                 .static_file_provider
                 .get_highest_static_file_tx(StaticFileSegment::Transactions);

diff --git a/crates/storage/provider/src/traits/block.rs b/crates/storage/provider/src/traits/block.rs
@@ -3,33 +3,10 @@ use reth_db_api::models::StoredBlockBodyIndices;
 use reth_execution_types::{Chain, ExecutionOutcome};
 use reth_node_types::NodePrimitives;
 use reth_primitives::SealedBlockWithSenders;
-use reth_storage_api::NodePrimitivesProvider;
+use reth_storage_api::{NodePrimitivesProvider, StorageLocation};
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{updates::TrieUpdates, HashedPostStateSorted};
 
-/// An enum that represents the storage location for a piece of data.
-#[derive(Debug, Copy, Clone, PartialEq, Eq)]
-pub enum StorageLocation {
-    /// Write only to static files.
-    StaticFiles,
-    /// Write only to the database.
-    Database,
-    /// Write to both the database and static files.
-    Both,
-}
-
-impl StorageLocation {
-    /// Returns true if the storage location includes static files.
-    pub const fn static_files(&self) -> bool {
-        matches!(self, Self::StaticFiles | Self::Both)
-    }
-
-    /// Returns true if the storage location includes the database.
-    pub const fn database(&self) -> bool {
-        matches!(self, Self::Database | Self::Both)
-    }
-}
-
 /// `BlockExecution` Writer
 pub trait BlockExecutionWriter:
     NodePrimitivesProvider<Primitives: NodePrimitives<Block = Self::Block>> + BlockWriter + Send + Sync
@@ -120,7 +97,7 @@ pub trait BlockWriter: Send + Sync {
     fn append_block_bodies(
         &self,
         bodies: Vec<(BlockNumber, Option<<Self::Block as reth_primitives_traits::Block>::Body>)>,
-        write_transactions_to: StorageLocation,
+        write_to: StorageLocation,
     ) -> ProviderResult<()>;
 
     /// Removes all blocks above the given block number from the database.
@@ -129,14 +106,14 @@ pub trait BlockWriter: Send + Sync {
     fn remove_blocks_above(
         &self,
         block: BlockNumber,
-        remove_transactions_from: StorageLocation,
+        remove_from: StorageLocation,
     ) -> ProviderResult<()>;
 
     /// Removes all block bodies above the given block number from the database.
     fn remove_bodies_above(
         &self,
         block: BlockNumber,
-        remove_transactions_from: StorageLocation,
+        remove_from: StorageLocation,
     ) -> ProviderResult<()>;
 
     /// Appends a batch of sealed blocks to the blockchain, including sender information, and

diff --git a/crates/storage/storage-api/src/chain.rs b/crates/storage/storage-api/src/chain.rs
@@ -1,4 +1,4 @@
-use crate::DBProvider;
+use crate::{DBProvider, StorageLocation};
 use alloy_primitives::BlockNumber;
 use reth_chainspec::{ChainSpecProvider, EthereumHardforks};
 use reth_db::{
@@ -22,13 +22,15 @@ pub trait BlockBodyWriter<Provider, Body: BlockBody> {
         &self,
         provider: &Provider,
         bodies: Vec<(BlockNumber, Option<Body>)>,
+        write_to: StorageLocation,
     ) -> ProviderResult<()>;
 
     /// Removes all block bodies above the given block number from the database.
     fn remove_block_bodies_above(
         &self,
         provider: &Provider,
         block: BlockNumber,
+        remove_from: StorageLocation,
     ) -> ProviderResult<()>;
 }
 
@@ -87,6 +89,7 @@ where
         &self,
         provider: &Provider,
         bodies: Vec<(u64, Option<reth_primitives::BlockBody>)>,
+        _write_to: StorageLocation,
     ) -> ProviderResult<()> {
         let mut ommers_cursor = provider.tx_ref().cursor_write::<tables::BlockOmmers>()?;
         let mut withdrawals_cursor =
@@ -116,6 +119,7 @@ where
         &self,
         provider: &Provider,
         block: BlockNumber,
+        _remove_from: StorageLocation,
     ) -> ProviderResult<()> {
         provider.tx_ref().unwind_table_by_num::<tables::BlockWithdrawals>(block)?;
         provider.tx_ref().unwind_table_by_num::<tables::BlockOmmers>(block)?;

diff --git a/crates/storage/storage-api/src/storage.rs b/crates/storage/storage-api/src/storage.rs
@@ -41,3 +41,26 @@ pub trait StorageChangeSetReader: Send + Sync {
         block_number: BlockNumber,
     ) -> ProviderResult<Vec<(BlockNumberAddress, StorageEntry)>>;
 }
+
+/// An enum that represents the storage location for a piece of data.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum StorageLocation {
+    /// Write only to static files.
+    StaticFiles,
+    /// Write only to the database.
+    Database,
+    /// Write to both the database and static files.
+    Both,
+}
+
+impl StorageLocation {
+    /// Returns true if the storage location includes static files.
+    pub const fn static_files(&self) -> bool {
+        matches!(self, Self::StaticFiles | Self::Both)
+    }
+
+    /// Returns true if the storage location includes the database.
+    pub const fn database(&self) -> bool {
+        matches!(self, Self::Database | Self::Both)
+    }
+}