From aa3aa5cc9c9305c569bf33398bd2e0cce20be790 Mon Sep 17 00:00:00 2001
From: HaoranYi <haoran.yi@gmail.com>
Date: Sat, 7 Jan 2023 09:25:04 -0600
Subject: [PATCH] extract max_size stat update for disk buckets (#29487)

---
 bucket_map/src/bucket.rs         |  6 +-----
 bucket_map/src/bucket_stats.rs   |  7 +++++++
 bucket_map/src/bucket_storage.rs | 10 +++++-----
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/bucket_map/src/bucket.rs b/bucket_map/src/bucket.rs
index 51d3c1dee..7de66d598 100644
--- a/bucket_map/src/bucket.rs
+++ b/bucket_map/src/bucket.rs
@@ -390,11 +390,7 @@ impl<T: Clone + Copy> Bucket<T> {
                     }
                 }
                 if valid {
-                    let sz = index.capacity();
-                    {
-                        let mut max = self.stats.index.max_size.lock().unwrap();
-                        *max = std::cmp::max(*max, sz);
-                    }
+                    self.stats.index.update_max_size(index.capacity());
                     let mut items = self.reallocated.items.lock().unwrap();
                     items.index = Some((random, index));
                     self.reallocated.add_reallocation();
diff --git a/bucket_map/src/bucket_stats.rs b/bucket_map/src/bucket_stats.rs
index 8521404fc..a01b3d7c7 100644
--- a/bucket_map/src/bucket_stats.rs
+++ b/bucket_map/src/bucket_stats.rs
@@ -11,6 +11,13 @@ pub struct BucketStats {
     pub find_entry_mut_us: AtomicU64,
 }
 
+impl BucketStats {
+    pub fn update_max_size(&self, size: u64) {
+        let mut max = self.max_size.lock().unwrap();
+        *max = std::cmp::max(*max, size);
+    }
+}
+
 #[derive(Debug, Default)]
 pub struct BucketMapStats {
     pub index: Arc<BucketStats>,
diff --git a/bucket_map/src/bucket_storage.rs b/bucket_map/src/bucket_storage.rs
index 373610115..69e3eecd2 100644
--- a/bucket_map/src/bucket_storage.rs
+++ b/bucket_map/src/bucket_storage.rs
@@ -341,6 +341,10 @@ impl BucketStorage {
         self.stats.resize_us.fetch_add(m.as_us(), Ordering::Relaxed);
     }
 
+    pub fn update_max_size(&self) {
+        self.stats.update_max_size(self.capacity());
+    }
+
     /// allocate a new bucket, copying data from 'bucket'
     pub fn new_resized(
         drives: &Arc<Vec<PathBuf>>,
@@ -365,11 +369,7 @@ impl BucketStorage {
         if let Some(bucket) = bucket {
             new_bucket.copy_contents(bucket);
         }
-        let sz = new_bucket.capacity();
-        {
-            let mut max = new_bucket.stats.max_size.lock().unwrap();
-            *max = std::cmp::max(*max, sz);
-        }
+        new_bucket.update_max_size();
         new_bucket
     }