From 5bbb0da7b80448655059465704beb9fee5183ab8 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)"
 <75863576+jeffwashington@users.noreply.github.com>
Date: Thu, 23 Sep 2021 08:14:53 -0500
Subject: [PATCH] AcctIdx: stats for in-mem (#20114)

---
 runtime/src/bucket_map_holder_stats.rs | 16 +++++++++++++---
 runtime/src/in_mem_accounts_index.rs   | 10 +++++++++-
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/runtime/src/bucket_map_holder_stats.rs b/runtime/src/bucket_map_holder_stats.rs
index 33fa2b088d..05342b5715 100644
--- a/runtime/src/bucket_map_holder_stats.rs
+++ b/runtime/src/bucket_map_holder_stats.rs
@@ -23,6 +23,7 @@ pub struct BucketMapHolderStats {
     pub keys: AtomicU64,
     pub deletes: AtomicU64,
     pub inserts: AtomicU64,
+    pub count: AtomicU64,
     pub count_in_mem: AtomicU64,
     pub per_bucket_count: Vec<AtomicU64>,
     pub flush_entries_updated_on_disk: AtomicU64,
@@ -47,14 +48,22 @@ impl BucketMapHolderStats {
         }
     }
 
-    pub fn insert_or_delete(&self, insert: bool, bin: usize) {
-        let per_bucket = self.per_bucket_count.get(bin);
+    pub fn insert_or_delete(&self, insert: bool, _bin: usize) {
         if insert {
             self.inserts.fetch_add(1, Ordering::Relaxed);
+            self.count.fetch_add(1, Ordering::Relaxed);
+        } else {
+            self.deletes.fetch_add(1, Ordering::Relaxed);
+            self.count.fetch_sub(1, Ordering::Relaxed);
+        }
+    }
+
+    pub fn insert_or_delete_mem(&self, insert: bool, bin: usize) {
+        let per_bucket = self.per_bucket_count.get(bin);
+        if insert {
             self.count_in_mem.fetch_add(1, Ordering::Relaxed);
             per_bucket.map(|count| count.fetch_add(1, Ordering::Relaxed));
         } else {
-            self.deletes.fetch_add(1, Ordering::Relaxed);
             self.count_in_mem.fetch_sub(1, Ordering::Relaxed);
             per_bucket.map(|count| count.fetch_sub(1, Ordering::Relaxed));
         }
@@ -109,6 +118,7 @@ impl BucketMapHolderStats {
                 self.count_in_mem.load(Ordering::Relaxed),
                 i64
             ),
+            ("count", self.count.load(Ordering::Relaxed), i64),
             ("min_in_bin", min, i64),
             ("max_in_bin", max, i64),
             ("count_from_bins", ct, i64),
diff --git a/runtime/src/in_mem_accounts_index.rs b/runtime/src/in_mem_accounts_index.rs
index 70116aff71..925344452c 100644
--- a/runtime/src/in_mem_accounts_index.rs
+++ b/runtime/src/in_mem_accounts_index.rs
@@ -152,7 +152,10 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
         let entry = map.entry(*pubkey);
         let result = match entry {
             Entry::Occupied(occupied) => Arc::clone(occupied.get()),
-            Entry::Vacant(vacant) => Arc::clone(vacant.insert(new_entry)),
+            Entry::Vacant(vacant) => {
+                stats.insert_or_delete_mem(true, self.bin);
+                Arc::clone(vacant.insert(new_entry))
+            }
         };
         Some(result)
     }
@@ -184,6 +187,7 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
                     // We have to have a write lock to the map here, which means nobody else can get
                     //  the arc, but someone may already have retreived a clone of it.
                     self.delete_disk_key(occupied.key());
+                    self.stats().insert_or_delete_mem(false, self.bin);
                     occupied.remove();
                 }
                 result
@@ -274,6 +278,7 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
                 };
                 assert!(new_value.dirty());
                 vacant.insert(new_value);
+                self.stats().insert_or_delete_mem(true, self.bin);
                 self.stats().insert_or_delete(true, self.bin);
             }
         }
@@ -424,6 +429,7 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
             Entry::Vacant(vacant) => {
                 // not in cache, look on disk
                 let disk_entry = self.load_account_entry_from_disk(vacant.key());
+                stats.insert_or_delete_mem(true, self.bin);
                 if let Some(disk_entry) = disk_entry {
                     // on disk, so insert into cache, then return cache value so caller will merge
                     let result = Some(Self::insert_returner(&disk_entry, vacant.key(), new_entry));
@@ -537,6 +543,7 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
                     }
                     Entry::Vacant(vacant) => {
                         vacant.insert(self.disk_to_cache_entry(item.slot_list, item.ref_count));
+                        self.stats().insert_or_delete_mem(true, self.bin);
                     }
                 }
             }
@@ -682,6 +689,7 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
                 }
 
                 // all conditions for removing succeeded, so really remove item from in-mem cache
+                self.stats().insert_or_delete_mem(false, self.bin);
                 occupied.remove();
             }
         }