chore(perf): backport #26766

crates/stages/stages/src/stages/index_account_history.rs

@@ -1,5 +1,8 @@
 use super::collect_account_history_indices;
-use crate::stages::utils::{collect_history_indices, load_account_history};
+use crate::stages::utils::{
+    collect_history_indices, load_account_history_append, prepare_account_history_writes,
+    write_prepared_history_shards,
+};
 use reth_config::config::{EtlConfig, IndexHistoryConfig};
 use reth_db_api::{models::ShardedKey, tables, transaction::DbTxMut, Tables};
 use reth_provider::{
@@ -8,7 +11,8 @@ use reth_provider::{
 };
 use reth_prune_types::{PruneCheckpoint, PruneMode, PrunePurpose, PruneSegment};
 use reth_stages_api::{
-    ExecInput, ExecOutput, Stage, StageCheckpoint, StageError, StageId, UnwindInput, UnwindOutput,
+    BlockRangeOutput, ExecInput, ExecOutput, Stage, StageCheckpoint, StageError, StageId,
+    UnwindInput, UnwindOutput,
 };
 use std::fmt::Debug;
 use tracing::info;
@@ -66,6 +70,10 @@ where
         provider: &Provider,
         mut input: ExecInput,
     ) -> Result<ExecOutput, StageError> {
+        let initial_sync = input.checkpoint.is_none();
+        let rebuild_from_empty = input.checkpoint().block_number == 0;
+        let mut prune_target_applied = false;
+
         if let Some((target_prunable_block, prune_mode)) = self
             .prune_mode
             .map(|mode| {
@@ -77,8 +85,10 @@ where
             })
             .transpose()?
             .flatten() &&
-            target_prunable_block > input.checkpoint().block_number
+            (target_prunable_block > input.checkpoint().block_number ||
+                target_prunable_block == 0 && input.checkpoint().block_number == 0)
         {
+            prune_target_applied = true;
             input.checkpoint = Some(StageCheckpoint::new(target_prunable_block));
 
             // Save prune checkpoint only if we don't have one already.
@@ -95,38 +105,50 @@ where
             }
         }
 
-        if input.target_reached() {
+        let use_rocksdb = provider.cached_storage_settings().storage_v2;
+        if input.target_reached() && prune_target_applied {
+            if use_rocksdb {
+                provider.rocksdb_provider().clear::<tables::AccountsHistory>()?;
+            } else {
+                provider.tx_ref().clear::<tables::AccountsHistory>()?;
+            }
+            return Ok(ExecOutput::done(input.checkpoint()))
+        }
+        if input.target_reached() && !initial_sync {
             return Ok(ExecOutput::done(input.checkpoint()))
         }
 
-        let mut range = input.next_block_range();
-        let first_sync = input.checkpoint().block_number == 0;
-        let use_rocksdb = provider.cached_storage_settings().storage_v2;
+        let BlockRangeOutput { mut block_range, is_final_range } =
+            if initial_sync && input.target() == 0 {
+                BlockRangeOutput { block_range: 0..=0, is_final_range: true }
+            } else {
+                input.next_block_range_with_threshold(self.commit_threshold.max(1))
+            };
 
-        // On first sync we might have history coming from genesis. We clear the table since it's
-        // faster to rebuild from scratch.
-        if first_sync {
+        // A zero checkpoint means no history range is durable. Clear stale rows before using the
+        // append-only loader, including when pruning advanced the in-memory checkpoint.
+        if rebuild_from_empty {
             if use_rocksdb {
-                // Note: RocksDB clear() executes immediately (not deferred to commit like MDBX),
-                // but this is safe for first_sync because if we crash before commit, the
-                // checkpoint stays at 0 and we'll just clear and rebuild again on restart. The
-                // source data (changesets) is intact.
+                // RocksDB clear executes immediately. A crash before commit leaves the durable
+                // checkpoint at zero, so the next attempt clears and rebuilds again.
                 provider.rocksdb_provider().clear::<tables::AccountsHistory>()?;
             } else {
                 provider.tx_ref().clear::<tables::AccountsHistory>()?;
             }
-            range = 0..=*input.next_block_range().end();
+            if input.checkpoint().block_number == 0 && !prune_target_applied {
+                block_range = 0..=*block_range.end();
+            }
         }
 
-        info!(target: "sync::stages::index_account_history::exec", ?first_sync, ?use_rocksdb, "Collecting indices");
+        info!(target: "sync::stages::index_account_history::exec", ?rebuild_from_empty, ?use_rocksdb, "Collecting indices");
 
-        let collector = if provider.cached_storage_settings().storage_v2 {
+        let collector = if use_rocksdb {
             // Use the provider-based collection that can read from static files.
-            collect_account_history_indices(provider, range.clone(), &self.etl_config)?
+            collect_account_history_indices(provider, block_range.clone(), &self.etl_config)?
         } else {
             collect_history_indices::<_, tables::AccountChangeSets, tables::AccountsHistory, _>(
                 provider,
-                range.clone(),
+                block_range.clone(),
                 ShardedKey::new,
                 |(index, value)| (index, value.address),
                 &self.etl_config,
@@ -135,19 +157,35 @@ where
 
         info!(target: "sync::stages::index_account_history::exec", "Loading indices into database");
 
-        provider.with_rocksdb_batch_auto_commit(|rocksdb_batch| {
-            let mut writer = EitherWriter::new_accounts_history(provider, rocksdb_batch)?;
-            load_account_history(collector, first_sync, &mut writer)
-                .map_err(|e| reth_provider::ProviderError::other(Box::new(e)))?;
-            Ok(((), writer.into_raw_rocksdb_batch()))
-        })?;
+        if rebuild_from_empty {
+            provider.with_rocksdb_batch_auto_commit(|rocksdb_batch| {
+                let mut writer = EitherWriter::new_accounts_history(provider, rocksdb_batch)?;
+                load_account_history_append(collector, &mut writer)
+                    .map_err(|e| reth_provider::ProviderError::other(Box::new(e)))?;
+                Ok(((), writer.into_raw_rocksdb_batch()))
+            })?;
+        } else {
+            let prepared =
+                prepare_account_history_writes(collector, provider, use_rocksdb, &self.etl_config)?;
+            provider.with_rocksdb_batch_auto_commit(|rocksdb_batch| {
+                let mut writer = EitherWriter::new_accounts_history(provider, rocksdb_batch)?;
+                write_prepared_history_shards::<tables::AccountsHistory>(
+                    prepared,
+                    |key, value| writer.upsert_account_history(key, value),
+                )?;
+                Ok(((), writer.into_raw_rocksdb_batch()))
+            })?;
+        }
 
         if use_rocksdb {
             provider.commit_pending_rocksdb_batches()?;
             provider.rocksdb_provider().flush(&[Tables::AccountsHistory.name()])?;
         }
 
-        Ok(ExecOutput { checkpoint: StageCheckpoint::new(*range.end()), done: true })
+        Ok(ExecOutput {
+            checkpoint: StageCheckpoint::new(*block_range.end()),
+            done: is_final_range,
+        })
     }
 
     /// Unwind the stage.
@@ -286,6 +324,56 @@ mod tests {
         assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![0])]));
     }
 
