Skip to content

Commit

Permalink
Reuse native histogram test functions from Prometheus (#6059)
Browse files Browse the repository at this point in the history 
* reuse native histogram test functions from Prometheus

Signed-off-by: Ben Ye <benye@amazon.com>

* lint

Signed-off-by: Ben Ye <benye@amazon.com>

---------

Signed-off-by: Ben Ye <benye@amazon.com>
  • Loading branch information
@yeya24
yeya24 authored Jul 8, 2024
1 parent 9965352 commit f623919
Show file tree
Hide file tree
Showing 10 changed files with 48 additions and 115 deletions.
6 changes: 3 additions & 3 deletions integration/e2e/util.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,11 @@ import (
"github.com/prometheus/prometheus/storage"
"github.com/prometheus/prometheus/storage/remote"
"github.com/prometheus/prometheus/tsdb"
"github.com/prometheus/prometheus/tsdb/tsdbutil"
"github.com/thanos-io/thanos/pkg/block/metadata"
"github.com/thanos-io/thanos/pkg/runutil"

cortex_tsdb "github.com/cortexproject/cortex/pkg/storage/tsdb"
histogram_util "github.com/cortexproject/cortex/pkg/util/histogram"
)

func RunCommandAndGetOutput(name string, args ...string) ([]byte, error) {
Expand Down Expand Up @@ -174,10 +174,10 @@ func GenerateHistogramSeries(name string, ts time.Time, floatHistogram bool, add
ph prompb.Histogram
)
if floatHistogram {
fh = histogram_util.GenerateTestFloatHistogram(int(i))
fh = tsdbutil.GenerateTestFloatHistogram(int(i))
ph = remote.FloatHistogramToHistogramProto(tsMillis, fh)
} else {
h = histogram_util.GenerateTestHistogram(int(i))
h = tsdbutil.GenerateTestHistogram(int(i))
ph = remote.HistogramToHistogramProto(tsMillis, h)
}

Expand Down
12 changes: 6 additions & 6 deletions pkg/distributor/distributor_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ import (
"github.com/prometheus/prometheus/model/labels"
"github.com/prometheus/prometheus/model/relabel"
"github.com/prometheus/prometheus/tsdb/chunkenc"
"github.com/prometheus/prometheus/tsdb/tsdbutil"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/weaveworks/common/httpgrpc"
Expand All @@ -45,7 +46,6 @@ import (
"github.com/cortexproject/cortex/pkg/util"
"github.com/cortexproject/cortex/pkg/util/chunkcompat"
"github.com/cortexproject/cortex/pkg/util/flagext"
histogram_util "github.com/cortexproject/cortex/pkg/util/histogram"
"github.com/cortexproject/cortex/pkg/util/limiter"
"github.com/cortexproject/cortex/pkg/util/services"
"github.com/cortexproject/cortex/pkg/util/test"
Expand Down Expand Up @@ -2643,7 +2643,7 @@ func mockWriteRequest(lbls []labels.Labels, value int64, timestampMs int64, hist
if histogram {
histograms = make([]cortexpb.Histogram, len(lbls))
for i := range lbls {
histograms[i] = cortexpb.HistogramToHistogramProto(timestampMs, histogram_util.GenerateTestHistogram(int(value)))
histograms[i] = cortexpb.HistogramToHistogramProto(timestampMs, tsdbutil.GenerateTestHistogram(int(value)))
}
} else {
samples = make([]cortexpb.Sample, len(lbls))
Expand Down Expand Up @@ -2882,7 +2882,7 @@ func makeWriteRequestTimeseries(labels []cortexpb.LabelAdapter, ts int64, value
},
}
if histogram {
t.Histograms = append(t.Histograms, cortexpb.HistogramToHistogramProto(ts, histogram_util.GenerateTestHistogram(value)))
t.Histograms = append(t.Histograms, cortexpb.HistogramToHistogramProto(ts, tsdbutil.GenerateTestHistogram(value)))
} else {
t.Samples = append(t.Samples, cortexpb.Sample{
TimestampMs: ts,
Expand All @@ -2908,7 +2908,7 @@ func makeWriteRequestHA(samples int, replica, cluster string, histogram bool) *c
}
if histogram {
ts.Histograms = []cortexpb.Histogram{
cortexpb.HistogramToHistogramProto(int64(i), histogram_util.GenerateTestHistogram(i)),
cortexpb.HistogramToHistogramProto(int64(i), tsdbutil.GenerateTestHistogram(i)),
}
} else {
ts.Samples = []cortexpb.Sample{
Expand Down Expand Up @@ -3354,8 +3354,8 @@ func TestDistributorValidation(t *testing.T) {
ctx := user.InjectOrgID(context.Background(), "1")
now := model.Now()
future, past := now.Add(5*time.Hour), now.Add(-25*time.Hour)
testHistogram := histogram_util.GenerateTestHistogram(1)
testFloatHistogram := histogram_util.GenerateTestFloatHistogram(1)
testHistogram := tsdbutil.GenerateTestHistogram(1)
testFloatHistogram := tsdbutil.GenerateTestFloatHistogram(1)

for i, tc := range []struct {
metadata []*cortexpb.MetricMetadata
Expand Down
14 changes: 7 additions & 7 deletions pkg/ingester/ingester_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ import (
"github.com/prometheus/prometheus/storage"
"github.com/prometheus/prometheus/tsdb"
"github.com/prometheus/prometheus/tsdb/chunkenc"
"github.com/prometheus/prometheus/tsdb/tsdbutil"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
Expand All @@ -53,7 +54,6 @@ import (
"github.com/cortexproject/cortex/pkg/storage/tsdb/bucketindex"
"github.com/cortexproject/cortex/pkg/util"
"github.com/cortexproject/cortex/pkg/util/chunkcompat"
histogram_util "github.com/cortexproject/cortex/pkg/util/histogram"
"github.com/cortexproject/cortex/pkg/util/services"
"github.com/cortexproject/cortex/pkg/util/test"
"github.com/cortexproject/cortex/pkg/util/validation"
Expand Down Expand Up @@ -658,8 +658,8 @@ func TestIngester_Push(t *testing.T) {
}
userID := "test"

testHistogram := cortexpb.HistogramToHistogramProto(10, histogram_util.GenerateTestHistogram(1))
testFloatHistogram := cortexpb.FloatHistogramToHistogramProto(11, histogram_util.GenerateTestFloatHistogram(1))
testHistogram := cortexpb.HistogramToHistogramProto(10, tsdbutil.GenerateTestHistogram(1))
testFloatHistogram := cortexpb.FloatHistogramToHistogramProto(11, tsdbutil.GenerateTestFloatHistogram(1))
tests := map[string]struct {
reqs []*cortexpb.WriteRequest
expectedErr error
Expand Down Expand Up @@ -953,7 +953,7 @@ func TestIngester_Push(t *testing.T) {
[]cortexpb.Sample{{Value: 1, TimestampMs: 9}},
nil,
[]cortexpb.Histogram{
cortexpb.HistogramToHistogramProto(9, histogram_util.GenerateTestHistogram(1)),
cortexpb.HistogramToHistogramProto(9, tsdbutil.GenerateTestHistogram(1)),
},
cortexpb.API),
},
Expand Down Expand Up @@ -1012,7 +1012,7 @@ func TestIngester_Push(t *testing.T) {
[]cortexpb.Sample{{Value: 1, TimestampMs: 1575043969 - (86400 * 1000)}},
nil,
[]cortexpb.Histogram{
cortexpb.HistogramToHistogramProto(1575043969-(86400*1000), histogram_util.GenerateTestHistogram(1)),
cortexpb.HistogramToHistogramProto(1575043969-(86400*1000), tsdbutil.GenerateTestHistogram(1)),
},
cortexpb.API),
},
Expand Down Expand Up @@ -3153,13 +3153,13 @@ func mockHistogramWriteRequest(t *testing.T, lbls labels.Labels, value int, time
c chunkenc.Chunk
)
if float {
fh = histogram_util.GenerateTestFloatHistogram(value)
fh = tsdbutil.GenerateTestFloatHistogram(value)
histograms = []cortexpb.Histogram{
cortexpb.FloatHistogramToHistogramProto(timestampMs, fh),
}
c = chunkenc.NewFloatHistogramChunk()
} else {
h = histogram_util.GenerateTestHistogram(value)
h = tsdbutil.GenerateTestHistogram(value)
histograms = []cortexpb.Histogram{
cortexpb.HistogramToHistogramProto(timestampMs, h),
}
Expand Down
2 changes: 1 addition & 1 deletion pkg/querier/batch/batch_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,7 @@ func BenchmarkNewChunkMergeIterator_Seek(b *testing.B) {
}

func TestSeekCorrectlyDealWithSinglePointChunks(t *testing.T) {
histograms := histogram_util.GenerateTestHistograms(1000, 1000, 1, 5, 20)
histograms := histogram_util.GenerateTestHistograms(1000, 1000, 1)
for _, enc := range []promchunk.Encoding{
promchunk.PrometheusXorChunk,
promchunk.PrometheusHistogramChunk,
Expand Down
4 changes: 2 additions & 2 deletions pkg/querier/batch/chunk_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ func mkGenericChunk(t require.TestingT, from model.Time, points int, enc promchu
}

func testIter(t require.TestingT, points int, iter chunkenc.Iterator, enc promchunk.Encoding) {
histograms := histogram_util.GenerateTestHistograms(0, 1000, points, 5, 20)
histograms := histogram_util.GenerateTestHistograms(0, 1000, points)
ets := model.TimeFromUnix(0)
for i := 0; i < points; i++ {
require.Equal(t, iter.Next(), enc.ChunkValueType(), strconv.Itoa(i))
Expand All @@ -77,7 +77,7 @@ func testIter(t require.TestingT, points int, iter chunkenc.Iterator, enc promch
}

func testSeek(t require.TestingT, points int, iter chunkenc.Iterator, enc promchunk.Encoding) {
histograms := histogram_util.GenerateTestHistograms(0, 1000, points, 5, 20)
histograms := histogram_util.GenerateTestHistograms(0, 1000, points)
for i := 0; i < points; i += points / 10 {
ets := int64(i * int(step/time.Millisecond))

Expand Down
21 changes: 9 additions & 12 deletions pkg/querier/blocks_store_queryable_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import (
"github.com/prometheus/prometheus/promql"
"github.com/prometheus/prometheus/storage"
"github.com/prometheus/prometheus/tsdb/chunkenc"
"github.com/prometheus/prometheus/tsdb/tsdbutil"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
Expand All @@ -41,7 +42,6 @@ import (
"github.com/cortexproject/cortex/pkg/storegateway"
"github.com/cortexproject/cortex/pkg/storegateway/storegatewaypb"
"github.com/cortexproject/cortex/pkg/util"
histogram_util "github.com/cortexproject/cortex/pkg/util/histogram"
"github.com/cortexproject/cortex/pkg/util/limiter"
"github.com/cortexproject/cortex/pkg/util/services"
"github.com/cortexproject/cortex/pkg/util/validation"
Expand All @@ -65,12 +65,12 @@ func TestBlocksStoreQuerier_Select(t *testing.T) {
series1Label = labels.Label{Name: "series", Value: "1"}
series2Label = labels.Label{Name: "series", Value: "2"}
noOpQueryLimiter = limiter.NewQueryLimiter(0, 0, 0, 0)
testHistogram1 = histogram_util.GenerateTestHistogram(1)
testHistogram2 = histogram_util.GenerateTestHistogram(2)
testHistogram3 = histogram_util.GenerateTestHistogram(3)
testFloatHistogram1 = histogram_util.GenerateTestFloatHistogram(1)
testFloatHistogram2 = histogram_util.GenerateTestFloatHistogram(2)
testFloatHistogram3 = histogram_util.GenerateTestFloatHistogram(3)
testHistogram1 = tsdbutil.GenerateTestHistogram(1)
testHistogram2 = tsdbutil.GenerateTestHistogram(2)
testHistogram3 = tsdbutil.GenerateTestHistogram(3)
testFloatHistogram1 = tsdbutil.GenerateTestFloatHistogram(1)
testFloatHistogram2 = tsdbutil.GenerateTestFloatHistogram(2)
testFloatHistogram3 = tsdbutil.GenerateTestFloatHistogram(3)
)

type valueResult struct {
Expand Down Expand Up @@ -2279,13 +2279,10 @@ func TestBlocksStoreQuerier_PromQLExecution(t *testing.T) {
h := h
// Check sample timestamp is expected.
require.Equal(t, h.T, int64(from)+int64(i)*15000)
expectedH := tsdbutil.GenerateTestGaugeFloatHistogram(int(h.T))
if enc == encoding.PrometheusHistogramChunk {
// GenerateTestHistogram will add 10 for the input value i so subtract 10 here.
expectedH := histogram_util.GenerateTestHistogram(int(h.T - 10))
require.Equal(t, expectedH.ToFloat(nil), h.H)
require.Equal(t, expectedH, h.H)
} else if enc == encoding.PrometheusFloatHistogramChunk {
// GenerateTestHistogram will add 10 for the input value i so subtract 10 here.
expectedH := histogram_util.GenerateTestFloatHistogram(int(h.T - 10))
require.Equal(t, expectedH, h.H)
}
}
Expand Down
10 changes: 5 additions & 5 deletions pkg/querier/querier_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import (
"github.com/prometheus/prometheus/promql"
"github.com/prometheus/prometheus/scrape"
"github.com/prometheus/prometheus/storage"
"github.com/prometheus/prometheus/tsdb/tsdbutil"
"github.com/prometheus/prometheus/util/annotations"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
Expand All @@ -34,7 +35,6 @@ import (
"github.com/cortexproject/cortex/pkg/util"
"github.com/cortexproject/cortex/pkg/util/chunkcompat"
"github.com/cortexproject/cortex/pkg/util/flagext"
histogram_util "github.com/cortexproject/cortex/pkg/util/histogram"
"github.com/cortexproject/cortex/pkg/util/validation"
)

Expand Down Expand Up @@ -140,7 +140,7 @@ var (
for i, point := range series.Histograms {
require.Equal(t, ts, point.T, strconv.Itoa(i))
// Convert expected value to float histogram.
expectedH := histogram_util.GenerateTestFloatHistogram(int(ts - 10))
expectedH := tsdbutil.GenerateTestGaugeFloatHistogram(int(ts))
require.Equal(t, expectedH, point.H, strconv.Itoa(i))
ts += int64(q.step / time.Millisecond)
}
Expand Down Expand Up @@ -202,7 +202,7 @@ var (
for i, point := range series.Histograms {
require.Equal(t, ts, point.T, strconv.Itoa(i))
// Convert expected value to float histogram.
expectedH := histogram_util.GenerateTestFloatHistogram(int(ts - 10))
expectedH := tsdbutil.GenerateTestGaugeFloatHistogram(int(ts))
require.Equal(t, expectedH, point.H, strconv.Itoa(i))
ts += int64(q.step / time.Millisecond)
}
Expand All @@ -222,7 +222,7 @@ var (
from, through := time.Unix(0, 0), end.Time()
require.Equal(t, q.samples(from, through, q.step), len(series.Floats))
for i, point := range series.Floats {
expectedFH := histogram_util.GenerateTestFloatHistogram(int(ts - 10))
expectedFH := tsdbutil.GenerateTestFloatHistogram(int(ts))
require.Equal(t, ts, point.T, strconv.Itoa(i))
require.Equal(t, expectedFH.Sum, point.F, strconv.Itoa(i))
ts += int64(q.step / time.Millisecond)
Expand All @@ -243,7 +243,7 @@ var (
from, through := time.Unix(0, 0), end.Time()
require.Equal(t, q.samples(from, through, q.step), len(series.Floats))
for i, point := range series.Floats {
expectedFH := histogram_util.GenerateTestFloatHistogram(int(ts - 10))
expectedFH := tsdbutil.GenerateTestFloatHistogram(int(ts))
require.Equal(t, ts, point.T, strconv.Itoa(i))
require.Equal(t, expectedFH.Count, point.F, strconv.Itoa(i))
ts += int64(q.step / time.Millisecond)
Expand Down
14 changes: 7 additions & 7 deletions pkg/ruler/compat_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,11 @@ import (
"github.com/prometheus/prometheus/model/value"
"github.com/prometheus/prometheus/promql"
"github.com/prometheus/prometheus/promql/parser"
"github.com/prometheus/prometheus/tsdb/tsdbutil"
"github.com/stretchr/testify/require"
"github.com/weaveworks/common/httpgrpc"

"github.com/cortexproject/cortex/pkg/cortexpb"
histogram_util "github.com/cortexproject/cortex/pkg/util/histogram"
"github.com/cortexproject/cortex/pkg/util/validation"
)

Expand All @@ -45,10 +45,10 @@ func TestPusherAppendable(t *testing.T) {
lbls2 := cortexpb.FromLabelsToLabelAdapters(labels.FromMap(map[string]string{labels.MetricName: "ALERTS", labels.AlertName: "boop"}))
lbls3 := cortexpb.FromLabelsToLabelAdapters(labels.FromMap(map[string]string{labels.MetricName: "ALERTS_FOR_STATE", labels.AlertName: "boop"}))

testHistogram := histogram_util.GenerateTestHistogram(1)
testFloatHistogram := histogram_util.GenerateTestFloatHistogram(2)
testHistogramWithNaN := histogram_util.GenerateTestHistogram(1)
testFloatHistogramWithNaN := histogram_util.GenerateTestFloatHistogram(1)
testHistogram := tsdbutil.GenerateTestHistogram(1)
testFloatHistogram := tsdbutil.GenerateTestFloatHistogram(2)
testHistogramWithNaN := tsdbutil.GenerateTestHistogram(1)
testFloatHistogramWithNaN := tsdbutil.GenerateTestFloatHistogram(1)
testHistogramWithNaN.Sum = math.Float64frombits(value.StaleNaN)
testFloatHistogramWithNaN.Sum = math.Float64frombits(value.StaleNaN)

Expand Down Expand Up @@ -470,9 +470,9 @@ func TestPusherErrors(t *testing.T) {
_, err = a.Append(0, lbls, int64(model.Now()), 123456)
require.NoError(t, err)

_, err = a.AppendHistogram(0, lbls, int64(model.Now()), histogram_util.GenerateTestHistogram(1), nil)
_, err = a.AppendHistogram(0, lbls, int64(model.Now()), tsdbutil.GenerateTestHistogram(1), nil)
require.NoError(t, err)
_, err = a.AppendHistogram(0, lbls, int64(model.Now()), nil, histogram_util.GenerateTestFloatHistogram(2))
_, err = a.AppendHistogram(0, lbls, int64(model.Now()), nil, tsdbutil.GenerateTestFloatHistogram(2))
require.NoError(t, err)

require.Equal(t, tc.returnedError, a.Commit())
Expand Down
76 changes: 6 additions & 70 deletions pkg/util/histogram/testutils.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,81 +13,17 @@

package histogram

import "github.com/prometheus/prometheus/model/histogram"
import (
"github.com/prometheus/prometheus/model/histogram"
"github.com/prometheus/prometheus/tsdb/tsdbutil"
)

// Adapted from Prometheus model/histogram/test_utils.go GenerateBigTestHistograms.
func GenerateTestHistograms(from, step, numHistograms, numSpans, numBuckets int) []*histogram.Histogram {
bucketsPerSide := numBuckets / 2
spanLength := uint32(bucketsPerSide / numSpans)
// Given all bucket deltas are 1, sum bucketsPerSide + 1.
observationCount := bucketsPerSide * (1 + bucketsPerSide)

func GenerateTestHistograms(from, step, numHistograms int) []*histogram.Histogram {
var histograms []*histogram.Histogram
for i := 0; i < numHistograms; i++ {
v := from + i*step
h := &histogram.Histogram{
CounterResetHint: histogram.GaugeType,
Count: uint64(v + observationCount),
ZeroCount: uint64(v),
ZeroThreshold: 1e-128,
Sum: 18.4 * float64(v+1),
Schema: 2,
NegativeSpans: make([]histogram.Span, numSpans),
PositiveSpans: make([]histogram.Span, numSpans),
NegativeBuckets: make([]int64, bucketsPerSide),
PositiveBuckets: make([]int64, bucketsPerSide),
}

for j := 0; j < numSpans; j++ {
s := histogram.Span{Offset: 1, Length: spanLength}
h.NegativeSpans[j] = s
h.PositiveSpans[j] = s
}

for j := 0; j < bucketsPerSide; j++ {
h.NegativeBuckets[j] = 1
h.PositiveBuckets[j] = 1
}

histograms = append(histograms, h)
histograms = append(histograms, tsdbutil.GenerateTestGaugeHistogram(v))
}
return histograms
}

func GenerateTestHistogram(i int) *histogram.Histogram {
bucketsPerSide := 10
spanLength := uint32(2)
// Given all bucket deltas are 1, sum bucketsPerSide + 1.
observationCount := bucketsPerSide * (1 + bucketsPerSide)

v := 10 + i
h := &histogram.Histogram{
CounterResetHint: histogram.GaugeType,
Count: uint64(v + observationCount),
ZeroCount: uint64(v),
ZeroThreshold: 1e-128,
Sum: 18.4 * float64(v+1),
Schema: 2,
NegativeSpans: make([]histogram.Span, 5),
PositiveSpans: make([]histogram.Span, 5),
NegativeBuckets: make([]int64, bucketsPerSide),
PositiveBuckets: make([]int64, bucketsPerSide),
}

for j := 0; j < 5; j++ {
s := histogram.Span{Offset: 1, Length: spanLength}
h.NegativeSpans[j] = s
h.PositiveSpans[j] = s
}

for j := 0; j < bucketsPerSide; j++ {
h.NegativeBuckets[j] = 1
h.PositiveBuckets[j] = 1
}

return h
}

func GenerateTestFloatHistogram(i int) *histogram.FloatHistogram {
return GenerateTestHistogram(i).ToFloat(nil)
}
Loading

0 comments on commit f623919

Please sign in to comment.