AcctIdx: stats for buckets on disk, add median (#20528)

bucket_map/src/bucket_api.rs

@@ -7,6 +7,7 @@ use solana_sdk::pubkey::Pubkey;
 use std::ops::RangeBounds;
 use std::path::PathBuf;
 
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 use std::sync::{RwLock, RwLockWriteGuard};
 
@@ -18,6 +19,7 @@ pub struct BucketApi<T: Clone + Copy> {
     pub stats: Arc<BucketMapStats>,
 
     bucket: LockedBucket<T>,
+    count: Arc<AtomicU64>,
 }
 
 impl<T: Clone + Copy> BucketApi<T> {
@@ -25,12 +27,14 @@ impl<T: Clone + Copy> BucketApi<T> {
         drives: Arc<Vec<PathBuf>>,
         max_search: MaxSearch,
         stats: Arc<BucketMapStats>,
+        count: Arc<AtomicU64>,
     ) -> Self {
         Self {
             drives,
             max_search,
             stats,
             bucket: RwLock::default(),
+            count,
         }
     }
 
@@ -90,7 +94,9 @@ impl<T: Clone + Copy> BucketApi<T> {
                 Arc::clone(&self.stats),
             ));
         } else {
-            bucket.as_mut().unwrap().handle_delayed_grows();
+            let write = bucket.as_mut().unwrap();
+            write.handle_delayed_grows();
+            self.count.store(write.bucket_len(), Ordering::Relaxed);
         }
         bucket
     }

bucket_map/src/bucket_map.rs

@@ -69,7 +69,6 @@ impl<T: Clone + Copy + Debug> BucketMap<T> {
             config.max_buckets.is_power_of_two(),
             "Max number of buckets must be a power of two"
         );
-        let stats = Arc::new(BucketMapStats::default());
         // this should be <= 1 << DEFAULT_CAPACITY or we end up searching the same items over and over - probably not a big deal since it is so small anyway
         const MAX_SEARCH: MaxSearch = 32;
         let max_search = config.max_search.unwrap_or(MAX_SEARCH);
@@ -84,14 +83,24 @@ impl<T: Clone + Copy + Debug> BucketMap<T> {
         });
         let drives = Arc::new(drives);
 
-        let mut buckets = Vec::with_capacity(config.max_buckets);
-        buckets.resize_with(config.max_buckets, || {
-            Arc::new(BucketApi::new(
-                Arc::clone(&drives),
-                max_search,
-                Arc::clone(&stats),
-            ))
+        let mut per_bucket_count = Vec::with_capacity(config.max_buckets);
+        per_bucket_count.resize_with(config.max_buckets, Arc::default);
+        let stats = Arc::new(BucketMapStats {
+            per_bucket_count,
+            ..BucketMapStats::default()
         });
+        let buckets = stats
+            .per_bucket_count
+            .iter()
+            .map(|per_bucket_count| {
+                Arc::new(BucketApi::new(
+                    Arc::clone(&drives),
+                    max_search,
+                    Arc::clone(&stats),
+                    Arc::clone(per_bucket_count),
+                ))
+            })
+            .collect();
 
         // A simple log2 function that is correct if x is a power of two
         let log2 = |x: usize| usize::BITS - x.leading_zeros() - 1;

bucket_map/src/bucket_stats.rs

@@ -11,8 +11,9 @@ pub struct BucketStats {
     pub mmap_us: AtomicU64,
 }
 
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Default)]
 pub struct BucketMapStats {
     pub index: Arc<BucketStats>,
     pub data: Arc<BucketStats>,
+    pub per_bucket_count: Vec<Arc<AtomicU64>>,
 }