+    #[tokio::test]
+    async fn initial_target_zero_indexes_genesis() {
+        let db = TestStageDB::default();
+        partial_setup(&db);
+
+        run(&db, 0, None);
+
+        let table = cast(db.table::<tables::AccountsHistory>().unwrap());
+        assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![0])]));
+    }
+
+    #[tokio::test]
+    async fn zero_commit_threshold_still_makes_progress() {
+        let db = TestStageDB::default();
+        partial_setup(&db);
+        let input = ExecInput { target: Some(1), checkpoint: None };
+        let mut stage = IndexAccountHistoryStage { commit_threshold: 0, ..Default::default() };
+        let provider = db.factory.database_provider_rw().unwrap();
+
+        let output = stage.execute(&provider, input).unwrap();
+
+        assert_eq!(output, ExecOutput { checkpoint: StageCheckpoint::new(1), done: true });
+        provider.commit().unwrap();
+        let table = cast(db.table::<tables::AccountsHistory>().unwrap());
+        assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![0, 1])]));
+    }
+
+    #[tokio::test]
+    async fn execute_respects_commit_threshold() {
+        let db = TestStageDB::default();
+        partial_setup(&db);
+        let mut stage = IndexAccountHistoryStage { commit_threshold: 2, ..Default::default() };
+
+        let mut checkpoint = None;
+        for (expected_checkpoint, done) in [(2, false), (4, false), (5, true)] {
+            let input = ExecInput { target: Some(5), checkpoint };
+            let provider = db.factory.database_provider_rw().unwrap();
+            let output = stage.execute(&provider, input).unwrap();
+            assert_eq!(
+                output,
+                ExecOutput { checkpoint: StageCheckpoint::new(expected_checkpoint), done }
+            );
+            provider.commit().unwrap();
+            checkpoint = Some(output.checkpoint);
+        }
+
+        let table = cast(db.table::<tables::AccountsHistory>().unwrap());
+        assert_eq!(table, BTreeMap::from([(shard(u64::MAX), (0..=5).collect())]));
+    }
+
     #[tokio::test]
     async fn insert_index_to_not_empty_shard() {
         // init
@@ -314,6 +402,49 @@ mod tests {
         assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![1, 2, 3])]));
     }
 
