From 2d3b601fae8d0de18ecd469dccec6049b808d6dc Mon Sep 17 00:00:00 2001
From: Arun Mahendra <arun.mahendra@shopify.com>
Date: Thu, 21 Jul 2022 09:14:41 -0400
Subject: [PATCH] explicitly adding +inf bucket to withExemplarsMetric

Signed-off-by: Arun Mahendra <arun.mahendra@shopify.com>
---
 prometheus/metric.go      | 14 +++++++++++---
 prometheus/metric_test.go | 20 ++++++++++++++++++--
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/prometheus/metric.go b/prometheus/metric.go
index 48d4a5d50..00bbd18e1 100644
--- a/prometheus/metric.go
+++ b/prometheus/metric.go
@@ -15,6 +15,7 @@ package prometheus
 
 import (
 	"errors"
+	"math"
 	"sort"
 	"strings"
 	"time"
@@ -183,9 +184,16 @@ func (m *withExemplarsMetric) Write(pb *dto.Metric) error {
 			})
 			if i < len(pb.Histogram.Bucket) {
 				pb.Histogram.Bucket[i].Exemplar = e
-			} else {
-				// This is not possible as last bucket is Inf.
-				panic("no bucket was found for given exemplar value")
+			} else { // +inf bucket should be explicitly added if there is an exemplar for it.
+				b := &dto.Bucket{
+					CumulativeCount: proto.Uint64(pb.Histogram.Bucket[len(pb.Histogram.GetBucket())-1].GetCumulativeCount()),
+					UpperBound:      proto.Float64(math.Inf(1)),
+					Exemplar:        e,
+				}
+				pb.Histogram.Bucket = append(pb.Histogram.Bucket, b)
+				break
+				// end looping after creating +inf bucket and adding one exemplar.
+				// there could be other exemplars that are in the "inf" range but those will be ignored.
 			}
 		}
 	default:
diff --git a/prometheus/metric_test.go b/prometheus/metric_test.go
index 61d807cdc..fb796d2e6 100644
--- a/prometheus/metric_test.go
+++ b/prometheus/metric_test.go
@@ -14,6 +14,7 @@
 package prometheus
 
 import (
+	"math"
 	"testing"
 
 	//nolint:staticcheck // Ignore SA1019. Need to keep deprecated package for compatibility.
@@ -56,16 +57,19 @@ func TestWithExemplarsMetric(t *testing.T) {
 			{Value: proto.Float64(89.0)},
 			{Value: proto.Float64(100.0)},
 			{Value: proto.Float64(157.0)},
+			{Value: proto.Float64(500.0)},
+			{Value: proto.Float64(2000.0)},
 		}}
 		metric := dto.Metric{}
 		if err := m.Write(&metric); err != nil {
 			t.Fatal(err)
 		}
-		if want, got := 4, len(metric.GetHistogram().Bucket); want != got {
+		if want, got := 5, len(metric.GetHistogram().Bucket); want != got {
 			t.Errorf("want %v, got %v", want, got)
 		}
 
-		expectedExemplarVals := []float64{24.0, 42.0, 100.0, 157.0}
+		// when there are more exemplars than there are buckets, a +inf bucket will be created and the last exemplar value will be added to the +inf bucket.
+		expectedExemplarVals := []float64{24.0, 42.0, 100.0, 157.0, 500.0}
 		for i, b := range metric.GetHistogram().Bucket {
 			if b.Exemplar == nil {
 				t.Errorf("Expected exemplar for bucket %v, got nil", i)
@@ -74,5 +78,17 @@ func TestWithExemplarsMetric(t *testing.T) {
 				t.Errorf("%v: want %v, got %v", i, want, got)
 			}
 		}
+
+		infBucket := metric.GetHistogram().Bucket[len(metric.GetHistogram().Bucket)-1].GetUpperBound()
+
+		if infBucket != math.Inf(1) {
+			t.Errorf("want %v, got %v", math.Inf(1), infBucket)
+		}
+
+		infBucketValue := metric.GetHistogram().Bucket[len(metric.GetHistogram().Bucket)-1].GetExemplar().GetValue()
+
+		if infBucketValue != 500.0 {
+			t.Errorf("want %v, got %v", 500.0, infBucketValue)
+		}
 	})
 }