make shrink use in-mem idx info (#2689)

* make shrink use in-mem idx info

* fix cli arg typo

* remove has_written_abnormal_entry_to_disk_flag.

* pr reviews: local counter

* remove the field that hold index entry in memory for shrinking

* clippy

* default test/bench config to use onlyabnormalwithverify scan option

* remove comments

* share shrink account populate code

---------

Co-authored-by: HaoranYi <haoran.yi@anza.xyz>

accounts-db/src/accounts_index.rs

@@ -78,6 +78,22 @@ pub(crate) struct GenerateIndexResult<T: IndexValue> {
     pub duplicates: Option<Vec<(Pubkey, (Slot, T))>>,
 }
 
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
+/// which accounts `scan` should load from disk
+pub enum ScanFilter {
+    /// Scan both in-memory and on-disk index
+    #[default]
+    All,
+
+    /// abnormal = ref_count != 1 or slot list.len() != 1
+    /// Scan only in-memory index and skip on-disk index
+    OnlyAbnormal,
+
+    /// Similar to `OnlyAbnormal but also check on-disk index to verify the
+    /// entry on-disk is indeed normal.
+    OnlyAbnormalWithVerify,
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 /// how accounts index 'upsert' should handle reclaims
 pub enum UpsertReclaim {
@@ -1399,6 +1415,7 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
         mut callback: F,
         avoid_callback_result: Option<AccountsIndexScanResult>,
         provide_entry_in_callback: bool,
+        filter: ScanFilter,
     ) where
         F: FnMut(
             &'a Pubkey,
@@ -1416,10 +1433,8 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
                 lock = Some(&self.account_maps[bin]);
                 last_bin = bin;
             }
-            // SAFETY: The caller must ensure that if `provide_entry_in_callback` is true, and
-            // if it's possible for `callback` to clone the entry Arc, then it must also add
-            // the entry to the in-mem cache if the entry is made dirty.
-            lock.as_ref().unwrap().get_internal(pubkey, |entry| {
+
+            let mut internal_callback = |entry: Option<&AccountMapEntry<T>>| {
                 let mut cache = false;
                 match entry {
                     Some(locked_entry) => {
@@ -1449,7 +1464,36 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
                     }
                 }
                 (cache, ())
-            });
+            };
+
+            match filter {
+                ScanFilter::All => {
+                    // SAFETY: The caller must ensure that if `provide_entry_in_callback` is true, and
+                    // if it's possible for `callback` to clone the entry Arc, then it must also add
+                    // the entry to the in-mem cache if the entry is made dirty.
+                    lock.as_ref()
+                        .unwrap()
+                        .get_internal(pubkey, internal_callback);
+                }
+                ScanFilter::OnlyAbnormal | ScanFilter::OnlyAbnormalWithVerify => {
+                    let found = lock
+                        .as_ref()
+                        .unwrap()
+                        .get_only_in_mem(pubkey, false, |entry| {
+                            internal_callback(entry);
+                            entry.is_some()
+                        });
+                    if !found && matches!(filter, ScanFilter::OnlyAbnormalWithVerify) {
+                        lock.as_ref().unwrap().get_internal(pubkey, |entry| {
+                            assert!(entry.is_some(), "{pubkey}, entry: {entry:?}");
+                            let entry = entry.unwrap();
+                            assert_eq!(entry.ref_count(), 1, "{pubkey}");
+                            assert_eq!(entry.slot_list.read().unwrap().len(), 1, "{pubkey}");
+                            (false, ())
+                        });
+                    }
+                }
+            }
         });
     }
 

accounts-db/src/accounts_index/in_mem_accounts_index.rs

@@ -281,7 +281,7 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
 
     /// lookup 'pubkey' by only looking in memory. Does not look on disk.
     /// callback is called whether pubkey is found or not
-    fn get_only_in_mem<RT>(
+    pub(super) fn get_only_in_mem<RT>(
         &self,
         pubkey: &K,
         update_age: bool,

accounts-db/src/ancient_append_vecs.rs

@@ -13,7 +13,7 @@ use {
             ShrinkCollectAliveSeparatedByRefs, ShrinkStatsSub,
         },
         accounts_file::AccountsFile,
-        accounts_index::AccountsIndexScanResult,
+        accounts_index::{AccountsIndexScanResult, ScanFilter},
         active_stats::ActiveStatItem,
         storable_accounts::{StorableAccounts, StorableAccountsBySlot},
     },
@@ -500,6 +500,7 @@ impl AccountsDb {
                     },
                     None,
                     true,
+                    ScanFilter::All,
                 );
             });
         });
@@ -3876,7 +3877,6 @@ pub mod tests {
                 alive_total_bytes: 0,
                 total_starting_accounts: 0,
                 all_are_zero_lamports: false,
-                _index_entries_being_shrunk: Vec::default(),
             };
             let accounts_to_combine = AccountsToCombine {
                 accounts_keep_slots: HashMap::default(),
@@ -3893,6 +3893,7 @@ pub mod tests {
                 },
                 None,
                 false,
+                ScanFilter::All,
             );
             // should have removed all of them
             assert!(expected_ref_counts.is_empty());

ledger-tool/src/args.rs

