fix(stages): include block range in headers checkpoint (#2950)

crates/primitives/src/stage/checkpoints.rs

@@ -1,6 +1,6 @@
 use crate::{
     trie::{hash_builder::HashBuilderState, StoredSubNode},
-    Address, BlockNumber, TxNumber, H256,
+    Address, BlockNumber, H256,
 };
 use bytes::{Buf, BufMut};
 use reth_codecs::{derive_arbitrary, main_codec, Compact};
@@ -136,6 +136,16 @@ pub struct ExecutionCheckpoint {
     pub progress: EntitiesCheckpoint,
 }
 
+/// Saves the progress of Headers stage.
+#[main_codec]
+#[derive(Default, Debug, Copy, Clone, PartialEq, Eq)]
+pub struct HeadersCheckpoint {
+    /// Block range which this checkpoint is valid for.
+    pub block_range: CheckpointBlockRange,
+    /// Progress measured in gas.
+    pub progress: EntitiesCheckpoint,
+}
+
 /// Saves the progress of abstract stage iterating over or downloading entities.
 #[main_codec]
 #[derive(Debug, Default, PartialEq, Eq, Clone, Copy)]
@@ -221,6 +231,14 @@ impl StageCheckpoint {
         }
     }
 
+    /// Returns the headers stage checkpoint, if any.
+    pub fn headers_stage_checkpoint(&self) -> Option<HeadersCheckpoint> {
+        match self.stage_checkpoint {
+            Some(StageUnitCheckpoint::Headers(checkpoint)) => Some(checkpoint),
+            _ => None,
+        }
+    }
+
     /// Sets the block number.
     pub fn with_block_number(mut self, block_number: BlockNumber) -> Self {
         self.block_number = block_number;
@@ -256,6 +274,12 @@ impl StageCheckpoint {
         self.stage_checkpoint = Some(StageUnitCheckpoint::Execution(checkpoint));
         self
     }
+
+    /// Sets the stage checkpoint to headers.
+    pub fn with_headers_stage_checkpoint(mut self, checkpoint: HeadersCheckpoint) -> Self {
+        self.stage_checkpoint = Some(StageUnitCheckpoint::Headers(checkpoint));
+        self
+    }
 }
 
 impl Display for StageCheckpoint {
@@ -269,9 +293,10 @@ impl Display for StageCheckpoint {
                     progress: entities, ..
                 }) |
                 StageUnitCheckpoint::Entities(entities) |
-                StageUnitCheckpoint::Execution(ExecutionCheckpoint { progress: entities, .. }),
+                StageUnitCheckpoint::Execution(ExecutionCheckpoint { progress: entities, .. }) |
+                StageUnitCheckpoint::Headers(HeadersCheckpoint { progress: entities, .. }),
             ) => entities.fmt(f),
-            Some(StageUnitCheckpoint::Transaction(_)) | None => write!(f, "{}", self.block_number),
+            None => write!(f, "{}", self.block_number),
         }
     }
 }
