From 839018719cc949c2b657a3310f7654f7feaf95d9 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Sat, 23 Nov 2019 23:58:18 +0200
Subject: [PATCH] Add Histogram.VisitNonZeroBuckets

---
 histogram.go      | 62 +++++++++++++++++++++++++++++++++++++++++------
 histogram_test.go | 28 +++++++++++++++------
 metrics_test.go   |  1 +
 3 files changed, 76 insertions(+), 15 deletions(-)

diff --git a/histogram.go b/histogram.go
index a72e6d6..0ee4400 100644
--- a/histogram.go
+++ b/histogram.go
@@ -91,6 +91,10 @@ func GetOrCreateHistogram(name string) *Histogram {
 //
 // v cannot be negative.
 func (h *Histogram) Update(v float64) {
+	if math.IsNaN(v) || v < 0 {
+		// Skip NaNs and negative values.
+		return
+	}
 	idx := getBucketIdx(v)
 	if idx >= uint(len(h.buckets)) {
 		panic(fmt.Errorf("BUG: idx cannot exceed %d; got %d", len(h.buckets), idx))
@@ -108,13 +112,60 @@ func (h *Histogram) UpdateDuration(startTime time.Time) {
 	h.Update(d)
 }
 
+// VisitNonZeroBuckets calls f for all buckets with non-zero counters.
+func (h *Histogram) VisitNonZeroBuckets(f func(vmrange string, count uint64)) {
+	for i, v := range h.buckets[:] {
+		if v == 0 {
+			continue
+		}
+		vmrange := getRangeForBucketIdx(uint(i))
+		f(vmrange, v)
+	}
+}
+
+func getRangeForBucketIdx(idx uint) string {
+	bucketRangesOnce.Do(initBucketRanges)
+	return bucketRanges[idx]
+}
+
+func initBucketRanges() {
+	start := "0"
+	for i := 0; i < bucketsCount; i++ {
+		end := getRangeEndFromBucketIdx(uint(i))
+		bucketRanges[i] = start + "..." + end
+		start = end
+	}
+}
+
+var (
+	bucketRanges     [bucketsCount]string
+	bucketRangesOnce sync.Once
+)
+
+func getTagForBucketIdx(idx uint) string {
+	bucketTagsOnce.Do(initBucketTags)
+	return bucketTags[idx]
+}
+
+func initBucketTags() {
+	for i := 0; i < bucketsCount; i++ {
+		vmrange := getRangeForBucketIdx(uint(i))
+		bucketTags[i] = fmt.Sprintf(`vmrange=%q`, vmrange)
+	}
+}
+
+var (
+	bucketTags     [bucketsCount]string
+	bucketTagsOnce sync.Once
+)
+
 func (h *Histogram) marshalTo(prefix string, w io.Writer) {
 	count := atomic.LoadUint64(&h.count)
 	if count == 0 {
 		return
 	}
 	for i := range h.buckets[:] {
-		h.marshalBucket(prefix, w, i)
+		h.marshalBucket(prefix, w, uint(i))
 	}
 	// Marshal `_sum` and `_count` metrics.
 	name, filters := splitMetricName(prefix)
@@ -129,17 +180,12 @@ func (h *Histogram) marshalTo(prefix string, w io.Writer) {
 	fmt.Fprintf(w, "%s_count%s %d\n", name, filters, count)
 }
 
-func (h *Histogram) marshalBucket(prefix string, w io.Writer, idx int) {
+func (h *Histogram) marshalBucket(prefix string, w io.Writer, idx uint) {
 	v := h.buckets[idx]
 	if v == 0 {
 		return
 	}
-	start := "0"
-	if idx > 0 {
-		start = getRangeEndFromBucketIdx(uint(idx - 1))
-	}
-	end := getRangeEndFromBucketIdx(uint(idx))
-	tag := fmt.Sprintf(`vmrange="%s...%s"`, start, end)
+	tag := getTagForBucketIdx(idx)
 	prefix = addTag(prefix, tag)
 	name, filters := splitMetricName(prefix)
 	fmt.Fprintf(w, "%s_bucket%s %d\n", name, filters, v)
diff --git a/histogram_test.go b/histogram_test.go
index 0556b5d..49264fb 100644
--- a/histogram_test.go
+++ b/histogram_test.go
@@ -4,18 +4,12 @@ import (
 	"bytes"
 	"fmt"
 	"math"
+	"reflect"
 	"strings"
 	"testing"
 	"time"
 )
 
-func TestHistogramUpdateNegativeValue(t *testing.T) {
-	h := NewHistogram("TestHisogramUpdateNegativeValue")
-	expectPanic(t, "negative value", func() {
-		h.Update(-123)
-	})
-}
-
 func TestGetBucketIdx(t *testing.T) {
 	f := func(v float64, idxExpected uint) {
 		t.Helper()
@@ -126,6 +120,11 @@ func TestHistogramSerial(t *testing.T) {
 	}
 	h.UpdateDuration(time.Now().Add(-time.Minute))
 
+	// Verify edge cases
+	h.Update(math.Inf(1))
+	h.Update(math.NaN())
+	h.Update(-123)
+
 	// Make sure the histogram becomes visible in the output of WritePrometheus,
 	// since now it contains values.
 	bb.Reset()
@@ -149,6 +148,21 @@ func TestHistogramConcurrent(t *testing.T) {
 		t.Fatal(err)
 	}
 	testMarshalTo(t, h, "prefix", "prefix_bucket{vmrange=\"0...0\"} 5\nprefix_bucket{vmrange=\"9e-1...1\"} 5\nprefix_bucket{vmrange=\"1...2\"} 5\nprefix_bucket{vmrange=\"2...3\"} 5\nprefix_bucket{vmrange=\"3...4\"} 5\nprefix_bucket{vmrange=\"4...5\"} 5\nprefix_bucket{vmrange=\"5...6\"} 5\nprefix_bucket{vmrange=\"6...7\"} 5\nprefix_bucket{vmrange=\"7...8\"} 5\nprefix_bucket{vmrange=\"8...9\"} 5\nprefix_sum 225\nprefix_count 50\n")
+
+	var labels []string
+	var values []uint64
+	h.VisitNonZeroBuckets(func(label string, value uint64) {
+		labels = append(labels, label)
+		values = append(values, value)
+	})
+	labelsExpected := []string{"0...0", "9e-1...1", "1...2", "2...3", "3...4", "4...5", "5...6", "6...7", "7...8", "8...9"}
+	if !reflect.DeepEqual(labels, labelsExpected) {
+		t.Fatalf("unexpected labels; got %v; want %v", labels, labelsExpected)
+	}
+	valuesExpected := []uint64{5, 5, 5, 5, 5, 5, 5, 5, 5, 5}
+	if !reflect.DeepEqual(values, valuesExpected) {
+		t.Fatalf("unexpected values; got %v; want %v", values, valuesExpected)
+	}
 }
 
 func TestHistogramWithTags(t *testing.T) {
diff --git a/metrics_test.go b/metrics_test.go
index 7d30b01..8a046e1 100644
--- a/metrics_test.go
+++ b/metrics_test.go
@@ -106,6 +106,7 @@ func testWritePrometheus() error {
 func expectPanic(t *testing.T, context string, f func()) {
 	t.Helper()
 	defer func() {
+		t.Helper()
 		if r := recover(); r == nil {
 			t.Fatalf("expecting panic in %s", context)
 		}