Adds stat for in-mem accounts index's capacity (#8067)

accounts-db/src/accounts_db.rs

@@ -6869,6 +6869,21 @@ impl AccountsDb {
 
         self.accounts_index.log_secondary_indexes();
 
+        // Now that the index is generated, get the total capacity of the in-mem maps
+        // across all the bins and set the initial value for the stat.
+        // We do this all at once, at the end, since getting the capacity requries iterating all
+        // the bins and grabbing a read lock, which we try to avoid whenever possible.
+        let index_capacity = self
+            .accounts_index
+            .account_maps
+            .iter()
+            .map(|bin| bin.capacity_for_startup())
+            .sum();
+        self.accounts_index
+            .bucket_map_holder_stats()
+            .capacity_in_mem
+            .store(index_capacity, Ordering::Relaxed);
+
         IndexGenerationInfo {
             accounts_data_len: total_accum.accounts_data_len,
             calculated_accounts_lt_hash: AccountsLtHash(total_accum.lt_hash),

accounts-db/src/accounts_index/in_mem_accounts_index.rs

@@ -387,8 +387,9 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
                 }
                 let disk_entry = disk_entry.unwrap();
                 let mut map = self.map_internal.write().unwrap();
+                let capacity_pre = map.capacity();
                 let entry = map.entry(*pubkey);
-                match entry {
+                let retval = match entry {
                     Entry::Occupied(occupied) => callback(Some(occupied.get())).1,
                     Entry::Vacant(vacant) => {
                         debug_assert!(!disk_entry.dirty());
@@ -402,7 +403,11 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
                         }
                         rt
                     }
-                }
+                };
+                let capacity_post = map.capacity();
+                drop(map);
+                stats.update_in_mem_capacity(capacity_pre, capacity_post);
+                retval
             }
         })
     }
@@ -472,12 +477,15 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
     pub fn remove_if_slot_list_empty(&self, pubkey: Pubkey) -> bool {
         let mut m = Measure::start("entry");
         let mut map = self.map_internal.write().unwrap();
+        let capacity_pre = map.capacity();
         let entry = map.entry(pubkey);
         m.stop();
         let found = matches!(entry, Entry::Occupied(_));
         let result = self.remove_if_slot_list_empty_entry(entry);
+        let capacity_post = map.capacity();
         drop(map);
-
+        self.stats()
+            .update_in_mem_capacity(capacity_pre, capacity_post);
         self.update_entry_stats(m, found);
         result
     }
@@ -559,8 +567,10 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
             if let Some(entry) = entry {
                 callback(entry);
             } else {
+                let stats = self.stats();
                 let mut m = Measure::start("entry");
                 let mut map = self.map_internal.write().unwrap();
+                let capacity_pre = map.capacity();
                 let entry = map.entry(*pubkey);
                 m.stop();
                 let found = matches!(entry, Entry::Occupied(_));
@@ -597,11 +607,12 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
                         );
                         assert!(new_value.dirty());
                         vacant.insert(new_value);
-                        self.stats().inc_mem_count();
+                        stats.inc_mem_count();
                     }
                 };
-
+                let capacity_post = map.capacity();
                 drop(map);
+                stats.update_in_mem_capacity(capacity_pre, capacity_post);
                 self.update_entry_stats(m, found);
             };
         });
@@ -1319,12 +1330,14 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
             return;
         }
 