@@ -282,8 +307,6 @@ impl Display for StageCheckpoint {
 #[derive_arbitrary(compact)]
 #[derive(Debug, PartialEq, Eq, Clone, Copy, Serialize, Deserialize)]
 pub enum StageUnitCheckpoint {
-    /// Saves the progress of transaction-indexed stages.
-    Transaction(TxNumber),
     /// Saves the progress of AccountHashing stage.
     Account(AccountHashingCheckpoint),
     /// Saves the progress of StorageHashing stage.
@@ -292,6 +315,8 @@ pub enum StageUnitCheckpoint {
     Entities(EntitiesCheckpoint),
     /// Saves the progress of Execution stage.
     Execution(ExecutionCheckpoint),
+    /// Saves the progress of Headers stage.
+    Headers(HeadersCheckpoint),
 }
 
 impl Compact for StageUnitCheckpoint {
@@ -300,23 +325,23 @@ impl Compact for StageUnitCheckpoint {
         B: BufMut + AsMut<[u8]>,
     {
         match self {
-            StageUnitCheckpoint::Transaction(data) => {
+            StageUnitCheckpoint::Account(data) => {
                 buf.put_u8(0);
                 1 + data.to_compact(buf)
             }
-            StageUnitCheckpoint::Account(data) => {
+            StageUnitCheckpoint::Storage(data) => {
                 buf.put_u8(1);
                 1 + data.to_compact(buf)
             }
-            StageUnitCheckpoint::Storage(data) => {
+            StageUnitCheckpoint::Entities(data) => {
                 buf.put_u8(2);
                 1 + data.to_compact(buf)
             }
-            StageUnitCheckpoint::Entities(data) => {
+            StageUnitCheckpoint::Execution(data) => {
                 buf.put_u8(3);
                 1 + data.to_compact(buf)
             }
-            StageUnitCheckpoint::Execution(data) => {
+            StageUnitCheckpoint::Headers(data) => {
                 buf.put_u8(4);
                 1 + data.to_compact(buf)
             }
@@ -329,25 +354,25 @@ impl Compact for StageUnitCheckpoint {
     {
         match buf[0] {
             0 => {
-                let (data, buf) = TxNumber::from_compact(&buf[1..], buf.len() - 1);
-                (Self::Transaction(data), buf)
-            }
-            1 => {
                 let (data, buf) = AccountHashingCheckpoint::from_compact(&buf[1..], buf.len() - 1);
                 (Self::Account(data), buf)
             }
-            2 => {
+            1 => {
                 let (data, buf) = StorageHashingCheckpoint::from_compact(&buf[1..], buf.len() - 1);
                 (Self::Storage(data), buf)
             }
-            3 => {
+            2 => {
                 let (data, buf) = EntitiesCheckpoint::from_compact(&buf[1..], buf.len() - 1);
                 (Self::Entities(data), buf)
             }
-            4 => {
+            3 => {
                 let (data, buf) = ExecutionCheckpoint::from_compact(&buf[1..], buf.len() - 1);
                 (Self::Execution(data), buf)
             }
+            4 => {
+                let (data, buf) = HeadersCheckpoint::from_compact(&buf[1..], buf.len() - 1);
+                (Self::Headers(data), buf)
+            }
             _ => unreachable!("Junk data in database: unknown StageUnitCheckpoint variant"),
         }
     }
@@ -383,7 +408,6 @@ mod tests {
     fn stage_unit_checkpoint_roundtrip() {
         let mut rng = rand::thread_rng();
         let checkpoints = vec![
-            StageUnitCheckpoint::Transaction(rng.gen()),
             StageUnitCheckpoint::Account(AccountHashingCheckpoint {
                 address: Some(Address::from_low_u64_be(rng.gen())),
                 block_range: CheckpointBlockRange { from: rng.gen(), to: rng.gen() },
@@ -412,6 +436,13 @@ mod tests {
                     total: Some(u32::MAX as u64 + rng.gen::<u64>()),
                 },
             }),
+            StageUnitCheckpoint::Headers(HeadersCheckpoint {
+                block_range: CheckpointBlockRange { from: rng.gen(), to: rng.gen() },
+                progress: EntitiesCheckpoint {
+                    processed: rng.gen::<u32>() as u64,
+                    total: Some(u32::MAX as u64 + rng.gen::<u64>()),
+                },
+            }),
         ];
 
         for checkpoint in checkpoints {

crates/primitives/src/stage/mod.rs

@@ -6,5 +6,6 @@ pub use id::StageId;
 mod checkpoints;
 pub use checkpoints::{
     AccountHashingCheckpoint, CheckpointBlockRange, EntitiesCheckpoint, ExecutionCheckpoint,
-    MerkleCheckpoint, StageCheckpoint, StageUnitCheckpoint, StorageHashingCheckpoint,
+    HeadersCheckpoint, MerkleCheckpoint, StageCheckpoint, StageUnitCheckpoint,
+    StorageHashingCheckpoint,
 };

crates/stages/src/pipeline/sync_metrics.rs

@@ -4,8 +4,8 @@ use reth_metrics::{
 };
 use reth_primitives::{
     stage::{
-        AccountHashingCheckpoint, EntitiesCheckpoint, ExecutionCheckpoint, StageCheckpoint,
-        StageId, StageUnitCheckpoint, StorageHashingCheckpoint,
+        AccountHashingCheckpoint, EntitiesCheckpoint, ExecutionCheckpoint, HeadersCheckpoint,
+        StageCheckpoint, StageId, StageUnitCheckpoint, StorageHashingCheckpoint,
     },
     BlockNumber,
 };
@@ -43,25 +43,24 @@ impl Metrics {
 
         let (processed, total) = match checkpoint.stage_checkpoint {
             Some(
-                StageUnitCheckpoint::Entities(progress @ EntitiesCheckpoint { .. }) |
-                StageUnitCheckpoint::Execution(ExecutionCheckpoint { progress, .. }) |
                 // Only report metrics for hashing stages if `total` is known, otherwise it means
-                // we're unwinding and operating on changesets, rather than accounts or storage slots.
+                // we're unwinding and operating on changesets, rather than accounts or storage
+                // slots.
                 StageUnitCheckpoint::Account(AccountHashingCheckpoint {
                     progress: progress @ EntitiesCheckpoint { total: Some(_), .. },
                     ..
                 }) |
                 StageUnitCheckpoint::Storage(StorageHashingCheckpoint {
                     progress: progress @ EntitiesCheckpoint { total: Some(_), .. },
                     ..
-                }),
+                }) |
+                StageUnitCheckpoint::Entities(progress @ EntitiesCheckpoint { .. }) |
+                StageUnitCheckpoint::Execution(ExecutionCheckpoint { progress, .. }) |
+                StageUnitCheckpoint::Headers(HeadersCheckpoint { progress, .. }),
             ) => (progress.processed, progress.total),
-            Some(
-                StageUnitCheckpoint::Transaction(_) |
-                StageUnitCheckpoint::Account(_) |
-                StageUnitCheckpoint::Storage(_),
-            ) |
-            None => (checkpoint.block_number, max_block_number),
+            Some(StageUnitCheckpoint::Account(_) | StageUnitCheckpoint::Storage(_)) | None => {
+                (checkpoint.block_number, max_block_number)
+            }
         };
 
         stage_metrics.entities_processed.set(processed as f64);