AcctIdx: respect disk idx mem size param (#22050)

runtime/src/in_mem_accounts_index.rs

@@ -806,31 +806,44 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
         thread_rng().gen_range(0, N) == 0
     }
 
+    /// assumes 1 entry in the slot list. Ignores overhead of the HashMap and such
+    fn approx_size_of_one_entry() -> usize {
+        std::mem::size_of::<T>()
+            + std::mem::size_of::<Pubkey>()
+            + std::mem::size_of::<AccountMapEntry<T>>()
+    }
+
     /// return true if 'entry' should be removed from the in-mem index
     fn should_remove_from_mem(
         &self,
         current_age: Age,
         entry: &AccountMapEntry<T>,
         startup: bool,
         update_stats: bool,
+        exceeds_budget: bool,
     ) -> bool {
         // this could be tunable dynamically based on memory pressure
         // we could look at more ages or we could throw out more items we are choosing to keep in the cache
         if startup || (current_age == entry.age()) {
-            // only read the slot list if we are planning to throw the item out
-            let slot_list = entry.slot_list.read().unwrap();
-            if slot_list.len() != 1 {
-                if update_stats {
-                    Self::update_stat(&self.stats().held_in_mem_slot_list_len, 1);
-                }
-                false // keep 0 and > 1 slot lists in mem. They will be cleaned or shrunk soon.
+            if exceeds_budget {
+                // if we are already holding too many items in-mem, then we need to be more aggressive at kicking things out
+                true
             } else {
-                // keep items with slot lists that contained cached items
-                let remove = !slot_list.iter().any(|(_, info)| info.is_cached());
-                if !remove && update_stats {
-                    Self::update_stat(&self.stats().held_in_mem_slot_list_cached, 1);
+                // only read the slot list if we are planning to throw the item out
+                let slot_list = entry.slot_list.read().unwrap();
+                if slot_list.len() != 1 {
+                    if update_stats {
+                        Self::update_stat(&self.stats().held_in_mem_slot_list_len, 1);
+                    }
+                    false // keep 0 and > 1 slot lists in mem. They will be cleaned or shrunk soon.
+                } else {
+                    // keep items with slot lists that contained cached items
+                    let remove = !slot_list.iter().any(|(_, info)| info.is_cached());
+                    if !remove && update_stats {
+                        Self::update_stat(&self.stats().held_in_mem_slot_list_cached, 1);
+                    }
+                    remove
                 }
-                remove
             }
         } else {
             false
@@ -848,6 +861,12 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
             return;
         }
 
+        let in_mem_count = self.stats().count_in_mem.load(Ordering::Relaxed);
+        let limit = self.storage.mem_budget_mb;
+        let exceeds_budget = limit
+            .map(|limit| in_mem_count * Self::approx_size_of_one_entry() >= limit * 1024 * 1024)
+            .unwrap_or_default();
+
         // may have to loop if disk has to grow and we have to restart
         loop {
             let mut removes;
@@ -863,7 +882,7 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
                 removes = Vec::with_capacity(map.len());
                 let m = Measure::start("flush_scan_and_update"); // we don't care about lock time in this metric - bg threads can wait
                 for (k, v) in map.iter() {
-                    if self.should_remove_from_mem(current_age, v, startup, true) {
+                    if self.should_remove_from_mem(current_age, v, startup, true, exceeds_budget) {
                         removes.push(*k);
                     } else if Self::random_chance_of_eviction() {
                         removes_random.push(*k);
@@ -902,9 +921,19 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
             let m = Measure::start("flush_remove_or_grow");
             match disk_resize {
                 Ok(_) => {
-                    if !self.flush_remove_from_cache(removes, current_age, startup, false)
-                        || !self.flush_remove_from_cache(removes_random, current_age, startup, true)
-                    {
+                    if !self.flush_remove_from_cache(
+                        removes,
+                        current_age,
+                        startup,
+                        false,
+                        exceeds_budget,
+                    ) || !self.flush_remove_from_cache(
+                        removes_random,
+                        current_age,
+                        startup,
+                        true,
+                        exceeds_budget,
+                    ) {
                         iterate_for_age = false; // did not make it all the way through this bucket, so didn't handle age completely
                     }
                     Self::update_time_stat(&self.stats().flush_remove_us, m);
@@ -934,6 +963,7 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
         current_age: Age,
         startup: bool,
         randomly_evicted: bool,
+        exceeds_budget: bool,
     ) -> bool {
         let mut completed_scan = true;
         if removes.is_empty() {
@@ -956,7 +986,13 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
 
                 if v.dirty()
                     || (!randomly_evicted
-                        && !self.should_remove_from_mem(current_age, v, startup, false))
+                        && !self.should_remove_from_mem(
+                            current_age,
+                            v,
+                            startup,
+                            false,
+                            exceeds_budget,
+                        ))
                 {
                     // marked dirty or bumped in age after we looked above
                     // these will be handled in later passes
@@ -1053,6 +1089,18 @@ mod tests {
             AccountMapEntryMeta::default(),
         ));
 
+        // exceeded budget
+        assert!(bucket.should_remove_from_mem(
+            current_age,
+            &Arc::new(AccountMapEntryInner::new(
+                vec![],
+                ref_count,
+                AccountMapEntryMeta::default()
+            )),
+            startup,
+            false,
+            true,
+        ));
         // empty slot list
         assert!(!bucket.should_remove_from_mem(
             current_age,
@@ -1063,13 +1111,15 @@ mod tests {
             )),
             startup,
             false,
+            false,
         ));
         // 1 element slot list
         assert!(bucket.should_remove_from_mem(
             current_age,
             &one_element_slot_list_entry,
             startup,
             false,
+            false,
         ));
         // 2 element slot list
         assert!(!bucket.should_remove_from_mem(
@@ -1081,6 +1131,7 @@ mod tests {
             )),
             startup,
             false,
+            false,
         ));
 
         {
@@ -1095,6 +1146,7 @@ mod tests {
                 )),
                 startup,
                 false,
+                false,
             ));
         }
 
@@ -1104,6 +1156,7 @@ mod tests {
             &one_element_slot_list_entry,
             startup,
             false,
+            false,
         ));
 
         // 1 element slot list, but not current age
@@ -1113,6 +1166,7 @@ mod tests {
             &one_element_slot_list_entry,
             startup,
             false,
+            false,
         ));
 
         // 1 element slot list, but at startup and age not current
@@ -1122,6 +1176,7 @@ mod tests {
             &one_element_slot_list_entry,
             startup,
             false,
+            false,
         ));
     }