+        let stats = self.stats();
         let next_age_on_failure = self.storage.future_age_to_flush(false);
         let mut failed = 0;
         let mut evicted = 0;
         // chunk these so we don't hold the write lock too long
         for evictions in evictions.chunks(50) {
             let mut map = self.map_internal.write().unwrap();
+            let capacity_pre = map.capacity();
             for k in evictions {
                 if let Entry::Occupied(occupied) = map.entry(*k) {
                     let v = occupied.get();
@@ -1359,10 +1372,13 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
             if map.is_empty() {
                 map.shrink_to_fit();
             }
+            let capacity_post = map.capacity();
+            drop(map);
+            stats.update_in_mem_capacity(capacity_pre, capacity_post);
         }
-        self.stats().sub_mem_count(evicted);
-        Self::update_stat(&self.stats().flush_entries_evicted_from_mem, evicted as u64);
-        Self::update_stat(&self.stats().failed_to_evict, failed as u64);
+        stats.sub_mem_count(evicted);
+        Self::update_stat(&stats.flush_entries_evicted_from_mem, evicted as u64);
+        Self::update_stat(&stats.failed_to_evict, failed as u64);
     }
 
     pub fn stats(&self) -> &BucketMapHolderStats {
@@ -1380,6 +1396,13 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
         let value = m.as_us();
         Self::update_stat(stat, value);
     }
+
+    /// Returns the capacity for this bin's map
+    ///
+    /// Only intended to be called at startup, since it grabs the map's read lock.
+    pub(crate) fn capacity_for_startup(&self) -> usize {
+        self.map_internal.read().unwrap().capacity()
+    }
 }
 
 /// An RAII implementation of a scoped lock for the `flushing_active` atomic flag in

accounts-db/src/bucket_map_holder_stats.rs

@@ -48,6 +48,7 @@ pub struct BucketMapHolderStats {
     pub bg_waiting_us: AtomicU64,
     pub bg_throttling_wait_us: AtomicU64,
     pub count_in_mem: AtomicUsize,
+    pub capacity_in_mem: AtomicUsize,
     pub flush_entries_updated_on_disk: AtomicU64,
     pub flush_entries_evicted_from_mem: AtomicU64,
     pub active_threads: AtomicU64,
@@ -102,6 +103,23 @@ impl BucketMapHolderStats {
         self.count_in_mem.fetch_sub(count, Ordering::Relaxed);
     }
 
+    /// Updates the 'in-mem capacity' stat, given a bin's pre and post values
+    pub fn update_in_mem_capacity(&self, pre: usize, post: usize) {
+        match post.cmp(&pre) {
+            std::cmp::Ordering::Equal => {
+                // nothing to do here
+            }
+            std::cmp::Ordering::Greater => {
+                self.capacity_in_mem
+                    .fetch_add(post - pre, Ordering::Relaxed);
+            }
+            std::cmp::Ordering::Less => {
+                self.capacity_in_mem
+                    .fetch_sub(pre - post, Ordering::Relaxed);
+            }
+        }
+    }
+
     fn ms_per_age<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>>(
         &self,
         storage: &BucketMapHolder<T, U>,
@@ -203,6 +221,9 @@ impl BucketMapHolderStats {
         let startup = storage.get_startup();
         let was_startup = self.last_was_startup.swap(startup, Ordering::Relaxed);
 
+        let count_in_mem = self.count_in_mem.load(Ordering::Relaxed);
+        let capacity_in_mem = self.capacity_in_mem.load(Ordering::Relaxed);
+
         // sum of elapsed time in each thread
         let mut thread_time_elapsed_ms = elapsed_ms * storage.threads as u64;
         if storage.is_disk_index_enabled() {
@@ -234,7 +255,6 @@ impl BucketMapHolderStats {
                     ),
                 );
             }
-            let count_in_mem = self.count_in_mem.load(Ordering::Relaxed);
             let held_in_mem_ref_count = self.held_in_mem.ref_count.swap(0, Ordering::Relaxed);
             let held_in_mem_slot_list_len =
                 self.held_in_mem.slot_list_len.swap(0, Ordering::Relaxed);
@@ -263,6 +283,7 @@ impl BucketMapHolderStats {
                     i64
                 ),
                 ("count_in_mem", count_in_mem, i64),
+                ("capacity_in_mem", capacity_in_mem, i64),
                 ("count", self.total_count(), i64),
                 (
                     "bg_waiting_percent",
@@ -544,11 +565,8 @@ impl BucketMapHolderStats {
                         * InMemAccountsIndex::<T, U>::approx_size_of_one_entry(),
                     i64
                 ),
-                (
-                    "count_in_mem",
-                    self.count_in_mem.load(Ordering::Relaxed),
-                    i64
-                ),
+                ("count_in_mem", count_in_mem, i64),
+                ("capacity_in_mem", capacity_in_mem, i64),
                 ("count", self.total_count(), i64),
                 (
                     "bg_waiting_percent",