From c3679ab9bde78ea2956a2996d169753c52fb7c10 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)"
 <75863576+jeffwashington@users.noreply.github.com>
Date: Tue, 21 Sep 2021 08:40:55 -0500
Subject: [PATCH] AcctIdx: first pass at aging buckets (#20008)

---
 runtime/src/bucket_map_holder.rs     |  8 ++++--
 runtime/src/in_mem_accounts_index.rs | 41 ++++++++++++++++++++++------
 2 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/runtime/src/bucket_map_holder.rs b/runtime/src/bucket_map_holder.rs
index f8208cd17..05367511f 100644
--- a/runtime/src/bucket_map_holder.rs
+++ b/runtime/src/bucket_map_holder.rs
@@ -88,12 +88,16 @@ impl<T: IndexValue> BucketMapHolder<T> {
     }
 
     pub fn bucket_flushed_at_current_age(&self) {
-        self.count_ages_flushed.fetch_add(1, Ordering::Relaxed);
+        self.count_ages_flushed.fetch_add(1, Ordering::Acquire);
     }
 
     // have all buckets been flushed at the current age?
     pub fn all_buckets_flushed_at_current_age(&self) -> bool {
-        self.count_ages_flushed.load(Ordering::Relaxed) >= self.bins
+        self.count_ages_flushed() >= self.bins
+    }
+
+    pub fn count_ages_flushed(&self) -> usize {
+        self.count_ages_flushed.load(Ordering::Relaxed)
     }
 
     pub fn maybe_advance_age(&self) -> bool {
diff --git a/runtime/src/in_mem_accounts_index.rs b/runtime/src/in_mem_accounts_index.rs
index 945da7ad6..a65a7d8b9 100644
--- a/runtime/src/in_mem_accounts_index.rs
+++ b/runtime/src/in_mem_accounts_index.rs
@@ -59,15 +59,19 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
 
     /// true if this bucket needs to call flush for the current age
     /// we need to scan each bucket once per value of age
-    fn get_should_age(&self) -> bool {
-        let last_age_flushed = self.last_age_flushed.load(Ordering::Relaxed);
-        let age = self.storage.age.load(Ordering::Relaxed);
-        last_age_flushed == age
+    fn get_should_age(&self, age: Age) -> bool {
+        let last_age_flushed = self.last_age_flushed();
+        last_age_flushed != age
     }
 
     /// called after flush scans this bucket at the current age
     fn set_has_aged(&self, age: Age) {
         self.last_age_flushed.store(age, Ordering::Relaxed);
+        self.storage.bucket_flushed_at_current_age();
+    }
+
+    fn last_age_flushed(&self) -> Age {
+        self.last_age_flushed.load(Ordering::Relaxed)
     }
 
     fn map(&self) -> &RwLock<HashMap<Pubkey, AccountMapEntry<T>>> {
@@ -411,8 +415,10 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
 
     fn flush_internal(&self) {
         let was_dirty = self.bin_dirty.swap(false, Ordering::Acquire);
-        if !was_dirty {
-            // wasn't dirty, no need to flush
+        let current_age = self.storage.current_age();
+        let iterate_for_age = self.get_should_age(current_age);
+        if !was_dirty && !iterate_for_age {
+            // wasn't dirty and no need to age, so no need to flush this bucket
             return;
         }
 
@@ -431,6 +437,11 @@ impl<T: IndexValue> InMemAccountsIndex<T> {
             }
             */
         }
+        if iterate_for_age {
+            // completed iteration of the buckets at the current age
+            assert_eq!(current_age, self.storage.current_age());
+            self.set_has_aged(current_age);
+        }
     }
 
     fn stats(&self) -> &BucketMapHolderStats {
@@ -518,8 +529,22 @@ mod tests {
     fn test_age() {
         solana_logger::setup();
         let test = new_for_test::<u64>();
-        assert!(!test.get_should_age());
+        assert!(test.get_should_age(test.storage.current_age()));
+        assert_eq!(test.storage.count_ages_flushed(), 0);
         test.set_has_aged(0);
-        assert!(test.get_should_age());
+        assert!(!test.get_should_age(test.storage.current_age()));
+        assert_eq!(test.storage.count_ages_flushed(), 1);
+        // simulate rest of buckets aging
+        for _ in 1..BINS_FOR_TESTING {
+            assert!(!test.storage.all_buckets_flushed_at_current_age());
+            test.storage.bucket_flushed_at_current_age();
+        }
+        assert!(test.storage.all_buckets_flushed_at_current_age());
+        // advance age
+        test.storage.increment_age();
+        assert_eq!(test.storage.current_age(), 1);
+        assert!(!test.storage.all_buckets_flushed_at_current_age());
+        assert!(test.get_should_age(test.storage.current_age()));
+        assert_eq!(test.storage.count_ages_flushed(), 0);
     }
 }