From da461caf6fc177688032e3fcf157158ac1a9388b Mon Sep 17 00:00:00 2001
From: jeff washington <jeff.washington@solana.com>
Date: Mon, 10 Apr 2023 14:15:57 -0500
Subject: [PATCH] disk index: correctly keep track of count during batch add

---
 bucket_map/src/bucket.rs     | 4 ++++
 bucket_map/src/bucket_map.rs | 5 +++++
 2 files changed, 9 insertions(+)

diff --git a/bucket_map/src/bucket.rs b/bucket_map/src/bucket.rs
index 14781a9a1a72e4..959ed1bbe760dc 100644
--- a/bucket_map/src/bucket.rs
+++ b/bucket_map/src/bucket.rs
@@ -329,6 +329,10 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
                 Ok(_result) => {
                     // everything added
                     self.set_anticipated_count(0);
+                    self.index.count.fetch_add(
+                        count.saturating_sub(duplicates.len()) as u64,
+                        Ordering::Relaxed,
+                    );
                     return duplicates;
                 }
                 Err(error) => {
diff --git a/bucket_map/src/bucket_map.rs b/bucket_map/src/bucket_map.rs
index bc61d5c67ac4de..7eb9cfdd730d85 100644
--- a/bucket_map/src/bucket_map.rs
+++ b/bucket_map/src/bucket_map.rs
@@ -400,9 +400,14 @@ mod tests {
             };
 
             let verify = || {
+                let expected_count = hash_map.read().unwrap().len();
                 let mut maps = maps
                     .iter()
                     .map(|map| {
+                        let total_entries = (0..map.num_buckets())
+                            .map(|bucket| map.get_bucket_from_index(bucket).bucket_len() as usize)
+                            .sum::<usize>();
+                        assert_eq!(total_entries, expected_count);
                         let mut r = vec![];
                         for bin in 0..map.num_buckets() {
                             r.append(