add --accounts-index-memory-limit-mb (#19269)

ledger-tool/src/main.rs

@@ -880,6 +880,12 @@ fn main() {
         .validator(is_bin)
         .takes_value(true)
         .help("Number of bins to divide the accounts index into");
+    let accounts_index_limit = Arg::with_name("accounts_index_memory_limit_mb")
+        .long("accounts-index-memory-limit-mb")
+        .value_name("MEGABYTES")
+        .validator(is_parsable::<usize>)
+        .takes_value(true)
+        .help("How much memory the accounts index can consume. If this is exceeded, some account index entries will be stored on disk. If missing, the entire index is stored in memory.");
     let account_paths_arg = Arg::with_name("account_paths")
         .long("accounts")
         .value_name("PATHS")
@@ -1192,6 +1198,7 @@ fn main() {
             .arg(&halt_at_slot_arg)
             .arg(&limit_load_slot_count_from_snapshot_arg)
             .arg(&accounts_index_bins)
+            .arg(&accounts_index_limit)
             .arg(&verify_index_arg)
             .arg(&hard_forks_arg)
             .arg(&no_accounts_db_caching_arg)
@@ -1913,6 +1920,10 @@ fn main() {
                 accounts_index_config.bins = Some(bins);
             }
 
+            if let Some(limit) = value_t!(matches, "accounts_index_memory_limit_mb", usize).ok() {
+                accounts_index_config.index_limit_mb = Some(limit);
+            }
+
             {
                 let mut accounts_index_paths = vec![]; // will be option
                 if accounts_index_paths.is_empty() {

runtime/src/accounts_index.rs

@@ -40,11 +40,13 @@ pub const ACCOUNTS_INDEX_CONFIG_FOR_TESTING: AccountsIndexConfig = AccountsIndex
     bins: Some(BINS_FOR_TESTING),
     flush_threads: Some(FLUSH_THREADS_TESTING),
     drives: None,
+    index_limit_mb: None,
 };
 pub const ACCOUNTS_INDEX_CONFIG_FOR_BENCHMARKS: AccountsIndexConfig = AccountsIndexConfig {
     bins: Some(BINS_FOR_BENCHMARKS),
     flush_threads: Some(FLUSH_THREADS_TESTING),
     drives: None,
+    index_limit_mb: None,
 };
 pub type ScanResult<T> = Result<T, ScanError>;
 pub type SlotList<T> = Vec<(Slot, T)>;
@@ -101,6 +103,7 @@ pub struct AccountsIndexConfig {
     pub bins: Option<usize>,
     pub flush_threads: Option<usize>,
     pub drives: Option<Vec<PathBuf>>,
+    pub index_limit_mb: Option<usize>,
 }
 
 #[derive(Debug, Default, Clone)]

runtime/src/bucket_map_holder.rs

@@ -17,6 +17,10 @@ pub struct BucketMapHolder<T: IndexValue> {
     next_bucket_to_flush: Mutex<usize>,
     bins: usize,
 
+    // how much mb are we allowed to keep in the in-mem index?
+    // Rest goes to disk.
+    pub mem_budget_mb: Option<usize>,
+
     /// startup is a special time for flush to focus on moving everything to disk as fast and efficiently as possible
     /// with less thread count limitations. LRU and access patterns are not important. Freeing memory
     /// and writing to disk in parallel are.
@@ -71,7 +75,7 @@ impl<T: IndexValue> BucketMapHolder<T> {
         self.count_ages_flushed.load(Ordering::Relaxed) >= self.bins
     }
 
-    pub fn new(bins: usize, _config: &Option<AccountsIndexConfig>) -> Self {
+    pub fn new(bins: usize, config: &Option<AccountsIndexConfig>) -> Self {
         Self {
             count_ages_flushed: AtomicUsize::default(),
             age: AtomicU8::default(),
@@ -80,6 +84,7 @@ impl<T: IndexValue> BucketMapHolder<T> {
             next_bucket_to_flush: Mutex::new(0),
             bins,
             startup: AtomicBool::default(),
+            mem_budget_mb: config.as_ref().and_then(|config| config.index_limit_mb),
             _phantom: std::marker::PhantomData::<T>::default(),
         }
     }

validator/src/main.rs

@@ -1990,6 +1990,14 @@ pub fn main() {
                 .help("Enables faster starting of validators by skipping shrink. \
                       This option is for use during testing."),
         )
+        .arg(
+            Arg::with_name("accounts_index_memory_limit_mb")
+                .long("accounts-index-memory-limit-mb")
+                .value_name("MEGABYTES")
+                .validator(is_parsable::<usize>)
+                .takes_value(true)
+                .help("How much memory the accounts index can consume. If this is exceeded, some account index entries will be stored on disk. If missing, the entire index is stored in memory."),
+        )
         .arg(
             Arg::with_name("accounts_index_bins")
                 .long("accounts-index-bins")
@@ -2512,6 +2520,10 @@ pub fn main() {
         accounts_index_config.bins = Some(bins);
     }
 
+    if let Some(limit) = value_t!(matches, "accounts_index_memory_limit_mb", usize).ok() {
+        accounts_index_config.index_limit_mb = Some(limit);
+    }
+
     {
         let mut accounts_index_paths = vec![]; // will be option soon
         if accounts_index_paths.is_empty() {