feat: use pipeline for reth stage unwind (#7085)

Co-authored-by: joshieDo <ranriver@protonmail.com>
Co-authored-by: joshieDo <93316087+joshieDo@users.noreply.github.com>
Co-authored-by: Alexey Shekhirin <a.shekhirin@gmail.com>

crates/primitives/src/static_file/mod.rs

@@ -44,6 +44,11 @@ impl HighestStaticFiles {
             StaticFileSegment::Receipts => &mut self.receipts,
         }
     }
+
+    /// Returns the maximum block of all segments.
+    pub fn max(&self) -> Option<u64> {
+        [self.headers, self.transactions, self.receipts].iter().filter_map(|&option| option).max()
+    }
 }
 
 /// Each static file has a fixed number of blocks. This gives out the range where the requested

crates/stages/src/pipeline/mod.rs

@@ -232,7 +232,7 @@ where
     ///
     /// CAUTION: This method locks the static file producer Mutex, hence can block the thread if the
     /// lock is occupied.
-    fn produce_static_files(&mut self) -> RethResult<()> {
+    pub fn produce_static_files(&mut self) -> RethResult<()> {
         let mut static_file_producer = self.static_file_producer.lock();
 
         let provider = self.provider_factory.provider()?;

crates/stages/src/stages/bodies.rs

@@ -13,9 +13,12 @@ use reth_interfaces::{
 };
 use reth_primitives::{
     stage::{EntitiesCheckpoint, StageCheckpoint, StageId},
-    StaticFileSegment,
+    StaticFileSegment, TxNumber,
+};
+use reth_provider::{
+    providers::{StaticFileProvider, StaticFileWriter},
+    BlockReader, DatabaseProviderRW, HeaderProvider, ProviderError, StatsReader,
 };
-use reth_provider::{providers::StaticFileWriter, DatabaseProviderRW, HeaderProvider, StatsReader};
 use std::{
     cmp::Ordering,
     task::{ready, Context, Poll},
@@ -145,17 +148,11 @@ impl<DB: Database, D: BodyDownloader> Stage<DB> for BodyStage<D> {
             // error will trigger an unwind, that will bring the database to the same height as the
             // static files.
             Ordering::Less => {
-                let last_block = static_file_provider
-                    .get_highest_static_file_block(StaticFileSegment::Transactions)
-                    .unwrap_or_default();
-
-                let missing_block =
-                    Box::new(provider.sealed_header(last_block + 1)?.unwrap_or_default());
-
-                return Err(StageError::MissingStaticFileData {
-                    block: missing_block,
-                    segment: StaticFileSegment::Transactions,
-                })
+                return Err(missing_static_data_error(
+                    next_static_file_tx_num.saturating_sub(1),
+                    static_file_provider,
+                    provider,
+                )?)
             }
             Ordering::Equal => {}
         }
@@ -311,17 +308,11 @@ impl<DB: Database, D: BodyDownloader> Stage<DB> for BodyStage<D> {
         // If there are more transactions on database, then we are missing static file data and we
         // need to unwind further.
         if db_tx_num > static_file_tx_num {
-            let last_block = static_file_provider
-                .get_highest_static_file_block(StaticFileSegment::Transactions)
-                .unwrap_or_default();
-
-            let missing_block =
-                Box::new(provider.sealed_header(last_block + 1)?.unwrap_or_default());
-
-            return Err(StageError::MissingStaticFileData {
-                block: missing_block,
-                segment: StaticFileSegment::Transactions,
-            })
+            return Err(missing_static_data_error(
+                static_file_tx_num,
+                static_file_provider,
+                provider,
+            )?)
         }
 
         // Unwinds static file
@@ -335,6 +326,37 @@ impl<DB: Database, D: BodyDownloader> Stage<DB> for BodyStage<D> {
     }
 }
 
+fn missing_static_data_error<DB: Database>(
+    last_tx_num: TxNumber,
+    static_file_provider: &StaticFileProvider,
+    provider: &DatabaseProviderRW<DB>,
+) -> Result<StageError, ProviderError> {
+    let mut last_block = static_file_provider
+        .get_highest_static_file_block(StaticFileSegment::Transactions)
+        .unwrap_or_default();
+
+    // To be extra safe, we make sure that the last tx num matches the last block from its indices.
+    // If not, get it.
+    loop {
+        if let Some(indices) = provider.block_body_indices(last_block)? {
+            if indices.last_tx_num() <= last_tx_num {
+                break
+            }
+        }
+        if last_block == 0 {
+            break
+        }
+        last_block -= 1;
+    }
+
+    let missing_block = Box::new(provider.sealed_header(last_block + 1)?.unwrap_or_default());
+
+    Ok(StageError::MissingStaticFileData {
+        block: missing_block,
+        segment: StaticFileSegment::Transactions,
+    })
+}
+
 // TODO(alexey): ideally, we want to measure Bodies stage progress in bytes, but it's hard to know
 //  beforehand how many bytes we need to download. So the good solution would be to measure the
 //  progress in gas as a proxy to size. Execution stage uses a similar approach.

crates/stages/src/stages/execution.rs

@@ -578,13 +578,30 @@ where
             start_block.saturating_sub(1),
         )?,
         Ordering::Less => {
-            let last_block = static_file_provider
+            let mut last_block = static_file_provider
                 .get_highest_static_file_block(StaticFileSegment::Receipts)
                 .unwrap_or(0);
 
-            let missing_block = Box::new(
-                tx.get::<tables::Headers>(last_block + 1)?.unwrap_or_default().seal_slow(),
-            );
+            let last_receipt_num = static_file_provider
+                .get_highest_static_file_tx(StaticFileSegment::Receipts)
+                .unwrap_or(0);
+
+            // To be extra safe, we make sure that the last receipt num matches the last block from
+            // its indices. If not, get it.
+            loop {
+                if let Some(indices) = provider.block_body_indices(last_block)? {
+                    if indices.last_tx_num() <= last_receipt_num {
+                        break
+                    }
+                }
+                if last_block == 0 {
+                    break
+                }
+                last_block -= 1;
+            }
+
+            let missing_block =
+                Box::new(provider.sealed_header(last_block + 1)?.unwrap_or_default());
 
             return Err(StageError::MissingStaticFileData {
                 block: missing_block,