+    #[tokio::test]
+    async fn insert_index_many_distinct_accounts_incremental() {
+        let db = TestStageDB::default();
+        const KEYS: u8 = 32;
+
+        db.commit(|tx| {
+            for block in 0..=MAX_BLOCK {
+                tx.put::<tables::BlockBodyIndices>(
+                    block,
+                    StoredBlockBodyIndices { tx_count: 3, ..Default::default() },
+                )?;
+            }
+            for i in 1..=KEYS {
+                let address = Address::repeat_byte(i);
+                tx.put::<tables::AccountsHistory>(
+                    ShardedKey::new(address, u64::MAX),
+                    list(&[1, 2, 3]),
+                )?;
+                for block in 0..=MAX_BLOCK {
+                    tx.put::<tables::AccountChangeSets>(
+                        block,
+                        AccountBeforeTx { address, info: None },
+                    )?;
+                }
+            }
+            Ok(())
+        })
+        .unwrap();
+
+        run(&db, 5, Some(3));
+
+        let table = db.table::<tables::AccountsHistory>().unwrap();
+        assert_eq!(table.len(), KEYS as usize);
+        for i in 1..=KEYS {
+            let address = Address::repeat_byte(i);
+            let shard = table
+                .iter()
+                .find(|(key, _)| key.key == address && key.highest_block_number == u64::MAX)
+                .map(|(_, list)| list.iter().collect::<Vec<_>>());
+            assert_eq!(shard, Some(vec![1, 2, 3, 4, 5]), "address {address}");
+        }
+    }
+
     #[tokio::test]
     async fn insert_index_to_full_shard() {
         // init
@@ -494,6 +625,7 @@ mod tests {
             tx.put::<tables::AccountChangeSets>(20, acc()).unwrap();
             tx.put::<tables::AccountChangeSets>(36, acc()).unwrap();
             tx.put::<tables::AccountChangeSets>(100, acc()).unwrap();
+            tx.put::<tables::AccountsHistory>(shard(u64::MAX), list(&[1])).unwrap();
             Ok(())
         })
         .unwrap();
@@ -521,6 +653,34 @@ mod tests {
         assert!(table.is_empty());
     }
 
+    #[tokio::test]
+    async fn prune_target_zero_excludes_genesis() {
+        let db = TestStageDB::default();
+        db.commit(|tx| {
+            tx.put::<tables::AccountChangeSets>(0, acc())?;
+            tx.put::<tables::AccountChangeSets>(1, acc())?;
+            Ok(())
+        })
+        .unwrap();
+
+        let input = ExecInput { target: Some(20_000), ..Default::default() };
+        let mut stage = IndexAccountHistoryStage {
+            prune_mode: Some(PruneMode::Before(1)),
+            ..Default::default()
+        };
+        let provider = db.factory.database_provider_rw().unwrap();
+        let output = stage.execute(&provider, input).unwrap();
+        assert_eq!(output, ExecOutput { checkpoint: StageCheckpoint::new(20_000), done: true });
+        provider.commit().unwrap();
+
+        let table = cast(db.table::<tables::AccountsHistory>().unwrap());
+        assert_eq!(table, BTreeMap::from([(shard(u64::MAX), vec![1])]));
+        let provider = db.factory.provider().unwrap();
+        let prune_checkpoint =
+            provider.get_prune_checkpoint(PruneSegment::AccountHistory).unwrap().unwrap();
+        assert_eq!(prune_checkpoint.block_number, Some(0));
+    }
+
     stage_test_suite_ext!(IndexAccountHistoryTestRunner, index_account_history);
 
     struct IndexAccountHistoryTestRunner {
@@ -724,6 +884,38 @@ mod tests {
             assert_eq!(blocks, (0..=10).collect::<Vec<_>>());
         }
 
+        #[tokio::test]
+        async fn prune_bump_rebuilds_empty_table_then_merges_later_chunks() {
+            let db = TestStageDB::default();
+            setup_v2_account_data(&db, 0..=10);
+            let rocksdb = db.factory.rocksdb_provider();
+            rocksdb.put::<tables::AccountsHistory>(shard(u64::MAX), &list(&[1])).unwrap();
+
+            let mut stage = IndexAccountHistoryStage {
+                commit_threshold: 2,
+                prune_mode: Some(PruneMode::Before(6)),
+                ..Default::default()
+            };
+            let mut checkpoint = None;
+            for (threshold, expected_checkpoint, done) in
+                [(2, 7, false), (2, 9, false), (u64::MAX, 20_000, true)]
+            {
+                stage.commit_threshold = threshold;
+                let input = ExecInput { target: Some(20_000), checkpoint };
+                let provider = db.factory.database_provider_rw().unwrap();
+                let output = stage.execute(&provider, input).unwrap();
+                assert_eq!(
+                    output,
+                    ExecOutput { checkpoint: StageCheckpoint::new(expected_checkpoint), done }
+                );
+                provider.commit().unwrap();
+                checkpoint = Some(output.checkpoint);
+            }
+
+            let result = rocksdb.get::<tables::AccountsHistory>(shard(u64::MAX)).unwrap().unwrap();
+            assert_eq!(result.iter().collect::<Vec<_>>(), (6..=10).collect::<Vec<_>>());
+        }
+
         /// Test that unwind works correctly when `account_history_in_rocksdb` is enabled.
         #[tokio::test]
         async fn unwind_works_when_rocksdb_enabled() {