From 4a34b77af5b6ae8addb8f60cfbda5dd2686b61d3 Mon Sep 17 00:00:00 2001
From: Russell Cohen <rcoh@amazon.com>
Date: Thu, 7 Nov 2024 02:45:36 -0500
Subject: [PATCH] metrics: fix bug with wrong number of buckets for the
 histogram (#6957)

---
 tokio/src/runtime/metrics/histogram.rs | 14 ++++++++++++--
 tokio/tests/rt_unstable_metrics.rs     | 17 ++++++++++++++++-
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/tokio/src/runtime/metrics/histogram.rs b/tokio/src/runtime/metrics/histogram.rs
index 5e612492418..6581ef10afc 100644
--- a/tokio/src/runtime/metrics/histogram.rs
+++ b/tokio/src/runtime/metrics/histogram.rs
@@ -264,14 +264,14 @@ impl HistogramBuilder {
             }
             None => self.histogram_type,
         };
-        let num_buckets = self.histogram_type.num_buckets();
+        let num_buckets = histogram_type.num_buckets();
 
         Histogram {
             buckets: (0..num_buckets)
                 .map(|_| MetricAtomicU64::new(0))
                 .collect::<Vec<_>>()
                 .into_boxed_slice(),
-            histogram_type: histogram_type,
+            histogram_type,
         }
     }
 }
@@ -303,6 +303,13 @@ mod test {
         .build()
     }
 
+    #[test]
+    fn test_legacy_builder() {
+        let mut builder = HistogramBuilder::new();
+        builder.legacy_mut(|b| b.num_buckets = 20);
+        assert_eq!(builder.build().num_buckets(), 20);
+    }
+
     #[test]
     fn log_scale_resolution_1() {
         let h = HistogramBuilder {
@@ -355,6 +362,9 @@ mod test {
 
         b.measure(4096, 1);
         assert_bucket_eq!(b, 9, 1);
+
+        b.measure(u64::MAX, 1);
+        assert_bucket_eq!(b, 9, 2);
     }
 
     #[test]
diff --git a/tokio/tests/rt_unstable_metrics.rs b/tokio/tests/rt_unstable_metrics.rs
index 6e3e97f1df4..df05bfbf7ce 100644
--- a/tokio/tests/rt_unstable_metrics.rs
+++ b/tokio/tests/rt_unstable_metrics.rs
@@ -13,7 +13,7 @@ use std::task::Poll;
 use std::thread;
 use tokio::macros::support::poll_fn;
 
-use tokio::runtime::{HistogramConfiguration, LogHistogram, Runtime};
+use tokio::runtime::{HistogramConfiguration, HistogramScale, LogHistogram, Runtime};
 use tokio::task::consume_budget;
 use tokio::time::{self, Duration};
 
@@ -424,6 +424,21 @@ fn log_histogram() {
     assert_eq!(N, n);
 }
 
+#[test]
+#[allow(deprecated)]
+fn legacy_log_histogram() {
+    let rt = tokio::runtime::Builder::new_multi_thread()
+        .enable_all()
+        .enable_metrics_poll_time_histogram()
+        .metrics_poll_count_histogram_scale(HistogramScale::Log)
+        .metrics_poll_count_histogram_resolution(Duration::from_micros(50))
+        .metrics_poll_count_histogram_buckets(20)
+        .build()
+        .unwrap();
+    let num_buckets = rt.metrics().poll_time_histogram_num_buckets();
+    assert_eq!(num_buckets, 20);
+}
+
 #[test]
 fn log_histogram_default_configuration() {
     let rt = tokio::runtime::Builder::new_current_thread()