From 2afded392b0ecace130d0be1862d291d3d117751 Mon Sep 17 00:00:00 2001
From: jeff washington <jeff.washington@solana.com>
Date: Mon, 3 Apr 2023 14:34:40 -0500
Subject: [PATCH] disk index: set_anticipated_count to optimally grow disk
 buckets at startup

---
 bucket_map/src/bucket.rs             | 19 ++++++++++++++++++-
 bucket_map/src/bucket_api.rs         | 13 +++++++++++++
 runtime/src/in_mem_accounts_index.rs |  5 +++++
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/bucket_map/src/bucket.rs b/bucket_map/src/bucket.rs
index bcca7b605c588b..b92231c6da80de 100644
--- a/bucket_map/src/bucket.rs
+++ b/bucket_map/src/bucket.rs
@@ -96,6 +96,10 @@ pub struct Bucket<T: Copy + 'static> {
     pub data: Vec<BucketStorage<DataBucket>>,
     stats: Arc<BucketMapStats>,
 
+    /// # entries caller expects the map to need to contain.
+    /// Used as a hint for the next time we need to grow.
+    anticipated_size: AtomicU64,
+
     pub reallocated: Reallocated<IndexBucket<T>, DataBucket>,
 }
 
@@ -123,6 +127,7 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
             data: vec![],
             stats,
             reallocated: Reallocated::default(),
+            anticipated_size: AtomicU64::default(),
         }
     }
 
@@ -420,8 +425,20 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
         }
     }
 
+    pub(crate) fn set_anticipated_count(&self, count: u64) {
+        self.anticipated_size.store(count, Ordering::Release);
+    }
+
     pub fn grow_index(&self, current_capacity_pow2: u8) {
         if self.index.capacity_pow2 == current_capacity_pow2 {
+            // make sure to grow to at least % more than the anticpated size
+            let mut starting_size_pow2 = self.index.capacity_pow2;
+            let anticipated_size = self.anticipated_size.load(Ordering::Acquire);
+            if anticipated_size > 0 {
+                // start the growth at the next pow2 larger than what would be required to hold `anticipated_size`.
+                // This will prevent unnecessary repeated grows at startup.
+                starting_size_pow2 = starting_size_pow2.max(anticipated_size.ilog2() as u8);
+            }
             let mut m = Measure::start("grow_index");
             //debug!("GROW_INDEX: {}", current_capacity_pow2);
             let increment = 1;
@@ -434,7 +451,7 @@ impl<'b, T: Clone + Copy + 'static> Bucket<T> {
                     1,
                     std::mem::size_of::<IndexEntry<T>>() as u64,
                     // *2 causes rapid growth of index buckets
-                    self.index.capacity_pow2 + i, // * 2,
+                    starting_size_pow2 + i, // * 2,
                     self.index.max_search,
                     Arc::clone(&self.stats.index),
                     Arc::clone(&self.index.count),
diff --git a/bucket_map/src/bucket_api.rs b/bucket_map/src/bucket_api.rs
index c4be1cd24e3d8b..6b36e2cfa5b54c 100644
--- a/bucket_map/src/bucket_api.rs
+++ b/bucket_map/src/bucket_api.rs
@@ -111,6 +111,19 @@ impl<T: Clone + Copy> BucketApi<T> {
         }
     }
 
+    /// caller can specify that the index needs to hold approximately `count` entries soon.
+    /// This gives a hint to the resizing algorithm and prevents repeated incremental resizes.
+    pub fn set_anticipated_count(&self, count: u64) {
+        if let Some(bucket) = self.bucket.read().unwrap().as_ref() {
+            bucket.set_anticipated_count(count);
+            return;
+        }
+        // caller is about to write a lot of elements to this bucket, so we need to allocate it first
+        // so we have somewhere to store the anticipated count
+        let mut bucket = self.get_write_bucket();
+        bucket.as_mut().unwrap().set_anticipated_count(count);
+    }
+
     pub fn update<F>(&self, key: &Pubkey, updatefn: F)
     where
         F: FnMut(Option<(&[T], RefCount)>) -> Option<(Vec<T>, RefCount)>,
diff --git a/runtime/src/in_mem_accounts_index.rs b/runtime/src/in_mem_accounts_index.rs
index 5ebdf760227592..8217f76489791d 100644
--- a/runtime/src/in_mem_accounts_index.rs
+++ b/runtime/src/in_mem_accounts_index.rs
@@ -1061,6 +1061,9 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
         // merge all items into the disk index now
         let disk = self.bucket.as_ref().unwrap();
         let mut count = 0;
+        let current_len = disk.bucket_len();
+        let anticipated = insert.len();
+        disk.set_anticipated_count((anticipated as u64).saturating_add(current_len));
         insert.into_iter().for_each(|(slot, k, v)| {
             let entry = (slot, v);
             let new_ref_count = u64::from(!v.is_cached());
@@ -1085,6 +1088,8 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> InMemAccountsIndex<T,
                 }
             });
         });
+        // remove the guidance for how many entries the bucket will eventually contain since we have added all we knew about
+        disk.set_anticipated_count(0);
         self.stats().inc_insert_count(count);
     }