From 2d233cf95ef815584217e66f974b6fc88c20d84b Mon Sep 17 00:00:00 2001
From: beorn7 <beorn@grafana.com>
Date: Tue, 3 May 2022 16:24:11 +0200
Subject: [PATCH] Histogram: Fix allFloatBucketIterator

If the zero threshold overlaps with the highest negative bucket and/or
the lowest positive bucket, its upper or lower boundary, respectively,
has to be adjusted. In valid histograms, only ever the highest
negative bucket and/or the lowest positive bucket may overlap with the
zero bucket. This is assumed in this code to simplify the checks.

Signed-off-by: beorn7 <beorn@grafana.com>
---
 model/histogram/float_histogram.go      | 10 +++-
 model/histogram/float_histogram_test.go | 72 ++++++++++++++++++++++++-
 2 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/model/histogram/float_histogram.go b/model/histogram/float_histogram.go
index 475eca013..af1bc0cdf 100644
--- a/model/histogram/float_histogram.go
+++ b/model/histogram/float_histogram.go
@@ -655,7 +655,9 @@ func (h *FloatHistogram) NegativeReverseBucketIterator() FloatBucketIterator {
 
 // AllBucketIterator returns a FloatBucketIterator to iterate over all negative,
 // zero, and positive buckets in ascending order (starting at the lowest bucket
-// and going up).
+// and going up). If the highest negative bucket or the lowest positive bucket
+// overlap with the zero bucket, their upper or lower boundary, respectively, is
+// set to the zero threshold.
 func (h *FloatHistogram) AllBucketIterator() FloatBucketIterator {
 	return &allFloatBucketIterator{
 		h:       h,
@@ -1071,6 +1073,9 @@ func (r *allFloatBucketIterator) Next() bool {
 	case -1:
 		if r.negIter.Next() {
 			r.currBucket = r.negIter.At()
+			if r.currBucket.Upper > -r.h.ZeroThreshold {
+				r.currBucket.Upper = -r.h.ZeroThreshold
+			}
 			return true
 		}
 		r.state = 0
@@ -1092,6 +1097,9 @@ func (r *allFloatBucketIterator) Next() bool {
 	case 1:
 		if r.posIter.Next() {
 			r.currBucket = r.posIter.At()
+			if r.currBucket.Lower < r.h.ZeroThreshold {
+				r.currBucket.Lower = r.h.ZeroThreshold
+			}
 			return true
 		}
 		r.state = 42
diff --git a/model/histogram/float_histogram_test.go b/model/histogram/float_histogram_test.go
index d4e9f8cc5..2b49b5b71 100644
--- a/model/histogram/float_histogram_test.go
+++ b/model/histogram/float_histogram_test.go
@@ -1729,6 +1729,66 @@ func TestAllFloatBucketIterator(t *testing.T) {
 			includeZero: false,
 			includePos:  true,
 		},
+		{
+			h: FloatHistogram{
+				Count:         447,
+				ZeroCount:     42,
+				ZeroThreshold: 0.5, // Coinciding with bucket boundary.
+				Sum:           1008.4,
+				Schema:        0,
+				PositiveSpans: []Span{
+					{Offset: 0, Length: 4},
+					{Offset: 1, Length: 0},
+					{Offset: 3, Length: 3},
+					{Offset: 3, Length: 0},
+					{Offset: 2, Length: 0},
+					{Offset: 5, Length: 3},
+				},
+				PositiveBuckets: []float64{100, 344, 123, 55, 3, 63, 2, 54, 235, 33},
+				NegativeSpans: []Span{
+					{Offset: 0, Length: 3},
+					{Offset: 1, Length: 0},
+					{Offset: 3, Length: 0},
+					{Offset: 3, Length: 4},
+					{Offset: 2, Length: 0},
+					{Offset: 5, Length: 3},
+				},
+				NegativeBuckets: []float64{10, 34, 1230, 54, 67, 63, 2, 554, 235, 33},
+			},
+			includeNeg:  true,
+			includeZero: true,
+			includePos:  true,
+		},
+		{
+			h: FloatHistogram{
+				Count:         447,
+				ZeroCount:     42,
+				ZeroThreshold: 0.6, // Within the bucket closest to zero.
+				Sum:           1008.4,
+				Schema:        0,
+				PositiveSpans: []Span{
+					{Offset: 0, Length: 4},
+					{Offset: 1, Length: 0},
+					{Offset: 3, Length: 3},
+					{Offset: 3, Length: 0},
+					{Offset: 2, Length: 0},
+					{Offset: 5, Length: 3},
+				},
+				PositiveBuckets: []float64{100, 344, 123, 55, 3, 63, 2, 54, 235, 33},
+				NegativeSpans: []Span{
+					{Offset: 0, Length: 3},
+					{Offset: 1, Length: 0},
+					{Offset: 3, Length: 0},
+					{Offset: 3, Length: 4},
+					{Offset: 2, Length: 0},
+					{Offset: 5, Length: 3},
+				},
+				NegativeBuckets: []float64{10, 34, 1230, 54, 67, 63, 2, 554, 235, 33},
+			},
+			includeNeg:  true,
+			includeZero: true,
+			includePos:  true,
+		},
 	}
 
 	for i, c := range cases {
@@ -1738,7 +1798,11 @@ func TestAllFloatBucketIterator(t *testing.T) {
 			if c.includeNeg {
 				it := c.h.NegativeReverseBucketIterator()
 				for it.Next() {
-					expBuckets = append(expBuckets, it.At())
+					b := it.At()
+					if c.includeZero && b.Upper > -c.h.ZeroThreshold {
+						b.Upper = -c.h.ZeroThreshold
+					}
+					expBuckets = append(expBuckets, b)
 				}
 			}
 			if c.includeZero {
@@ -1753,7 +1817,11 @@ func TestAllFloatBucketIterator(t *testing.T) {
 			if c.includePos {
 				it := c.h.PositiveBucketIterator()
 				for it.Next() {
-					expBuckets = append(expBuckets, it.At())
+					b := it.At()
+					if c.includeZero && b.Lower < c.h.ZeroThreshold {
+						b.Lower = c.h.ZeroThreshold
+					}
+					expBuckets = append(expBuckets, b)
 				}
 			}