add --accounts-index-scan-results-limit-mb to allow scans to abort (#21327)

* ScanConfig -> &ScanConfig

* add --accounts-index-scan-results-limit-mb to allow scans to abort

* feedback

runtime/src/accounts.rs

@@ -4,7 +4,7 @@ use crate::{
         LoadHint, LoadedAccount, ScanStorageResult, ACCOUNTS_DB_CONFIG_FOR_BENCHMARKS,
         ACCOUNTS_DB_CONFIG_FOR_TESTING,
     },
-    accounts_index::{AccountSecondaryIndexes, IndexKey, ScanConfig, ScanResult},
+    accounts_index::{AccountSecondaryIndexes, IndexKey, ScanConfig, ScanError, ScanResult},
     accounts_update_notifier_interface::AccountsUpdateNotifier,
     ancestors::Ancestors,
     bank::{
@@ -41,6 +41,7 @@ use std::{
     collections::{hash_map, BinaryHeap, HashMap, HashSet},
     ops::RangeBounds,
     path::PathBuf,
+    sync::atomic::{AtomicUsize, Ordering},
     sync::{Arc, Mutex},
 };
 
@@ -665,7 +666,7 @@ impl Accounts {
                     collector.push(Reverse((account.lamports(), *pubkey)));
                 }
             },
-            ScanConfig::default(),
+            &ScanConfig::default(),
         )?;
         Ok(account_balances
             .into_sorted_vec()
@@ -743,7 +744,7 @@ impl Accounts {
         ancestors: &Ancestors,
         bank_id: BankId,
         program_id: &Pubkey,
-        config: ScanConfig,
+        config: &ScanConfig,
     ) -> ScanResult<Vec<(Pubkey, AccountSharedData)>> {
         self.accounts_db.scan_accounts(
             ancestors,
@@ -763,7 +764,7 @@ impl Accounts {
         bank_id: BankId,
         program_id: &Pubkey,
         filter: F,
-        config: ScanConfig,
+        config: &ScanConfig,
     ) -> ScanResult<Vec<(Pubkey, AccountSharedData)>> {
         self.accounts_db.scan_accounts(
             ancestors,
@@ -783,21 +784,51 @@ impl Accounts {
         bank_id: BankId,
         index_key: &IndexKey,
         filter: F,
-        config: ScanConfig,
+        config: &ScanConfig,
+        byte_limit_for_scan: Option<usize>,
     ) -> ScanResult<Vec<(Pubkey, AccountSharedData)>> {
-        self.accounts_db
+        let sum = AtomicUsize::default();
+        let config = ScanConfig {
+            abort: Some(config.abort.as_ref().map(Arc::clone).unwrap_or_default()),
+            collect_all_unsorted: config.collect_all_unsorted,
+        };
+        let result = self
+            .accounts_db
             .index_scan_accounts(
                 ancestors,
                 bank_id,
                 *index_key,
                 |collector: &mut Vec<(Pubkey, AccountSharedData)>, some_account_tuple| {
                     Self::load_while_filtering(collector, some_account_tuple, |account| {
-                        filter(account)
-                    })
+                        let use_account = filter(account);
+                        if use_account {
+                            if let Some(byte_limit_for_scan) = byte_limit_for_scan.as_ref() {
+                                let added = account.data().len()
+                                    + std::mem::size_of::<AccountSharedData>()
+                                    + std::mem::size_of::<Pubkey>();
+                                if sum
+                                    .fetch_add(added, Ordering::Relaxed)
+                                    .saturating_add(added)
+                                    > *byte_limit_for_scan
+                                {
+                                    // total size of results exceeds size limit, so abort scan
+                                    config.abort();
+                                }
+                            }
+                        }
+                        use_account
+                    });
                 },
-                config,
+                &config,
             )
-            .map(|result| result.0)
+            .map(|result| result.0);
+        if config.is_aborted() {
+            ScanResult::Err(ScanError::Aborted(
+                "The accumulated scan results exceeded the limit".to_string(),
+            ))
+        } else {
+            result
+        }
     }
 
     pub fn account_indexes_include_key(&self, key: &Pubkey) -> bool {
@@ -819,7 +850,7 @@ impl Accounts {
                     collector.push((*pubkey, account, slot))
                 }
             },
-            ScanConfig::default(),
+            &ScanConfig::default(),
         )
     }
 
@@ -841,7 +872,7 @@ impl Accounts {
             "load_to_collect_rent_eagerly_scan_elapsed",
             ancestors,
             range,
-            ScanConfig::new(true),
+            &ScanConfig::new(true),
             |collector: &mut Vec<(Pubkey, AccountSharedData)>, option| {
                 Self::load_while_filtering(collector, option, |_| true)
             },