@@ -4,7 +4,7 @@ use {
     solana_accounts_db::{
         accounts_db::{AccountsDb, AccountsDbConfig, CreateAncientStorage},
         accounts_file::StorageAccess,
-        accounts_index::{AccountsIndexConfig, IndexLimitMb},
+        accounts_index::{AccountsIndexConfig, IndexLimitMb, ScanFilter},
         partitioned_rewards::TestPartitionedEpochRewards,
         utils::create_and_canonicalize_directories,
     },
@@ -78,6 +78,20 @@ pub fn accounts_db_args<'a, 'b>() -> Box<[Arg<'a, 'b>]> {
                 clean",
             )
             .hidden(hidden_unless_forced()),
+        Arg::with_name("accounts_db_scan_filter_for_shrinking")
+            .long("accounts-db-scan-filter-for-shrinking")
+            .takes_value(true)
+            .possible_values(&["all", "only-abnormal", "only-abnormal-with-verify"])
+            .help(
+                "Debug option to use different type of filtering for accounts index scan in \
+                shrinking. \"all\" will scan both in-memory and on-disk accounts index, which is the default. \
+                \"only-abnormal\" will scan in-memory accounts index only for abnormal entries and \
+                skip scanning on-disk accounts index by assuming that on-disk accounts index contains \
+                only normal accounts index entry. \"only-abnormal-with-verify\" is similar to \
+                \"only-abnormal\", which will scan in-memory index for abnormal entries, but will also \
+                verify that on-disk account entries are indeed normal.",
+            )
+            .hidden(hidden_unless_forced()),
         Arg::with_name("accounts_db_test_skip_rewrites")
             .long("accounts-db-test-skip-rewrites")
             .help(
@@ -296,6 +310,19 @@ pub fn get_accounts_db_config(
         })
         .unwrap_or_default();
 
+    let scan_filter_for_shrinking = arg_matches
+        .value_of("accounts_db_scan_filter_for_shrinking")
+        .map(|filter| match filter {
+            "all" => ScanFilter::All,
+            "only-abnormal" => ScanFilter::OnlyAbnormal,
+            "only-abnormal-with-verify" => ScanFilter::OnlyAbnormalWithVerify,
+            _ => {
+                // clap will enforce one of the above values is given
+                unreachable!("invalid value given to accounts_db_scan_filter_for_shrinking")
+            }
+        })
+        .unwrap_or_default();
+
     AccountsDbConfig {
         index: Some(accounts_index_config),
         base_working_path: Some(ledger_tool_ledger_path),
@@ -309,6 +336,7 @@ pub fn get_accounts_db_config(
             .is_present("accounts_db_test_skip_rewrites"),
         create_ancient_storage,
         storage_access,
+        scan_filter_for_shrinking,
         ..AccountsDbConfig::default()
     }
 }

validator/src/cli.rs

@@ -1313,6 +1313,22 @@ pub fn app<'a>(version: &'a str, default_args: &'a DefaultArgs) -> App<'a, 'a> {
                 )
                 .hidden(hidden_unless_forced()),
         )
+        .arg(
+            Arg::with_name("accounts_db_scan_filter_for_shrinking")
+                .long("accounts-db-scan-filter-for-shrinking")
+                .takes_value(true)
+                .possible_values(&["all", "only-abnormal", "only-abnormal-with-verify"])
+                .help(
+                    "Debug option to use different type of filtering for accounts index scan in \
+                    shrinking. \"all\" will scan both in-memory and on-disk accounts index, which is the default. \
+                    \"only-abnormal\" will scan in-memory accounts index only for abnormal entries and \
+                    skip scanning on-disk accounts index by assuming that on-disk accounts index contains \
+                    only normal accounts index entry. \"only-abnormal-with-verify\" is similar to \
+                    \"only-abnormal\", which will scan in-memory index for abnormal entries, but will also \
+                    verify that on-disk account entries are indeed normal.",
+                )
+                .hidden(hidden_unless_forced()),
+        )
         .arg(
             Arg::with_name("accounts_db_test_skip_rewrites")
                 .long("accounts-db-test-skip-rewrites")

validator/src/main.rs

@@ -21,7 +21,7 @@ use {
         accounts_file::StorageAccess,
         accounts_index::{
             AccountIndex, AccountSecondaryIndexes, AccountSecondaryIndexesIncludeExclude,
-            AccountsIndexConfig, IndexLimitMb,
+            AccountsIndexConfig, IndexLimitMb, ScanFilter,
         },
         partitioned_rewards::TestPartitionedEpochRewards,
         utils::{create_all_accounts_run_and_snapshot_dirs, create_and_canonicalize_directories},
@@ -1271,6 +1271,19 @@ pub fn main() {
         })
         .unwrap_or_default();
 
+    let scan_filter_for_shrinking = matches
+        .value_of("accounts_db_scan_filter_for_shrinking")
+        .map(|filter| match filter {
+            "all" => ScanFilter::All,
+            "only-abnormal" => ScanFilter::OnlyAbnormal,
+            "only-abnormal-with-verify" => ScanFilter::OnlyAbnormalWithVerify,
+            _ => {
+                // clap will enforce one of the above values is given
+                unreachable!("invalid value given to accounts_db_scan_filter_for_shrinking")
+            }
+        })
+        .unwrap_or_default();
+
     let accounts_db_config = AccountsDbConfig {
         index: Some(accounts_index_config),
         base_working_path: Some(ledger_path.clone()),
@@ -1287,6 +1300,7 @@ pub fn main() {
         test_skip_rewrites_but_include_in_bank_hash: matches
             .is_present("accounts_db_test_skip_rewrites"),
         storage_access,
+        scan_filter_for_shrinking,
         ..AccountsDbConfig::default()
     };