Refactor tsdbutil for tests/native histograms (#11948)

* Add float histograms to ChunkFromSamplesGeneric Signed-off-by: Justin Lei <justin.lei@grafana.com> * Add Generate*Samples functions to tsdbutil Signed-off-by: Justin Lei <justin.lei@grafana.com> * PR responses Signed-off-by: Justin Lei <justin.lei@grafana.com> --------- Signed-off-by: Justin Lei <justin.lei@grafana.com>
prometheus · Feb 10, 2023 · af1d9e0 · af1d9e0
1 parent 36b04ef
commit af1d9e0
Show file tree

Hide file tree

Showing 8 changed files with 211 additions and 152 deletions.
diff --git a/promql/engine_test.go b/promql/engine_test.go
@@ -25,7 +25,7 @@ import (
 
 	"github.com/go-kit/log"
 
-	"github.com/prometheus/prometheus/util/stats"
+	"github.com/prometheus/prometheus/tsdb/tsdbutil"
 
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
@@ -35,7 +35,7 @@ import (
 	"github.com/prometheus/prometheus/model/timestamp"
 	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/prometheus/prometheus/storage"
-	"github.com/prometheus/prometheus/tsdb"
+	"github.com/prometheus/prometheus/util/stats"
 )
 
 func TestMain(m *testing.M) {
@@ -3136,7 +3136,7 @@ func TestNativeHistogramRate(t *testing.T) {
 	lbls := labels.FromStrings("__name__", seriesName)
 
 	app := test.Storage().Appender(context.TODO())
-	for i, h := range tsdb.GenerateTestHistograms(100) {
+	for i, h := range tsdbutil.GenerateTestHistograms(100) {
 		_, err := app.AppendHistogram(0, lbls, int64(i)*int64(15*time.Second/time.Millisecond), h, nil)
 		require.NoError(t, err)
 	}
@@ -3181,7 +3181,7 @@ func TestNativeFloatHistogramRate(t *testing.T) {
 	lbls := labels.FromStrings("__name__", seriesName)
 
 	app := test.Storage().Appender(context.TODO())
-	for i, fh := range tsdb.GenerateTestFloatHistograms(100) {
+	for i, fh := range tsdbutil.GenerateTestFloatHistograms(100) {
 		_, err := app.AppendHistogram(0, lbls, int64(i)*int64(15*time.Second/time.Millisecond), nil, fh)
 		require.NoError(t, err)
 	}

diff --git a/rules/manager_test.go b/rules/manager_test.go
@@ -38,8 +38,8 @@ import (
 	"github.com/prometheus/prometheus/promql"
 	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/prometheus/prometheus/storage"
-	"github.com/prometheus/prometheus/tsdb"
 	"github.com/prometheus/prometheus/tsdb/chunkenc"
+	"github.com/prometheus/prometheus/tsdb/tsdbutil"
 	"github.com/prometheus/prometheus/util/teststorage"
 )
 
@@ -1348,7 +1348,7 @@ func TestNativeHistogramsInRecordingRules(t *testing.T) {
 
 	// Add some histograms.
 	db := suite.TSDB()
-	hists := tsdb.GenerateTestHistograms(5)
+	hists := tsdbutil.GenerateTestHistograms(5)
 	ts := time.Now()
 	app := db.Appender(context.Background())
 	for i, h := range hists {

diff --git a/storage/merge_test.go b/storage/merge_test.go
@@ -23,7 +23,6 @@ import (
 
 	"github.com/stretchr/testify/require"
 
-	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/tsdb/chunkenc"
 	"github.com/prometheus/prometheus/tsdb/tsdbutil"
@@ -387,18 +386,11 @@ func TestCompactingChunkSeriesMerger(t *testing.T) {
 
 	// histogramSample returns a histogram that is unique to the ts.
 	histogramSample := func(ts int64) sample {
-		idx := ts + 1
-		return sample{t: ts, h: &histogram.Histogram{
-			Schema:          2,
-			ZeroThreshold:   0.001,
-			ZeroCount:       2 * uint64(idx),
-			Count:           5 * uint64(idx),
-			Sum:             12.34 * float64(idx),
-			PositiveSpans:   []histogram.Span{{Offset: 1, Length: 2}, {Offset: 2, Length: 1}},
-			NegativeSpans:   []histogram.Span{{Offset: 2, Length: 1}, {Offset: 1, Length: 2}},
-			PositiveBuckets: []int64{1 * idx, -1 * idx, 3 * idx},
-			NegativeBuckets: []int64{1 * idx, 2 * idx, 3 * idx},
-		}}
+		return sample{t: ts, h: tsdbutil.GenerateTestHistogram(int(ts + 1))}
+	}
+
+	floatHistogramSample := func(ts int64) sample {
+		return sample{t: ts, fh: tsdbutil.GenerateTestFloatHistogram(int(ts + 1))}
 	}
 
 	for _, tc := range []struct {
@@ -529,6 +521,46 @@ func TestCompactingChunkSeriesMerger(t *testing.T) {
 				[]tsdbutil.Sample{histogramSample(15)},
 			),
 		},
+		{
+			name: "float histogram chunks overlapping",
+			input: []ChunkSeries{
+				NewListChunkSeriesFromSamples(labels.FromStrings("bar", "baz"), []tsdbutil.Sample{floatHistogramSample(0), floatHistogramSample(5)}, []tsdbutil.Sample{floatHistogramSample(10), floatHistogramSample(15)}),
+				NewListChunkSeriesFromSamples(labels.FromStrings("bar", "baz"), []tsdbutil.Sample{floatHistogramSample(2), floatHistogramSample(20)}, []tsdbutil.Sample{floatHistogramSample(25), floatHistogramSample(30)}),
+				NewListChunkSeriesFromSamples(labels.FromStrings("bar", "baz"), []tsdbutil.Sample{floatHistogramSample(18), floatHistogramSample(26)}, []tsdbutil.Sample{floatHistogramSample(31), floatHistogramSample(35)}),
+			},
+			expected: NewListChunkSeriesFromSamples(labels.FromStrings("bar", "baz"),
+				[]tsdbutil.Sample{floatHistogramSample(0), floatHistogramSample(2), floatHistogramSample(5), floatHistogramSample(10), floatHistogramSample(15), floatHistogramSample(18), floatHistogramSample(20), floatHistogramSample(25), floatHistogramSample(26), floatHistogramSample(30)},
+				[]tsdbutil.Sample{floatHistogramSample(31), floatHistogramSample(35)},
+			),
+		},
+		{
+			name: "float histogram chunks overlapping with float chunks",
+			input: []ChunkSeries{
+				NewListChunkSeriesFromSamples(labels.FromStrings("bar", "baz"), []tsdbutil.Sample{floatHistogramSample(0), floatHistogramSample(5)}, []tsdbutil.Sample{floatHistogramSample(10), floatHistogramSample(15)}),
+				NewListChunkSeriesFromSamples(labels.FromStrings("bar", "baz"), []tsdbutil.Sample{sample{1, 1, nil, nil}, sample{12, 12, nil, nil}}, []tsdbutil.Sample{sample{14, 14, nil, nil}}),
+			},
+			expected: NewListChunkSeriesFromSamples(labels.FromStrings("bar", "baz"),
+				[]tsdbutil.Sample{floatHistogramSample(0)},
+				[]tsdbutil.Sample{sample{1, 1, nil, nil}},
+				[]tsdbutil.Sample{floatHistogramSample(5), floatHistogramSample(10)},
+				[]tsdbutil.Sample{sample{12, 12, nil, nil}, sample{14, 14, nil, nil}},
+				[]tsdbutil.Sample{floatHistogramSample(15)},
+			),
+		},
+		{
+			name: "float histogram chunks overlapping with histogram chunks",
+			input: []ChunkSeries{
+				NewListChunkSeriesFromSamples(labels.FromStrings("bar", "baz"), []tsdbutil.Sample{floatHistogramSample(0), floatHistogramSample(5)}, []tsdbutil.Sample{floatHistogramSample(10), floatHistogramSample(15)}),
+				NewListChunkSeriesFromSamples(labels.FromStrings("bar", "baz"), []tsdbutil.Sample{histogramSample(1), histogramSample(12)}, []tsdbutil.Sample{histogramSample(14)}),
+			},
+			expected: NewListChunkSeriesFromSamples(labels.FromStrings("bar", "baz"),
+				[]tsdbutil.Sample{floatHistogramSample(0)},
+				[]tsdbutil.Sample{histogramSample(1)},
+				[]tsdbutil.Sample{floatHistogramSample(5), floatHistogramSample(10)},
+				[]tsdbutil.Sample{histogramSample(12), histogramSample(14)},
+				[]tsdbutil.Sample{floatHistogramSample(15)},
+			),
+		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			merged := m(tc.input...)

diff --git a/tsdb/agent/db_test.go b/tsdb/agent/db_test.go
@@ -54,10 +54,10 @@ func TestDB_InvalidSeries(t *testing.T) {
 	})
 
 	t.Run("Histograms", func(t *testing.T) {
-		_, err := app.AppendHistogram(0, labels.Labels{}, 0, tsdb.GenerateTestHistograms(1)[0], nil)
+		_, err := app.AppendHistogram(0, labels.Labels{}, 0, tsdbutil.GenerateTestHistograms(1)[0], nil)
 		require.ErrorIs(t, err, tsdb.ErrInvalidSample, "should reject empty labels")
 
-		_, err = app.AppendHistogram(0, labels.FromStrings("a", "1", "a", "2"), 0, tsdb.GenerateTestHistograms(1)[0], nil)
+		_, err = app.AppendHistogram(0, labels.FromStrings("a", "1", "a", "2"), 0, tsdbutil.GenerateTestHistograms(1)[0], nil)
 		require.ErrorIs(t, err, tsdb.ErrInvalidSample, "should reject duplicate labels")
 	})
 
@@ -151,7 +151,7 @@ func TestCommit(t *testing.T) {
 	for _, l := range lbls {
 		lset := labels.New(l...)
 
-		histograms := tsdb.GenerateTestHistograms(numHistograms)
+		histograms := tsdbutil.GenerateTestHistograms(numHistograms)
 
 		for i := 0; i < numHistograms; i++ {
 			_, err := app.AppendHistogram(0, lset, int64(i), histograms[i], nil)
@@ -163,7 +163,7 @@ func TestCommit(t *testing.T) {
 	for _, l := range lbls {
 		lset := labels.New(l...)
 
-		floatHistograms := tsdb.GenerateTestFloatHistograms(numHistograms)
+		floatHistograms := tsdbutil.GenerateTestFloatHistograms(numHistograms)
 
 		for i := 0; i < numHistograms; i++ {
 			_, err := app.AppendHistogram(0, lset, int64(i), nil, floatHistograms[i])
@@ -257,7 +257,7 @@ func TestRollback(t *testing.T) {
 	for _, l := range lbls {
 		lset := labels.New(l...)
 
-		histograms := tsdb.GenerateTestHistograms(numHistograms)
+		histograms := tsdbutil.GenerateTestHistograms(numHistograms)
 
 		for i := 0; i < numHistograms; i++ {
 			_, err := app.AppendHistogram(0, lset, int64(i), histograms[i], nil)
@@ -269,7 +269,7 @@ func TestRollback(t *testing.T) {
 	for _, l := range lbls {
 		lset := labels.New(l...)
 
-		floatHistograms := tsdb.GenerateTestFloatHistograms(numHistograms)
+		floatHistograms := tsdbutil.GenerateTestFloatHistograms(numHistograms)
 
 		for i := 0; i < numHistograms; i++ {
 			_, err := app.AppendHistogram(0, lset, int64(i), nil, floatHistograms[i])
@@ -374,7 +374,7 @@ func TestFullTruncateWAL(t *testing.T) {
 	for _, l := range lbls {
 		lset := labels.New(l...)
 
-		histograms := tsdb.GenerateTestHistograms(numHistograms)
+		histograms := tsdbutil.GenerateTestHistograms(numHistograms)
 
 		for i := 0; i < numHistograms; i++ {
 			_, err := app.AppendHistogram(0, lset, int64(lastTs), histograms[i], nil)
@@ -387,7 +387,7 @@ func TestFullTruncateWAL(t *testing.T) {
 	for _, l := range lbls {
 		lset := labels.New(l...)
 
-		floatHistograms := tsdb.GenerateTestFloatHistograms(numHistograms)
+		floatHistograms := tsdbutil.GenerateTestFloatHistograms(numHistograms)
 
 		for i := 0; i < numHistograms; i++ {
 			_, err := app.AppendHistogram(0, lset, int64(lastTs), nil, floatHistograms[i])
@@ -436,7 +436,7 @@ func TestPartialTruncateWAL(t *testing.T) {
 	for _, l := range lbls {
 		lset := labels.New(l...)
 
-		histograms := tsdb.GenerateTestHistograms(numDatapoints)
+		histograms := tsdbutil.GenerateTestHistograms(numDatapoints)
 
 		for i := 0; i < numDatapoints; i++ {
 			_, err := app.AppendHistogram(0, lset, lastTs, histograms[i], nil)
@@ -449,7 +449,7 @@ func TestPartialTruncateWAL(t *testing.T) {
 	for _, l := range lbls {
 		lset := labels.New(l...)
 
-		floatHistograms := tsdb.GenerateTestFloatHistograms(numDatapoints)
+		floatHistograms := tsdbutil.GenerateTestFloatHistograms(numDatapoints)
 
 		for i := 0; i < numDatapoints; i++ {
 			_, err := app.AppendHistogram(0, lset, lastTs, nil, floatHistograms[i])
@@ -475,7 +475,7 @@ func TestPartialTruncateWAL(t *testing.T) {
 	for _, l := range lbls {
 		lset := labels.New(l...)
 
-		histograms := tsdb.GenerateTestHistograms(numDatapoints)
+		histograms := tsdbutil.GenerateTestHistograms(numDatapoints)
 
 		for i := 0; i < numDatapoints; i++ {
 			_, err := app.AppendHistogram(0, lset, lastTs, histograms[i], nil)
@@ -488,7 +488,7 @@ func TestPartialTruncateWAL(t *testing.T) {
 	for _, l := range lbls {
 		lset := labels.New(l...)
 
-		floatHistograms := tsdb.GenerateTestFloatHistograms(numDatapoints)
+		floatHistograms := tsdbutil.GenerateTestFloatHistograms(numDatapoints)
 
 		for i := 0; i < numDatapoints; i++ {
 			_, err := app.AppendHistogram(0, lset, lastTs, nil, floatHistograms[i])
@@ -529,7 +529,7 @@ func TestWALReplay(t *testing.T) {
 	for _, l := range lbls {
 		lset := labels.New(l...)
 
-		histograms := tsdb.GenerateTestHistograms(numHistograms)
+		histograms := tsdbutil.GenerateTestHistograms(numHistograms)
 
 		for i := 0; i < numHistograms; i++ {
 			_, err := app.AppendHistogram(0, lset, lastTs, histograms[i], nil)
@@ -541,7 +541,7 @@ func TestWALReplay(t *testing.T) {
 	for _, l := range lbls {
 		lset := labels.New(l...)
 
-		floatHistograms := tsdb.GenerateTestFloatHistograms(numHistograms)
+		floatHistograms := tsdbutil.GenerateTestFloatHistograms(numHistograms)
 
 		for i := 0; i < numHistograms; i++ {
 			_, err := app.AppendHistogram(0, lset, lastTs, nil, floatHistograms[i])
@@ -622,7 +622,7 @@ func Test_ExistingWAL_NextRef(t *testing.T) {
 	}
 
 	histogramCount := 10
-	histograms := tsdb.GenerateTestHistograms(histogramCount)
+	histograms := tsdbutil.GenerateTestHistograms(histogramCount)
 	// Append <histogramCount> series
 	for i := 0; i < histogramCount; i++ {
 		lset := labels.FromStrings(model.MetricNameLabel, fmt.Sprintf("histogram_%d", i))

diff --git a/tsdb/head.go b/tsdb/head.go
@@ -17,7 +17,6 @@ import (
 	"fmt"
 	"io"
 	"math"
-	"math/rand"
 	"path/filepath"
 	"sync"
 	"time"
@@ -2037,93 +2036,3 @@ func (h *Head) updateWALReplayStatusRead(current int) {
 
 	h.stats.WALReplayStatus.Current = current
 }
-
-func GenerateTestHistograms(n int) (r []*histogram.Histogram) {
-	for i := 0; i < n; i++ {
-		h := GenerateTestHistogram(i)
-		if i > 0 {
-			h.CounterResetHint = histogram.NotCounterReset
-		}
-		r = append(r, h)
-	}
-	return r
-}
-
-// Generates a test histogram, it is up to the user to set any known counter reset hint.
-func GenerateTestHistogram(i int) *histogram.Histogram {
-	return &histogram.Histogram{
-		Count:         10 + uint64(i*8),
-		ZeroCount:     2 + uint64(i),
-		ZeroThreshold: 0.001,
-		Sum:           18.4 * float64(i+1),
-		Schema:        1,
-		PositiveSpans: []histogram.Span{
-			{Offset: 0, Length: 2},
-			{Offset: 1, Length: 2},
-		},
-		PositiveBuckets: []int64{int64(i + 1), 1, -1, 0},
-		NegativeSpans: []histogram.Span{
-			{Offset: 0, Length: 2},
-			{Offset: 1, Length: 2},
-		},
-		NegativeBuckets: []int64{int64(i + 1), 1, -1, 0},
-	}
-}
-
-func GenerateTestGaugeHistograms(n int) (r []*histogram.Histogram) {
-	for x := 0; x < n; x++ {
-		r = append(r, GenerateTestGaugeHistogram(rand.Intn(n)))
-	}
-	return r
-}
-
-func GenerateTestGaugeHistogram(i int) *histogram.Histogram {
-	h := GenerateTestHistogram(i)
-	h.CounterResetHint = histogram.GaugeType
-	return h
-}
-
-func GenerateTestFloatHistograms(n int) (r []*histogram.FloatHistogram) {
-	for i := 0; i < n; i++ {
-		h := GenerateTestFloatHistogram(i)
-		if i > 0 {
-			h.CounterResetHint = histogram.NotCounterReset
-		}
-		r = append(r, h)
-	}
-	return r
-}
-
-// Generates a test float histogram, it is up to the user to set any known counter reset hint.
-func GenerateTestFloatHistogram(i int) *histogram.FloatHistogram {
-	return &histogram.FloatHistogram{
-		Count:         10 + float64(i*8),
-		ZeroCount:     2 + float64(i),
-		ZeroThreshold: 0.001,
-		Sum:           18.4 * float64(i+1),
-		Schema:        1,
-		PositiveSpans: []histogram.Span{
-			{Offset: 0, Length: 2},
-			{Offset: 1, Length: 2},
-		},
-		PositiveBuckets: []float64{float64(i + 1), float64(i + 2), float64(i + 1), float64(i + 1)},
-		NegativeSpans: []histogram.Span{
-			{Offset: 0, Length: 2},
-			{Offset: 1, Length: 2},
-		},
-		NegativeBuckets: []float64{float64(i + 1), float64(i + 2), float64(i + 1), float64(i + 1)},
-	}
-}
-
-func GenerateTestGaugeFloatHistograms(n int) (r []*histogram.FloatHistogram) {
-	for x := 0; x < n; x++ {
-		r = append(r, GenerateTestGaugeFloatHistogram(rand.Intn(n)))
-	}
-	return r
-}
-
-func GenerateTestGaugeFloatHistogram(i int) *histogram.FloatHistogram {
-	h := GenerateTestFloatHistogram(i)
-	h.CounterResetHint = histogram.GaugeType
-	return h
-}