From 947810b0f282ef5d1860072d835722b9702dd1c7 Mon Sep 17 00:00:00 2001
From: beorn7 <beorn@grafana.com>
Date: Wed, 15 Dec 2021 17:39:32 +0100
Subject: [PATCH] promql: Tweak histogramQuantile

- Simplify the code a bit.

- Cover more corner cases.

- Remove TODO for negative buckets. (I think they are handled. Tests
  will reveal if not.)

Signed-off-by: beorn7 <beorn@grafana.com>
---
 promql/quantile.go | 33 ++++++++++++++++++---------------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/promql/quantile.go b/promql/quantile.go
index cb9783ced..aa19bde73 100644
--- a/promql/quantile.go
+++ b/promql/quantile.go
@@ -151,32 +151,35 @@ func histogramQuantile(q float64, h *histogram.FloatHistogram) float64 {
 	}
 
 	var (
-		bucket *histogram.FloatBucket
+		bucket histogram.FloatBucket
 		count  float64
 		it     = h.AllBucketIterator()
 		rank   = q * h.Count
-		idx    = -1
 	)
-	// TODO(codesome): Do we need any special handling for negative buckets?
 	for it.Next() {
-		idx++
-		b := it.At()
-		count += b.Count
+		bucket = it.At()
+		count += bucket.Count
 		if count >= rank {
-			bucket = &b
 			break
 		}
 	}
-	if bucket == nil {
-		panic("histogramQuantile: not possible")
-	}
-
-	if idx == 0 && bucket.Lower < 0 && bucket.Upper > 0 {
-		// Zero bucket has the result and it happens to be the first
-		// bucket of this histogram.  So we consider 0 to be the lower
-		// bound.
+	if bucket.Lower < 0 && bucket.Upper > 0 && len(h.NegativeBuckets) == 0 {
+		// The result is in the zero bucket and the histogram has no
+		// negative buckets. So we consider 0 to be the lower bound.
 		bucket.Lower = 0
 	}
+	// Due to numerical inaccuracies, we could end up with a higher count
+	// than h.Count. Thus, make sure count is never higher than h.Count.
+	if count > h.Count {
+		count = h.Count
+	}
+	// We could have hit the highest bucket without even reaching the rank
+	// (observations not counted in any bucket are considered "overflow"
+	// observations above the highest bucket), in which case we simple
+	// return the upper limit of the highest explicit bucket.
+	if count < rank {
+		return bucket.Upper
+	}
 
 	rank -= count - bucket.Count
 	// TODO(codesome): Use a better estimation than linear.