chore: turn static file field generic on StorageWriter (#9980)

Co-authored-by: Alexey Shekhirin <a.shekhirin@gmail.com>

crates/storage/errors/src/writer.rs

@@ -1,4 +1,5 @@
 use crate::db::DatabaseError;
+use reth_primitives::StaticFileSegment;
 
 /// `StorageWriter` related errors
 #[derive(Clone, Debug, thiserror_no_std::Error, PartialEq, Eq)]
@@ -9,6 +10,9 @@ pub enum StorageWriterError {
     /// Static file writer is missing
     #[error("Static file writer is missing")]
     MissingStaticFileWriter,
+    /// Static file writer is of wrong segment
+    #[error("Static file writer is of wrong segment: got {0}, expected {1}")]
+    IncorrectStaticFileWriter(StaticFileSegment, StaticFileSegment),
     /// Database-related errors.
     #[error(transparent)]
     Database(#[from] DatabaseError),

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

@@ -1,16 +1,16 @@
 use crate::{
-    providers::StaticFileProviderRWRefMut, DatabaseProvider, DatabaseProviderRO,
-    DatabaseProviderRW, StateChangeWriter, StateWriter,
+    providers::StaticFileProviderRWRefMut, DatabaseProvider, StateChangeWriter, StateWriter,
 };
 use reth_db::{
     cursor::DbCursorRO,
     tables,
     transaction::{DbTx, DbTxMut},
-    Database,
 };
 use reth_errors::{ProviderError, ProviderResult};
 use reth_execution_types::ExecutionOutcome;
-use reth_primitives::{BlockNumber, Header, TransactionSignedNoHash, B256, U256};
+use reth_primitives::{
+    BlockNumber, Header, StaticFileSegment, TransactionSignedNoHash, B256, U256,
+};
 use reth_storage_api::{HeaderProvider, ReceiptWriter};
 use reth_storage_errors::writer::StorageWriterError;
 use revm::db::OriginalValuesKnown;
@@ -29,96 +29,101 @@ enum StorageType<C = (), S = ()> {
 /// [`StorageWriter`] is responsible for managing the writing to either database, static file or
 /// both.
 #[derive(Debug)]
-pub struct StorageWriter<'a, 'b, TX> {
-    database_writer: Option<&'a DatabaseProvider<TX>>,
-    static_file_writer: Option<StaticFileProviderRWRefMut<'b>>,
+pub struct StorageWriter<'a, TX, SF> {
+    database: Option<&'a DatabaseProvider<TX>>,
+    static_file: Option<SF>,
 }
 
-impl<'a, 'b, TX> StorageWriter<'a, 'b, TX> {
+impl<'a, TX, SF> StorageWriter<'a, TX, SF> {
     /// Creates a new instance of [`StorageWriter`].
     ///
     /// # Parameters
-    /// - `database_writer`: An optional reference to a database writer.
-    /// - `static_file_writer`: An optional mutable reference to a static file writer.
-    pub const fn new(
-        database_writer: Option<&'a DatabaseProvider<TX>>,
-        static_file_writer: Option<StaticFileProviderRWRefMut<'b>>,
-    ) -> Self {
-        Self { database_writer, static_file_writer }
+    /// - `database`: An optional reference to a database provider.
+    /// - `static_file`: An optional mutable reference to a static file instance.
+    pub const fn new(database: Option<&'a DatabaseProvider<TX>>, static_file: Option<SF>) -> Self {
+        Self { database, static_file }
     }
 
-    /// Creates a new instance of [`StorageWriter`] from a static file writer.
-    pub const fn from_static_file_writer(
-        static_file_writer: StaticFileProviderRWRefMut<'b>,
-    ) -> Self {
-        Self::new(None, Some(static_file_writer))
+    /// Creates a new instance of [`StorageWriter`] from a static file instance.
+    pub const fn from_static_file(static_file: SF) -> Self {
+        Self::new(None, Some(static_file))
     }
 
-    /// Creates a new instance of [`StorageWriter`] from a read-only database provider.
-    pub const fn from_database_provider_ro<DB>(
-        database: &'a DatabaseProviderRO<DB>,
-    ) -> StorageWriter<'_, '_, <DB as Database>::TX>
-    where
-        DB: Database,
-    {
-        StorageWriter::new(Some(database), None)
-    }
-
-    /// Creates a new instance of [`StorageWriter`] from a read-write database provider.
-    pub fn from_database_provider_rw<DB>(
-        database: &'a DatabaseProviderRW<DB>,
-    ) -> StorageWriter<'_, '_, <DB as Database>::TXMut>
-    where
-        DB: Database,
-    {
+    /// Creates a new instance of [`StorageWriter`] from a database provider.
+    pub const fn from_database(database: &'a DatabaseProvider<TX>) -> Self {
         StorageWriter::new(Some(database), None)
     }
 
     /// Returns a reference to the database writer.
     ///
     /// # Panics
-    /// If the database writer is not set.
-    fn database_writer(&self) -> &DatabaseProvider<TX> {
-        self.database_writer.as_ref().expect("should exist")
+    /// If the database provider is not set.
+    fn database(&self) -> &DatabaseProvider<TX> {
+        self.database.as_ref().expect("should exist")
     }
 
-    /// Returns a mutable reference to the static file writer.
+    /// Returns a mutable reference to the static file instance.
     ///
     /// # Panics
-    /// If the static file writer is not set.
-    fn static_file_writer(&mut self) -> &mut StaticFileProviderRWRefMut<'b> {
-        self.static_file_writer.as_mut().expect("should exist")
+    /// If the static file instance is not set.
+    fn static_file(&mut self) -> &mut SF {
+        self.static_file.as_mut().expect("should exist")
     }
 
-    /// Ensures that the database writer is set.
+    /// Ensures that the database provider is set.
     ///
     /// # Returns
-    /// - `Ok(())` if the database writer is set.
-    /// - `Err(StorageWriterError::MissingDatabaseWriter)` if the database writer is not set.
-    const fn ensure_database_writer(&self) -> Result<(), StorageWriterError> {
-        if self.database_writer.is_none() {
+    /// - `Ok(())` if the database provider is set.
+    /// - `Err(StorageWriterError::MissingDatabaseWriter)` if the database provider is not set.
+    const fn ensure_database(&self) -> Result<(), StorageWriterError> {
+        if self.database.is_none() {
             return Err(StorageWriterError::MissingDatabaseWriter)
         }
         Ok(())
     }
 
-    /// Ensures that the static file writer is set.
+    /// Ensures that the static file instance is set.
     ///
     /// # Returns
-    /// - `Ok(())` if the static file writer is set.
-    /// - `Err(StorageWriterError::MissingStaticFileWriter)` if the static file writer is not set.
-    const fn ensure_static_file_writer(&self) -> Result<(), StorageWriterError> {
-        if self.static_file_writer.is_none() {
+    /// - `Ok(())` if the static file instance is set.
+    /// - `Err(StorageWriterError::MissingStaticFileWriter)` if the static file instance is not set.
+    #[allow(unused)]
+    const fn ensure_static_file(&self) -> Result<(), StorageWriterError> {
+        if self.static_file.is_none() {
             return Err(StorageWriterError::MissingStaticFileWriter)
         }
         Ok(())
     }
 }
 
-impl<'a, 'b, TX> StorageWriter<'a, 'b, TX>
+impl<'a, 'b, TX> StorageWriter<'a, TX, StaticFileProviderRWRefMut<'b>>
 where
     TX: DbTx,
 {
+    /// Ensures that the static file writer is set and of the right [`StaticFileSegment`] variant.
+    ///
+    /// # Returns
+    /// - `Ok(())` if the static file writer is set.
+    /// - `Err(StorageWriterError::MissingStaticFileWriter)` if the static file instance is not set.
+    fn ensure_static_file_segment(
+        &self,
+        segment: StaticFileSegment,
+    ) -> Result<(), StorageWriterError> {
+        match &self.static_file {
+            Some(writer) => {
+                if writer.user_header().segment() != segment {
+                    Err(StorageWriterError::IncorrectStaticFileWriter(
+                        writer.user_header().segment(),
+                        segment,
+                    ))
+                } else {
+                    Ok(())
+                }
+            }
+            None => Err(StorageWriterError::MissingStaticFileWriter),
+        }
+    }
+
     /// Appends headers to static files, using the
     /// [`HeaderTerminalDifficulties`](tables::HeaderTerminalDifficulties) table to determine the
     /// total difficulty of the parent block during header insertion.
@@ -134,19 +139,19 @@ where
         I: Borrow<(H, B256)>,
         H: Borrow<Header>,
     {
-        self.ensure_database_writer()?;
-        self.ensure_static_file_writer()?;
+        self.ensure_database()?;
+        self.ensure_static_file_segment(StaticFileSegment::Headers)?;
 
         let mut td = self
-            .database_writer()
+            .database()
             .header_td_by_number(initial_block_number)?
             .ok_or(ProviderError::TotalDifficultyNotFound(initial_block_number))?;
 
         for pair in headers {
             let (header, hash) = pair.borrow();
             let header = header.borrow();
             td += header.difficulty;
-            self.static_file_writer().append_header(header, td, hash)?;
+            self.static_file().append_header(header, td, hash)?;
         }
 
         Ok(td)
@@ -166,11 +171,11 @@ where
     where
         T: Borrow<Vec<TransactionSignedNoHash>>,
     {
-        self.ensure_database_writer()?;
-        self.ensure_static_file_writer()?;
+        self.ensure_database()?;
+        self.ensure_static_file_segment(StaticFileSegment::Transactions)?;
 
         let mut bodies_cursor =
-            self.database_writer().tx_ref().cursor_read::<tables::BlockBodyIndices>()?;
+            self.database().tx_ref().cursor_read::<tables::BlockBodyIndices>()?;
 
         let mut last_tx_idx = None;
         for (idx, transactions) in transactions.enumerate() {
@@ -191,11 +196,11 @@ where
                 .ok_or_else(|| ProviderError::BlockBodyIndicesNotFound(block_number))?;
 
             for tx in transactions.borrow() {
-                self.static_file_writer().append_transaction(tx_index, tx)?;
+                self.static_file().append_transaction(tx_index, tx)?;
                 tx_index += 1;
             }
 
-            self.static_file_writer().increment_block(block_number)?;
+            self.static_file().increment_block(block_number)?;
 
             // update index
             last_tx_idx = Some(tx_index);
@@ -204,7 +209,7 @@ where
     }
 }
 
-impl<'a, 'b, TX> StorageWriter<'a, 'b, TX>
+impl<'a, 'b, TX> StorageWriter<'a, TX, StaticFileProviderRWRefMut<'b>>
 where
     TX: DbTxMut + DbTx,
 {
@@ -225,23 +230,21 @@ where
         initial_block_number: BlockNumber,
         blocks: impl Iterator<Item = Vec<Option<reth_primitives::Receipt>>>,
     ) -> ProviderResult<()> {
-        self.ensure_database_writer()?;
+        self.ensure_database()?;
         let mut bodies_cursor =
-            self.database_writer().tx_ref().cursor_read::<tables::BlockBodyIndices>()?;
+            self.database().tx_ref().cursor_read::<tables::BlockBodyIndices>()?;
 
         // We write receipts to database in two situations:
         // * If we are in live sync. In this case, `StorageWriter` is built without a static file
         //   writer.
         // * If there is any kind of receipt pruning
-        let mut storage_type = if self.static_file_writer.is_none() ||
-            self.database_writer().prune_modes_ref().has_receipts_pruning()
+        let mut storage_type = if self.static_file.is_none() ||
+            self.database().prune_modes_ref().has_receipts_pruning()
         {
-            StorageType::Database(
-                self.database_writer().tx_ref().cursor_write::<tables::Receipts>()?,
-            )
+            StorageType::Database(self.database().tx_ref().cursor_write::<tables::Receipts>()?)
         } else {
-            self.ensure_static_file_writer()?;
-            StorageType::StaticFile(self.static_file_writer())
+            self.ensure_static_file_segment(StaticFileSegment::Receipts)?;
+            StorageType::StaticFile(self.static_file())
         };
 
         let mut last_tx_idx = None;
@@ -287,7 +290,7 @@ where
     }
 }
 
-impl<'a, 'b, TX> StateWriter for StorageWriter<'a, 'b, TX>
+impl<'a, 'b, TX> StateWriter for StorageWriter<'a, TX, StaticFileProviderRWRefMut<'b>>
 where
     TX: DbTxMut + DbTx,
 {
@@ -298,18 +301,18 @@ where
         execution_outcome: ExecutionOutcome,
         is_value_known: OriginalValuesKnown,
     ) -> ProviderResult<()> {
-        self.ensure_database_writer()?;
+        self.ensure_database()?;
         let (plain_state, reverts) =
             execution_outcome.bundle.into_plain_state_and_reverts(is_value_known);
 
-        self.database_writer().write_state_reverts(reverts, execution_outcome.first_block)?;
+        self.database().write_state_reverts(reverts, execution_outcome.first_block)?;
 
         self.append_receipts_from_blocks(
             execution_outcome.first_block,
             execution_outcome.receipts.into_iter(),
         )?;
 
-        self.database_writer().write_state_changes(plain_state)?;
+        self.database().write_state_changes(plain_state)?;
 
         Ok(())
     }