promql: Add histograms to TestQueryStatistics

Also, fix the bugs exposed by the tests.

Signed-off-by: beorn7 <beorn@grafana.com>

promql/engine.go

@@ -1542,13 +1542,12 @@ func (ev *evaluator) eval(expr parser.Expr) (parser.Value, annotations.Annotatio
 			histSamples := totalHPointSize(ss.Histograms)
 
 			if len(ss.Floats)+histSamples > 0 {
-				if ev.currentSamples+len(ss.Floats)+histSamples <= ev.maxSamples {
+				if ev.currentSamples+len(ss.Floats)+histSamples > ev.maxSamples {
-					mat = append(mat, ss)
-					prevSS = &mat[len(mat)-1]
-					ev.currentSamples += len(ss.Floats) + histSamples
-				} else {
 					ev.error(ErrTooManySamples(env))
 				}
+				mat = append(mat, ss)
+				prevSS = &mat[len(mat)-1]
+				ev.currentSamples += len(ss.Floats) + histSamples
 			}
 			ev.samplesStats.UpdatePeak(ev.currentSamples)
 
@@ -1711,26 +1710,28 @@ func (ev *evaluator) eval(expr parser.Expr) (parser.Value, annotations.Annotatio
 				step++
 				_, f, h, ok := ev.vectorSelectorSingle(it, e, ts)
 				if ok {
-					if ev.currentSamples < ev.maxSamples {
+					if h == nil {
-						if h == nil {
+						ev.currentSamples++
-							if ss.Floats == nil {
+						ev.samplesStats.IncrementSamplesAtStep(step, 1)
-								ss.Floats = reuseOrGetFPointSlices(prevSS, numSteps)
+						if ev.currentSamples > ev.maxSamples {
-							}
+							ev.error(ErrTooManySamples(env))
-							ss.Floats = append(ss.Floats, FPoint{F: f, T: ts})
+						}
-							ev.currentSamples++
+						if ss.Floats == nil {
-							ev.samplesStats.IncrementSamplesAtStep(step, 1)
+							ss.Floats = reuseOrGetFPointSlices(prevSS, numSteps)
-						} else {
-							if ss.Histograms == nil {
-								ss.Histograms = reuseOrGetHPointSlices(prevSS, numSteps)
-							}
-							point := HPoint{H: h, T: ts}
-							ss.Histograms = append(ss.Histograms, point)
-							histSize := point.size()
-							ev.currentSamples += histSize
-							ev.samplesStats.IncrementSamplesAtStep(step, int64(histSize))
 						}
+						ss.Floats = append(ss.Floats, FPoint{F: f, T: ts})
 					} else {
-						ev.error(ErrTooManySamples(env))
+						point := HPoint{H: h, T: ts}
+						histSize := point.size()
+						ev.currentSamples += histSize
+						ev.samplesStats.IncrementSamplesAtStep(step, int64(histSize))
+						if ev.currentSamples > ev.maxSamples {
+							ev.error(ErrTooManySamples(env))
+						}
+						if ss.Histograms == nil {
+							ss.Histograms = reuseOrGetHPointSlices(prevSS, numSteps)
+						}
+						ss.Histograms = append(ss.Histograms, point)
 					}
 				}
 			}
@@ -2170,10 +2171,10 @@ loop:
 					histograms = histograms[:n]
 					continue loop
 				}
-				if ev.currentSamples >= ev.maxSamples {
+				ev.currentSamples += histograms[n].size()
+				if ev.currentSamples > ev.maxSamples {
 					ev.error(ErrTooManySamples(env))
 				}
-				ev.currentSamples += histograms[n].size()
 			}
 		case chunkenc.ValFloat:
 			t, f := buf.At()
@@ -2182,10 +2183,10 @@ loop:
 			}
 			// Values in the buffer are guaranteed to be smaller than maxt.
 			if t >= mintFloats {
-				if ev.currentSamples >= ev.maxSamples {
+				ev.currentSamples++
+				if ev.currentSamples > ev.maxSamples {
 					ev.error(ErrTooManySamples(env))
 				}
-				ev.currentSamples++
 				if floats == nil {
 					floats = getFPointSlice(16)
 				}
@@ -2213,22 +2214,22 @@ loop:
 			histograms = histograms[:n]
 			break
 		}
-		if ev.currentSamples >= ev.maxSamples {
+		ev.currentSamples += histograms[n].size()
+		if ev.currentSamples > ev.maxSamples {
 			ev.error(ErrTooManySamples(env))
 		}
-		ev.currentSamples += histograms[n].size()
 
 	case chunkenc.ValFloat:
 		t, f := it.At()
 		if t == maxt && !value.IsStaleNaN(f) {
-			if ev.currentSamples >= ev.maxSamples {
+			ev.currentSamples++
+			if ev.currentSamples > ev.maxSamples {
 				ev.error(ErrTooManySamples(env))
 			}
 			if floats == nil {
 				floats = getFPointSlice(16)
 			}
 			floats = append(floats, FPoint{T: t, F: f})
-			ev.currentSamples++
 		}
 	}
 	ev.samplesStats.UpdatePeak(ev.currentSamples)

promql/engine_test.go

@@ -755,7 +755,7 @@ load 10s
   metricWith3SampleEvery10Seconds{a="1",b="1"} 1+1x100
   metricWith3SampleEvery10Seconds{a="2",b="2"} 1+1x100
   metricWith3SampleEvery10Seconds{a="3",b="2"} 1+1x100
-  metricWith1HistogramsEvery10Seconds {{schema:1 count:5 sum:20 buckets:[1 2 1 1]}}+{{schema:1 count:10 sum:5 buckets:[1 2 3 4]}}x100
+  metricWith1HistogramEvery10Seconds {{schema:1 count:5 sum:20 buckets:[1 2 1 1]}}+{{schema:1 count:10 sum:5 buckets:[1 2 3 4]}}x100
 `)
 	t.Cleanup(func() { storage.Close() })
 
@@ -799,10 +799,10 @@ load 10s
 		{
 			Query:        "metricWith1HistogramEvery10Seconds",
 			Start:        time.Unix(21, 0),
-			PeakSamples:  1,
+			PeakSamples:  12,
-			TotalSamples: 1, // 1 sample / 10 seconds
+			TotalSamples: 12, // 1 histogram sample of size 12 / 10 seconds
 			TotalSamplesPerStep: stats.TotalSamplesPerStep{
-				21000: 1,
+				21000: 12,
 			},
 		},
 		{
@@ -815,6 +815,15 @@ load 10s
 				21000: 1,
 			},
 		},
+		{
+			Query:        "timestamp(metricWith1HistogramEvery10Seconds)",
+			Start:        time.Unix(21, 0),
+			PeakSamples:  13, // histogram size 12 + 1 extra because of timestamp
+			TotalSamples: 1,  // 1 float sample (because of timestamp) / 10 seconds
+			TotalSamplesPerStep: stats.TotalSamplesPerStep{
+				21000: 1,
+			},
+		},
 		{
 			Query:        "metricWith1SampleEvery10Seconds",
 			Start:        time.Unix(22, 0),
@@ -887,11 +896,20 @@ load 10s
 				201000: 6,
 			},
 		},
+		{
+			Query:        "metricWith1HistogramEvery10Seconds[60s]",
+			Start:        time.Unix(201, 0),
+			PeakSamples:  72,
+			TotalSamples: 72, // 1 histogram (size 12) / 10 seconds * 60 seconds
+			TotalSamplesPerStep: stats.TotalSamplesPerStep{
+				201000: 72,
+			},
+		},
 		{
 			Query:        "max_over_time(metricWith1SampleEvery10Seconds[59s])[20s:5s]",
 			Start:        time.Unix(201, 0),
 			PeakSamples:  10,
-			TotalSamples: 24, // (1 sample / 10 seconds * 60 seconds) * 60/5 (using 59s so we always return 6 samples
+			TotalSamples: 24, // (1 sample / 10 seconds * 60 seconds) * 20/5 (using 59s so we always return 6 samples
 			// as if we run a query on 00 looking back 60 seconds we will return 7 samples;
 			// see next test).
 			TotalSamplesPerStep: stats.TotalSamplesPerStep{
@@ -902,12 +920,22 @@ load 10s
 			Query:        "max_over_time(metricWith1SampleEvery10Seconds[60s])[20s:5s]",
 			Start:        time.Unix(201, 0),
 			PeakSamples:  11,
-			TotalSamples: 26, // (1 sample / 10 seconds * 60 seconds) + 2 as
+			TotalSamples: 26, // (1 sample / 10 seconds * 60 seconds) * 4 + 2 as
 			// max_over_time(metricWith1SampleEvery10Seconds[60s]) @ 190 and 200 will return 7 samples.
 			TotalSamplesPerStep: stats.TotalSamplesPerStep{
 				201000: 26,
 			},
 		},
+		{
+			Query:        "max_over_time(metricWith1HistogramEvery10Seconds[60s])[20s:5s]",
+			Start:        time.Unix(201, 0),
+			PeakSamples:  72,
+			TotalSamples: 312, // (1 histogram (size 12) / 10 seconds * 60 seconds) * 4 + 2 * 12 as
+			// max_over_time(metricWith1SampleEvery10Seconds[60s]) @ 190 and 200 will return 7 samples.
+			TotalSamplesPerStep: stats.TotalSamplesPerStep{
+				201000: 312,
+			},
+		},
 		{
 			Query:        "metricWith1SampleEvery10Seconds[60s] @ 30",
 			Start:        time.Unix(201, 0),
@@ -917,6 +945,15 @@ load 10s
 				201000: 4,
 			},
 		},
+		{
+			Query:        "metricWith1HistogramEvery10Seconds[60s] @ 30",
+			Start:        time.Unix(201, 0),
+			PeakSamples:  48,
+			TotalSamples: 48, // @ modifier force the evaluation to at 30 seconds - So it brings 4 datapoints (0, 10, 20, 30 seconds) * 1 series
+			TotalSamplesPerStep: stats.TotalSamplesPerStep{
+				201000: 48,
+			},
+		},
 		{
 			Query:        "sum(max_over_time(metricWith3SampleEvery10Seconds[60s] @ 30))",
 			Start:        time.Unix(201, 0),
@@ -1045,7 +1082,21 @@ load 10s
 			},
 		},
 		{
-			// timestamp function as a special handling
+			Query:        `metricWith1HistogramEvery10Seconds`,
+			Start:        time.Unix(204, 0),
+			End:          time.Unix(223, 0),
+			Interval:     5 * time.Second,
+			PeakSamples:  48,
+			TotalSamples: 48, // 1 histogram (size 12) per query * 4 steps
+			TotalSamplesPerStep: stats.TotalSamplesPerStep{
+				204000: 12, // aligned to the step time, not the sample time
+				209000: 12,
+				214000: 12,
+				219000: 12,
+			},
+		},
+		{
+			// timestamp function has a special handling
 			Query:        "timestamp(metricWith1SampleEvery10Seconds)",
 			Start:        time.Unix(201, 0),
 			End:          time.Unix(220, 0),
@@ -1059,6 +1110,21 @@ load 10s
 				216000: 1,
 			},
 		},
+		{
+			// timestamp function has a special handling
+			Query:        "timestamp(metricWith1HistogramEvery10Seconds)",
+			Start:        time.Unix(201, 0),
+			End:          time.Unix(220, 0),
+			Interval:     5 * time.Second,
+			PeakSamples:  16,
+			TotalSamples: 4, // 1 sample per query * 4 steps
+			TotalSamplesPerStep: stats.TotalSamplesPerStep{
+				201000: 1,
+				206000: 1,
+				211000: 1,
+				216000: 1,
+			},
+		},
 		{
 			Query:        `max_over_time(metricWith3SampleEvery10Seconds{a="1"}[10s])`,
 			Start:        time.Unix(991, 0),