From e9a1e26ab700c1206c749b0b069fff99c8426aac Mon Sep 17 00:00:00 2001 From: Jeanette Tan Date: Sat, 22 Apr 2023 02:27:15 +0800 Subject: [PATCH] Perform integer/float histogram type checking on conversions, and use a consistent method for determining integer vs float histogram Signed-off-by: Jeanette Tan --- prompb/custom.go | 5 +++++ storage/remote/codec.go | 24 +++++++++++++++--------- storage/remote/codec_test.go | 2 +- storage/remote/queue_manager_test.go | 2 +- storage/remote/write_handler.go | 2 +- storage/remote/write_handler_test.go | 2 +- 6 files changed, 24 insertions(+), 13 deletions(-) diff --git a/prompb/custom.go b/prompb/custom.go index 4b07187bd2..13d6e0f0cd 100644 --- a/prompb/custom.go +++ b/prompb/custom.go @@ -20,6 +20,11 @@ import ( func (m Sample) T() int64 { return m.Timestamp } func (m Sample) V() float64 { return m.Value } +func (h Histogram) IsFloatHistogram() bool { + _, ok := h.GetCount().(*Histogram_CountFloat) + return ok +} + func (r *ChunkedReadResponse) PooledMarshal(p *sync.Pool) ([]byte, error) { size := r.Size() data, ok := p.Get().(*[]byte) diff --git a/storage/remote/codec.go b/storage/remote/codec.go index 02c84a3e6c..9a683b908f 100644 --- a/storage/remote/codec.go +++ b/storage/remote/codec.go @@ -455,10 +455,10 @@ func (c *concreteSeriesIterator) Seek(t int64) chunkenc.ValueType { } func getHistogramValType(h *prompb.Histogram) chunkenc.ValueType { - if _, isInt := h.GetCount().(*prompb.Histogram_CountInt); isInt { - return chunkenc.ValHistogram + if h.IsFloatHistogram() { + return chunkenc.ValFloatHistogram } - return chunkenc.ValFloatHistogram + return chunkenc.ValHistogram } // At implements chunkenc.Iterator. @@ -624,9 +624,11 @@ func exemplarProtoToExemplar(ep prompb.Exemplar) exemplar.Exemplar { } // HistogramProtoToHistogram extracts a (normal integer) Histogram from the -// provided proto message. The caller has to make sure that the proto message -// represents an integer histogram and not a float histogram. +// provided proto message. func HistogramProtoToHistogram(hp prompb.Histogram) *histogram.Histogram { + if hp.IsFloatHistogram() { + panic("don't call HistogramProtoToHistogram on a float histogram") + } return &histogram.Histogram{ CounterResetHint: histogram.CounterResetHint(hp.ResetHint), Schema: hp.Schema, @@ -642,9 +644,11 @@ func HistogramProtoToHistogram(hp prompb.Histogram) *histogram.Histogram { } // FloatHistogramProtoToFloatHistogram extracts a float Histogram from the -// provided proto message to a Float Histogram. The caller has to make sure that -// the proto message represents a float histogram and not an integer histogram. +// provided proto message to a Float Histogram. func FloatHistogramProtoToFloatHistogram(hp prompb.Histogram) *histogram.FloatHistogram { + if !hp.IsFloatHistogram() { + panic("don't call FloatHistogramProtoToFloatHistogram on an integer histogram") + } return &histogram.FloatHistogram{ CounterResetHint: histogram.CounterResetHint(hp.ResetHint), Schema: hp.Schema, @@ -660,9 +664,11 @@ func FloatHistogramProtoToFloatHistogram(hp prompb.Histogram) *histogram.FloatHi } // HistogramProtoToFloatHistogram extracts and converts a (normal integer) histogram from the provided proto message -// to a float histogram. The caller has to make sure that the proto message represents an integer histogram and not a -// float histogram. +// to a float histogram. func HistogramProtoToFloatHistogram(hp prompb.Histogram) *histogram.FloatHistogram { + if hp.IsFloatHistogram() { + panic("don't call HistogramProtoToFloatHistogram on a float histogram") + } return &histogram.FloatHistogram{ CounterResetHint: histogram.CounterResetHint(hp.ResetHint), Schema: hp.Schema, diff --git a/storage/remote/codec_test.go b/storage/remote/codec_test.go index 27e2cc704d..dbd5cec219 100644 --- a/storage/remote/codec_test.go +++ b/storage/remote/codec_test.go @@ -528,7 +528,7 @@ func TestNilHistogramProto(*testing.T) { // This function will panic if it impromperly handles nil // values, causing the test to fail. HistogramProtoToHistogram(prompb.Histogram{}) - FloatHistogramProtoToFloatHistogram(prompb.Histogram{}) + HistogramProtoToFloatHistogram(prompb.Histogram{}) } func exampleHistogram() histogram.Histogram { diff --git a/storage/remote/queue_manager_test.go b/storage/remote/queue_manager_test.go index a57c3bf7b1..b43258ff06 100644 --- a/storage/remote/queue_manager_test.go +++ b/storage/remote/queue_manager_test.go @@ -802,7 +802,7 @@ func (c *TestWriteClient) Store(_ context.Context, req []byte) error { for _, histogram := range ts.Histograms { count++ - if histogram.GetCountFloat() > 0 || histogram.GetZeroCountFloat() > 0 { + if histogram.IsFloatHistogram() { c.receivedFloatHistograms[seriesName] = append(c.receivedFloatHistograms[seriesName], histogram) } else { c.receivedHistograms[seriesName] = append(c.receivedHistograms[seriesName], histogram) diff --git a/storage/remote/write_handler.go b/storage/remote/write_handler.go index 5aa1130888..1f4c43e595 100644 --- a/storage/remote/write_handler.go +++ b/storage/remote/write_handler.go @@ -125,7 +125,7 @@ func (h *writeHandler) write(ctx context.Context, req *prompb.WriteRequest) (err } for _, hp := range ts.Histograms { - if hp.GetCountFloat() > 0 || hp.GetZeroCountFloat() > 0 { // It is a float histogram. + if hp.IsFloatHistogram() { fhs := FloatHistogramProtoToFloatHistogram(hp) _, err = app.AppendHistogram(0, labels, hp.Timestamp, nil, fhs) } else { diff --git a/storage/remote/write_handler_test.go b/storage/remote/write_handler_test.go index e7a88ddc23..3bce5f1d88 100644 --- a/storage/remote/write_handler_test.go +++ b/storage/remote/write_handler_test.go @@ -68,7 +68,7 @@ func TestRemoteWriteHandler(t *testing.T) { } for _, hp := range ts.Histograms { - if hp.GetCountFloat() > 0 || hp.GetZeroCountFloat() > 0 { // It is a float histogram. + if hp.IsFloatHistogram() { fh := FloatHistogramProtoToFloatHistogram(hp) require.Equal(t, mockHistogram{labels, hp.Timestamp, nil, fh}, appendable.histograms[k]) } else {