From a129f62aaa4ccfa9926498804fa5b70bf6f72784 Mon Sep 17 00:00:00 2001
From: Varun Doshi <61531351+varun-doshi@users.noreply.github.com>
Date: Mon, 14 Oct 2024 14:49:04 +0530
Subject: [PATCH] feat: reset pruned numbers on stage drop (#11491)

Co-authored-by: Alexey Shekhirin <a.shekhirin@gmail.com>
---
 crates/cli/commands/src/stage/drop.rs | 114 ++++++++++++--------------
 1 file changed, 51 insertions(+), 63 deletions(-)

diff --git a/crates/cli/commands/src/stage/drop.rs b/crates/cli/commands/src/stage/drop.rs
index e324c98515..9e0396404b 100644
--- a/crates/cli/commands/src/stage/drop.rs
+++ b/crates/cli/commands/src/stage/drop.rs
@@ -4,7 +4,7 @@ use clap::Parser;
 use itertools::Itertools;
 use reth_chainspec::{EthChainSpec, EthereumHardforks};
 use reth_cli::chainspec::ChainSpecParser;
-use reth_db::{static_file::iter_static_files, tables};
+use reth_db::{mdbx::tx::Tx, static_file::iter_static_files, tables, DatabaseError};
 use reth_db_api::transaction::{DbTx, DbTxMut};
 use reth_db_common::{
     init::{insert_genesis_header, insert_genesis_history, insert_genesis_state},
@@ -69,42 +69,28 @@ impl<C: ChainSpecParser<ChainSpec: EthChainSpec + EthereumHardforks>> Command<C>
                 tx.clear::<tables::Headers>()?;
                 tx.clear::<tables::HeaderTerminalDifficulties>()?;
                 tx.clear::<tables::HeaderNumbers>()?;
-                tx.put::<tables::StageCheckpoints>(
-                    StageId::Headers.to_string(),
-                    Default::default(),
-                )?;
+                reset_stage_checkpoint(tx, StageId::Headers)?;
+
                 insert_genesis_header(&provider_rw.0, &static_file_provider, &self.env.chain)?;
             }
             StageEnum::Bodies => {
                 tx.clear::<tables::BlockBodyIndices>()?;
                 tx.clear::<tables::Transactions>()?;
+                reset_prune_checkpoint(tx, PruneSegment::Transactions)?;
+
                 tx.clear::<tables::TransactionBlocks>()?;
                 tx.clear::<tables::BlockOmmers>()?;
                 tx.clear::<tables::BlockWithdrawals>()?;
                 tx.clear::<tables::BlockRequests>()?;
-                tx.put::<tables::StageCheckpoints>(
-                    StageId::Bodies.to_string(),
-                    Default::default(),
-                )?;
+                reset_stage_checkpoint(tx, StageId::Bodies)?;
+
                 insert_genesis_header(&provider_rw.0, &static_file_provider, &self.env.chain)?;
             }
             StageEnum::Senders => {
                 tx.clear::<tables::TransactionSenders>()?;
                 // Reset pruned numbers to not count them in the next rerun's stage progress
-                if let Some(mut prune_checkpoint) =
-                    tx.get::<tables::PruneCheckpoints>(PruneSegment::SenderRecovery)?
-                {
-                    prune_checkpoint.block_number = None;
-                    prune_checkpoint.tx_number = None;
-                    tx.put::<tables::PruneCheckpoints>(
-                        PruneSegment::SenderRecovery,
-                        prune_checkpoint,
-                    )?;
-                }
-                tx.put::<tables::StageCheckpoints>(
-                    StageId::SenderRecovery.to_string(),
-                    Default::default(),
-                )?;
+                reset_prune_checkpoint(tx, PruneSegment::SenderRecovery)?;
+                reset_stage_checkpoint(tx, StageId::SenderRecovery)?;
             }
             StageEnum::Execution => {
                 tx.clear::<tables::PlainAccountState>()?;
@@ -113,53 +99,38 @@ impl<C: ChainSpecParser<ChainSpec: EthChainSpec + EthereumHardforks>> Command<C>
                 tx.clear::<tables::StorageChangeSets>()?;
                 tx.clear::<tables::Bytecodes>()?;
                 tx.clear::<tables::Receipts>()?;
-                tx.put::<tables::StageCheckpoints>(
-                    StageId::Execution.to_string(),
-                    Default::default(),
-                )?;
+
+                reset_prune_checkpoint(tx, PruneSegment::Receipts)?;
+                reset_prune_checkpoint(tx, PruneSegment::ContractLogs)?;
+                reset_stage_checkpoint(tx, StageId::Execution)?;
+
                 let alloc = &self.env.chain.genesis().alloc;
                 insert_genesis_state(&provider_rw.0, alloc.iter())?;
             }
             StageEnum::AccountHashing => {
                 tx.clear::<tables::HashedAccounts>()?;
-                tx.put::<tables::StageCheckpoints>(
-                    StageId::AccountHashing.to_string(),
-                    Default::default(),
-                )?;
+                reset_stage_checkpoint(tx, StageId::AccountHashing)?;
             }
             StageEnum::StorageHashing => {
                 tx.clear::<tables::HashedStorages>()?;
-                tx.put::<tables::StageCheckpoints>(
-                    StageId::StorageHashing.to_string(),
-                    Default::default(),
-                )?;
+                reset_stage_checkpoint(tx, StageId::StorageHashing)?;
             }
             StageEnum::Hashing => {
                 // Clear hashed accounts
                 tx.clear::<tables::HashedAccounts>()?;
-                tx.put::<tables::StageCheckpoints>(
-                    StageId::AccountHashing.to_string(),
-                    Default::default(),
-                )?;
+                reset_stage_checkpoint(tx, StageId::AccountHashing)?;
 
                 // Clear hashed storages
                 tx.clear::<tables::HashedStorages>()?;
-                tx.put::<tables::StageCheckpoints>(
-                    StageId::StorageHashing.to_string(),
-                    Default::default(),
-                )?;
+                reset_stage_checkpoint(tx, StageId::StorageHashing)?;
             }
             StageEnum::Merkle => {
                 tx.clear::<tables::AccountsTrie>()?;
                 tx.clear::<tables::StoragesTrie>()?;
-                tx.put::<tables::StageCheckpoints>(
-                    StageId::MerkleExecute.to_string(),
-                    Default::default(),
-                )?;
-                tx.put::<tables::StageCheckpoints>(
-                    StageId::MerkleUnwind.to_string(),
-                    Default::default(),
-                )?;
+
+                reset_stage_checkpoint(tx, StageId::MerkleExecute)?;
+                reset_stage_checkpoint(tx, StageId::MerkleUnwind)?;
+
                 tx.delete::<tables::StageCheckpointProgresses>(
                     StageId::MerkleExecute.to_string(),
                     None,
@@ -168,22 +139,17 @@ impl<C: ChainSpecParser<ChainSpec: EthChainSpec + EthereumHardforks>> Command<C>
             StageEnum::AccountHistory | StageEnum::StorageHistory => {
                 tx.clear::<tables::AccountsHistory>()?;
                 tx.clear::<tables::StoragesHistory>()?;
-                tx.put::<tables::StageCheckpoints>(
-                    StageId::IndexAccountHistory.to_string(),
-                    Default::default(),
-                )?;
-                tx.put::<tables::StageCheckpoints>(
-                    StageId::IndexStorageHistory.to_string(),
-                    Default::default(),
-                )?;
+
+                reset_stage_checkpoint(tx, StageId::IndexAccountHistory)?;
+                reset_stage_checkpoint(tx, StageId::IndexStorageHistory)?;
+
                 insert_genesis_history(&provider_rw.0, self.env.chain.genesis().alloc.iter())?;
             }
             StageEnum::TxLookup => {
                 tx.clear::<tables::TransactionHashNumbers>()?;
-                tx.put::<tables::StageCheckpoints>(
-                    StageId::TransactionLookup.to_string(),
-                    Default::default(),
-                )?;
+                reset_prune_checkpoint(tx, PruneSegment::TransactionLookup)?;
+
+                reset_stage_checkpoint(tx, StageId::TransactionLookup)?;
                 insert_genesis_header(&provider_rw.0, &static_file_provider, &self.env.chain)?;
             }
         }
@@ -195,3 +161,25 @@ impl<C: ChainSpecParser<ChainSpec: EthChainSpec + EthereumHardforks>> Command<C>
         Ok(())
     }
 }
+
+fn reset_prune_checkpoint(
+    tx: &Tx<reth_db::mdbx::RW>,
+    prune_segment: PruneSegment,
+) -> Result<(), DatabaseError> {
+    if let Some(mut prune_checkpoint) = tx.get::<tables::PruneCheckpoints>(prune_segment)? {
+        prune_checkpoint.block_number = None;
+        prune_checkpoint.tx_number = None;
+        tx.put::<tables::PruneCheckpoints>(prune_segment, prune_checkpoint)?;
+    }
+
+    Ok(())
+}
+
+fn reset_stage_checkpoint(
+    tx: &Tx<reth_db::mdbx::RW>,
+    stage_id: StageId,
+) -> Result<(), DatabaseError> {
+    tx.put::<tables::StageCheckpoints>(stage_id.to_string(), Default::default())?;
+
+    Ok(())
+}