From cddb9da4f0f6551f340fb635192aa250504e1d05 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)"
 <75863576+jeffwashington@users.noreply.github.com>
Date: Fri, 17 Sep 2021 15:11:07 -0500
Subject: [PATCH] AcctIdx: central age functions/state (#19980)

---
 runtime/src/bucket_map_holder.rs | 66 ++++++++++++++++++++++++++++++++
 1 file changed, 66 insertions(+)

diff --git a/runtime/src/bucket_map_holder.rs b/runtime/src/bucket_map_holder.rs
index 0b314cfdc0..2f5175e5bc 100644
--- a/runtime/src/bucket_map_holder.rs
+++ b/runtime/src/bucket_map_holder.rs
@@ -2,10 +2,14 @@ use crate::accounts_index::IndexValue;
 use crate::bucket_map_holder_stats::BucketMapHolderStats;
 use crate::waitable_condvar::WaitableCondvar;
 use std::fmt::Debug;
+use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering};
 use std::sync::Mutex;
+pub type Age = u8;
 
 // will eventually hold the bucket map
 pub struct BucketMapHolder<T: IndexValue> {
+    pub count_ages_flushed: AtomicUsize,
+    pub age: AtomicU8,
     pub stats: BucketMapHolderStats,
 
     // used by bg processing to know when any bucket has become dirty
@@ -23,8 +27,32 @@ impl<T: IndexValue> Debug for BucketMapHolder<T> {
 
 #[allow(clippy::mutex_atomic)]
 impl<T: IndexValue> BucketMapHolder<T> {
+    pub fn increment_age(&self) {
+        // fetch_add is defined to wrap.
+        // That's what we want. 0..255, then back to 0.
+        self.age.fetch_add(1, Ordering::Relaxed);
+        // since we changed age, there are now 0 buckets that have been flushed at this age
+        let previous = self.count_ages_flushed.swap(0, Ordering::Relaxed);
+        assert!(previous >= self.bins); // we should not have increased age before previous age was fully flushed
+    }
+
+    pub fn current_age(&self) -> Age {
+        self.age.load(Ordering::Relaxed)
+    }
+
+    pub fn bucket_flushed_at_current_age(&self) {
+        self.count_ages_flushed.fetch_add(1, Ordering::Relaxed);
+    }
+
+    // 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
+    }
+
     pub fn new(bins: usize) -> Self {
         Self {
+            count_ages_flushed: AtomicUsize::default(),
+            age: AtomicU8::default(),
             stats: BucketMapHolderStats::default(),
             wait_dirty_bucket: WaitableCondvar::default(),
             next_bucket_to_flush: Mutex::new(0),
@@ -74,4 +102,42 @@ pub mod tests {
             assert_eq!(visited.load(Ordering::Relaxed), expected, "bin: {}", bin)
         });
     }
+
+    #[test]
+    fn test_age_increment() {
+        solana_logger::setup();
+        let bins = 4;
+        let test = BucketMapHolder::<u64>::new(bins);
+        for age in 0..513 {
+            assert_eq!(test.current_age(), (age % 256) as Age);
+
+            // inc all
+            for _ in 0..bins {
+                assert!(!test.all_buckets_flushed_at_current_age());
+                test.bucket_flushed_at_current_age();
+            }
+
+            test.increment_age();
+        }
+    }
+
+    #[test]
+    fn test_age_broad() {
+        solana_logger::setup();
+        let bins = 4;
+        let test = BucketMapHolder::<u64>::new(bins);
+        assert_eq!(test.current_age(), 0);
+        assert!(!test.all_buckets_flushed_at_current_age());
+        // inc all but 1
+        for _ in 1..bins {
+            test.bucket_flushed_at_current_age();
+            assert!(!test.all_buckets_flushed_at_current_age());
+        }
+        test.bucket_flushed_at_current_age();
+        assert!(test.all_buckets_flushed_at_current_age());
+        test.increment_age();
+
+        assert_eq!(test.current_age(), 1);
+        assert!(!test.all_buckets_flushed_at_current_age());
+    }
 }