From 640663e1bc886478f0e5e332e0a4e1704951d39c Mon Sep 17 00:00:00 2001
From: Christopher Parratto <cparratto@apple.com>
Date: Wed, 14 Feb 2024 11:10:23 -0600
Subject: [PATCH] Adds prometheus builder support for specifying bucket count

---
 metrics-exporter-prometheus/src/builder.rs    | 19 +++++++++++++++++
 metrics-exporter-prometheus/src/common.rs     |  4 ++++
 .../src/distribution.rs                       | 21 ++++++++++++-------
 3 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/metrics-exporter-prometheus/src/builder.rs b/metrics-exporter-prometheus/src/builder.rs
index ef02b8c5..0510e47f 100644
--- a/metrics-exporter-prometheus/src/builder.rs
+++ b/metrics-exporter-prometheus/src/builder.rs
@@ -97,6 +97,7 @@ pub struct PrometheusBuilder {
     allowed_addresses: Option<Vec<IpNet>>,
     quantiles: Vec<Quantile>,
     bucket_duration: Option<Duration>,
+    bucket_count: Option<u32>,
     buckets: Option<Vec<f64>>,
     bucket_overrides: Option<HashMap<Matcher, Vec<f64>>>,
     idle_timeout: Option<Duration>,
@@ -122,6 +123,7 @@ impl PrometheusBuilder {
             allowed_addresses: None,
             quantiles,
             bucket_duration: None,
+            bucket_count: None,
             buckets: None,
             bucket_overrides: None,
             idle_timeout: None,
@@ -257,6 +259,22 @@ impl PrometheusBuilder {
         Ok(self)
     }
 
+    /// Sets the default bucket count for rolling summaries
+    ///
+    /// Count number buckets are created to store summary information
+    ///
+    /// ## Errors
+    ///
+    /// If `value` less than 1 error will be thrown
+    pub fn set_bucket_count(mut self, count: u32) -> Result<Self, BuildError> {
+        if count == 0 {
+            return Err(BuildError::ZeroBucketCount);
+        }
+
+        self.bucket_count = Some(count);
+        Ok(self)
+    }
+
     /// Sets the buckets to use when rendering histograms.
     ///
     /// Buckets values represent the higher bound of each buckets.  If buckets are set, then all
@@ -556,6 +574,7 @@ impl PrometheusBuilder {
                 self.quantiles,
                 self.bucket_duration,
                 self.buckets,
+                self.bucket_count,
                 self.bucket_overrides,
             ),
             descriptions: RwLock::new(HashMap::new()),
diff --git a/metrics-exporter-prometheus/src/common.rs b/metrics-exporter-prometheus/src/common.rs
index 94ff6ab6..23e0a814 100644
--- a/metrics-exporter-prometheus/src/common.rs
+++ b/metrics-exporter-prometheus/src/common.rs
@@ -78,6 +78,10 @@ pub enum BuildError {
     /// Bucket duration cannot be zero
     #[error("bucket durations cannot be set to zero")]
     ZeroBucketDuration,
+
+    /// Bucket count cannot be zero
+    #[error("bucket count cannot be set to zero")]
+    ZeroBucketCount,
 }
 
 pub struct Snapshot {
diff --git a/metrics-exporter-prometheus/src/distribution.rs b/metrics-exporter-prometheus/src/distribution.rs
index 476014ce..1cd0c4ec 100644
--- a/metrics-exporter-prometheus/src/distribution.rs
+++ b/metrics-exporter-prometheus/src/distribution.rs
@@ -36,13 +36,14 @@ impl Distribution {
     /// Creates a summary distribution.
     pub fn new_summary(
         quantiles: Arc<Vec<Quantile>>,
-        bucket_duration: Option<Duration>,
+        bucket_duration: Duration,
+        bucket_count: u32,
     ) -> Distribution {
-        let summary = bucket_duration.map_or(
-            RollingSummary::new(NonZeroU32::new(3).unwrap(), Duration::from_secs(20)),
-            |duration| RollingSummary::new(NonZeroU32::new(3).unwrap(), duration),
-        );
-        Distribution::Summary(summary, quantiles, 0.0)
+        Distribution::Summary(
+            RollingSummary::new(NonZeroU32::new(bucket_count).unwrap(), bucket_duration),
+            quantiles,
+            0.0,
+        )
     }
 
     /// Records the given `samples` in the current distribution.
@@ -67,6 +68,7 @@ pub struct DistributionBuilder {
     quantiles: Arc<Vec<Quantile>>,
     buckets: Option<Vec<f64>>,
     bucket_duration: Option<Duration>,
+    bucket_count: Option<u32>,
     bucket_overrides: Option<Vec<(Matcher, Vec<f64>)>>,
 }
 
@@ -76,12 +78,14 @@ impl DistributionBuilder {
         quantiles: Vec<Quantile>,
         bucket_duration: Option<Duration>,
         buckets: Option<Vec<f64>>,
+        bucket_count: Option<u32>,
         bucket_overrides: Option<HashMap<Matcher, Vec<f64>>>,
     ) -> DistributionBuilder {
         DistributionBuilder {
             quantiles: Arc::new(quantiles),
             bucket_duration,
             buckets,
+            bucket_count,
             bucket_overrides: bucket_overrides.map(|entries| {
                 let mut matchers = entries.into_iter().collect::<Vec<_>>();
                 matchers.sort_by(|a, b| a.0.cmp(&b.0));
@@ -104,7 +108,10 @@ impl DistributionBuilder {
             return Distribution::new_histogram(buckets);
         }
 
-        Distribution::new_summary(self.quantiles.clone(), self.bucket_duration)
+        let b_duration = self.bucket_duration.map_or(Duration::from_secs(20), |d| d);
+        let b_count = self.bucket_count.map_or(3, |c| c);
+
+        Distribution::new_summary(self.quantiles.clone(), b_duration, b_count)
     }
 
     /// Returns the distribution type for the given metric key.