add --limit_load_slot_count_from_snapshot to ledger-tool (#17417)

core/tests/snapshots.rs

@@ -164,6 +164,7 @@ mod tests {
             None,
             AccountSecondaryIndexes::default(),
             false,
+            None,
         )
         .unwrap();
 

ledger-tool/src/main.rs

@@ -771,6 +771,12 @@ fn main() {
         .validator(is_slot)
         .takes_value(true)
         .help("Halt processing at the given slot");
+    let limit_load_slot_count_from_snapshot_arg = Arg::with_name("limit_load_slot_count_from_snapshot")
+        .long("limit-load-slot-count-from-snapshot")
+        .value_name("SLOT")
+        .validator(is_slot)
+        .takes_value(true)
+        .help("For debugging and profiling with large snapshots, artificially limit how many slots are loaded from a snapshot.");
     let hard_forks_arg = Arg::with_name("hard_forks")
         .long("hard-fork")
         .value_name("SLOT")
@@ -1040,6 +1046,7 @@ fn main() {
             .arg(&no_snapshot_arg)
             .arg(&account_paths_arg)
             .arg(&halt_at_slot_arg)
+            .arg(&limit_load_slot_count_from_snapshot_arg)
             .arg(&hard_forks_arg)
             .arg(&no_accounts_db_caching_arg)
             .arg(&accounts_db_test_hash_calculation_arg)
@@ -1744,6 +1751,12 @@ fn main() {
                 poh_verify: !arg_matches.is_present("skip_poh_verify"),
                 bpf_jit: !matches.is_present("no_bpf_jit"),
                 accounts_db_caching_enabled: !arg_matches.is_present("no_accounts_db_caching"),
+                limit_load_slot_count_from_snapshot: value_t!(
+                    arg_matches,
+                    "limit_load_slot_count_from_snapshot",
+                    usize
+                )
+                .ok(),
                 allow_dead_slots: arg_matches.is_present("allow_dead_slots"),
                 accounts_db_test_hash_calculation: arg_matches
                     .is_present("accounts_db_test_hash_calculation"),

ledger/src/bank_forks_utils.rs

@@ -71,6 +71,7 @@ pub fn load(
                     Some(&crate::builtins::get(process_options.bpf_jit)),
                     process_options.account_indexes.clone(),
                     process_options.accounts_db_caching_enabled,
+                    process_options.limit_load_slot_count_from_snapshot,
                 )
                 .expect("Load from snapshot failed");
                 if let Some(shrink_paths) = shrink_paths {

ledger/src/blockstore_processor.rs

@@ -368,6 +368,7 @@ pub struct ProcessOptions {
     pub debug_keys: Option<Arc<HashSet<Pubkey>>>,
     pub account_indexes: AccountSecondaryIndexes,
     pub accounts_db_caching_enabled: bool,
+    pub limit_load_slot_count_from_snapshot: Option<usize>,
     pub allow_dead_slots: bool,
     pub accounts_db_test_hash_calculation: bool,
 }

runtime/src/accounts_db.rs

@@ -5123,12 +5123,15 @@ impl AccountsDb {
     }
 
     #[allow(clippy::needless_collect)]
-    pub fn generate_index(&self) {
+    pub fn generate_index(&self, limit_load_slot_count_from_snapshot: Option<usize>) {
         type AccountsMap<'a> =
             HashMap<Pubkey, (StoredMetaWriteVersion, AppendVecId, StoredAccountMeta<'a>)>;
         let mut slots = self.storage.all_slots();
         #[allow(clippy::stable_sort_primitive)]
         slots.sort();
+        if let Some(limit) = limit_load_slot_count_from_snapshot {
+            slots.truncate(limit); // get rid of the newer slots and keep just the older
+        }
         let total_processed_slots_across_all_threads = AtomicU64::new(0);
         let outer_slots_len = slots.len();
         let chunk_size = (outer_slots_len / 7) + 1; // approximately 400k slots in a snapshot

runtime/src/serde_snapshot.rs

@@ -135,6 +135,7 @@ pub(crate) fn bank_from_stream<R>(
     additional_builtins: Option<&Builtins>,
     account_indexes: AccountSecondaryIndexes,
     caching_enabled: bool,
+    limit_load_slot_count_from_snapshot: Option<usize>,
 ) -> std::result::Result<Bank, Error>
 where
     R: Read,
@@ -154,6 +155,7 @@ where
                 additional_builtins,
                 account_indexes,
                 caching_enabled,
+                limit_load_slot_count_from_snapshot,
             )?;
             Ok(bank)
         }};
@@ -243,6 +245,7 @@ fn reconstruct_bank_from_fields<E>(
     additional_builtins: Option<&Builtins>,
     account_indexes: AccountSecondaryIndexes,
     caching_enabled: bool,
+    limit_load_slot_count_from_snapshot: Option<usize>,
 ) -> Result<Bank, Error>
 where
     E: SerializableStorage,
@@ -254,6 +257,7 @@ where
         &genesis_config.cluster_type,
         account_indexes,
         caching_enabled,
+        limit_load_slot_count_from_snapshot,
     )?;
     accounts_db.freeze_accounts(
         &Ancestors::from(&bank_fields.ancestors),
@@ -279,6 +283,7 @@ fn reconstruct_accountsdb_from_fields<E>(
     cluster_type: &ClusterType,
     account_indexes: AccountSecondaryIndexes,
     caching_enabled: bool,
+    limit_load_slot_count_from_snapshot: Option<usize>,
 ) -> Result<AccountsDb, Error>
 where
     E: SerializableStorage,
@@ -371,6 +376,6 @@ where
     accounts_db
         .write_version
         .fetch_add(version, Ordering::Relaxed);
-    accounts_db.generate_index();
+    accounts_db.generate_index(limit_load_slot_count_from_snapshot);
     Ok(accounts_db)
 }

runtime/src/serde_snapshot/tests.rs

@@ -73,6 +73,7 @@ where
         &ClusterType::Development,
         AccountSecondaryIndexes::default(),
         false,
+        None,
     )
 }
 
@@ -226,6 +227,7 @@ fn test_bank_serialize_style(serde_style: SerdeStyle) {
         None,
         AccountSecondaryIndexes::default(),
         false,
+        None,
     )
     .unwrap();
     dbank.src = ref_sc;

runtime/src/snapshot_utils.rs

@@ -603,6 +603,7 @@ pub fn bank_from_archive<P: AsRef<Path>>(
     additional_builtins: Option<&Builtins>,
     account_indexes: AccountSecondaryIndexes,
     accounts_db_caching_enabled: bool,
+    limit_load_slot_count_from_snapshot: Option<usize>,
 ) -> Result<Bank> {
     let unpack_dir = tempfile::Builder::new()
         .prefix(TMP_SNAPSHOT_PREFIX)
@@ -633,6 +634,7 @@ pub fn bank_from_archive<P: AsRef<Path>>(
         additional_builtins,
         account_indexes,
         accounts_db_caching_enabled,
+        limit_load_slot_count_from_snapshot,
     )?;
 
     if !bank.verify_snapshot_bank() {
@@ -791,6 +793,7 @@ fn rebuild_bank_from_snapshots(
     additional_builtins: Option<&Builtins>,
     account_indexes: AccountSecondaryIndexes,
     accounts_db_caching_enabled: bool,
+    limit_load_slot_count_from_snapshot: Option<usize>,
 ) -> Result<Bank> {
     info!("snapshot version: {}", snapshot_version);
 
@@ -826,6 +829,7 @@ fn rebuild_bank_from_snapshots(
                 additional_builtins,
                 account_indexes,
                 accounts_db_caching_enabled,
+                limit_load_slot_count_from_snapshot,
             ),
         }?)
     })?;