fix: Fix slicelabels corruption when used with proto decoding

Alternative to https://github.com/prometheus/prometheus/pull/16957/

Signed-off-by: bwplotka <bwplotka@gmail.com>

# Conflicts:
#	model/textparse/protobufparse.go
#	model/textparse/protobufparse_test.go
This commit is contained in:
bwplotka 2025-09-05 14:36:30 +01:00
parent 0fc2547740
commit 4912fcaab5
12 changed files with 333 additions and 55 deletions

View File

@ -40,7 +40,7 @@ jobs:
- uses: prometheus/promci@443c7fc2397e946bc9f5029e313a9c3441b9b86d # v0.4.7 - uses: prometheus/promci@443c7fc2397e946bc9f5029e313a9c3441b9b86d # v0.4.7
- uses: ./.github/promci/actions/setup_environment - uses: ./.github/promci/actions/setup_environment
- run: go test --tags=dedupelabels ./... - run: go test --tags=dedupelabels ./...
- run: go test --tags=slicelabels -race ./cmd/prometheus ./prompb/io/prometheus/client - run: go test --tags=slicelabels -race ./cmd/prometheus ./model/textparse ./prompb/...
- run: go test --tags=forcedirectio -race ./tsdb/ - run: go test --tags=forcedirectio -race ./tsdb/
- run: GOARCH=386 go test ./... - run: GOARCH=386 go test ./...
- uses: ./.github/promci/actions/check_proto - uses: ./.github/promci/actions/check_proto

View File

@ -775,6 +775,13 @@ func (b *ScratchBuilder) SetSymbolTable(s *SymbolTable) {
b.syms = s b.syms = s
} }
// SetUnsafeAdd allows turning on/off the assumptions that added strings are unsafe
// for reuse. ScratchBuilder implementations that do reuse strings, must clone
// the strings.
//
// Dedupelabels implementation, does not reuse added strings, so this operation is noop.
func (ScratchBuilder) SetUnsafeAdd(bool) {}
func (b *ScratchBuilder) Reset() { func (b *ScratchBuilder) Reset() {
b.add = b.add[:0] b.add = b.add[:0]
b.output = EmptyLabels() b.output = EmptyLabels()
@ -786,12 +793,6 @@ func (b *ScratchBuilder) Add(name, value string) {
b.add = append(b.add, Label{Name: name, Value: value}) b.add = append(b.add, Label{Name: name, Value: value})
} }
// UnsafeAddBytes adds a name/value pair, using []byte instead of string to reduce memory allocations.
// The values must remain live until Labels() is called.
func (b *ScratchBuilder) UnsafeAddBytes(name, value []byte) {
b.add = append(b.add, Label{Name: yoloString(name), Value: yoloString(value)})
}
// Sort the labels added so far by name. // Sort the labels added so far by name.
func (b *ScratchBuilder) Sort() { func (b *ScratchBuilder) Sort() {
slices.SortFunc(b.add, func(a, b Label) int { return strings.Compare(a.Name, b.Name) }) slices.SortFunc(b.add, func(a, b Label) int { return strings.Compare(a.Name, b.Name) })

View File

@ -19,6 +19,7 @@ import (
"bytes" "bytes"
"slices" "slices"
"strings" "strings"
"unique"
"unsafe" "unsafe"
"github.com/cespare/xxhash/v2" "github.com/cespare/xxhash/v2"
@ -437,7 +438,8 @@ func (b *Builder) Labels() Labels {
// ScratchBuilder allows efficient construction of a Labels from scratch. // ScratchBuilder allows efficient construction of a Labels from scratch.
type ScratchBuilder struct { type ScratchBuilder struct {
add Labels add Labels
unsafeAdd bool
} }
// SymbolTable is no-op, just for api parity with dedupelabels. // SymbolTable is no-op, just for api parity with dedupelabels.
@ -466,6 +468,15 @@ func (*ScratchBuilder) SetSymbolTable(*SymbolTable) {
// no-op // no-op
} }
// SetUnsafeAdd allows turning on/off the assumptions that added strings are unsafe
// for reuse. ScratchBuilder implementations that do reuse strings, must clone
// the strings.
//
// SliceLabels will clone all added strings when this option is true.
func (b *ScratchBuilder) SetUnsafeAdd(unsafeAdd bool) {
b.unsafeAdd = unsafeAdd
}
func (b *ScratchBuilder) Reset() { func (b *ScratchBuilder) Reset() {
b.add = b.add[:0] b.add = b.add[:0]
} }
@ -473,16 +484,15 @@ func (b *ScratchBuilder) Reset() {
// Add a name/value pair. // Add a name/value pair.
// Note if you Add the same name twice you will get a duplicate label, which is invalid. // Note if you Add the same name twice you will get a duplicate label, which is invalid.
func (b *ScratchBuilder) Add(name, value string) { func (b *ScratchBuilder) Add(name, value string) {
if b.unsafeAdd {
// Underlying label structure for slicelabels shares memory, so we need to
// copy it if the input is unsafe.
name = unique.Make(name).Value()
value = unique.Make(value).Value()
}
b.add = append(b.add, Label{Name: name, Value: value}) b.add = append(b.add, Label{Name: name, Value: value})
} }
// UnsafeAddBytes adds a name/value pair, using []byte instead of string.
// The default version of this function is unsafe, hence the name.
// This version is safe - it copies the strings immediately - but we keep the same name so everything compiles.
func (b *ScratchBuilder) UnsafeAddBytes(name, value []byte) {
b.add = append(b.add, Label{Name: string(name), Value: string(value)})
}
// Sort the labels added so far by name. // Sort the labels added so far by name.
func (b *ScratchBuilder) Sort() { func (b *ScratchBuilder) Sort() {
slices.SortFunc(b.add, func(a, b Label) int { return strings.Compare(a.Name, b.Name) }) slices.SortFunc(b.add, func(a, b Label) int { return strings.Compare(a.Name, b.Name) })

View File

@ -15,6 +15,12 @@
package labels package labels
import (
"testing"
"github.com/stretchr/testify/require"
)
var expectedSizeOfLabels = []uint64{ // Values must line up with testCaseLabels. var expectedSizeOfLabels = []uint64{ // Values must line up with testCaseLabels.
72, 72,
0, 0,
@ -25,3 +31,55 @@ var expectedSizeOfLabels = []uint64{ // Values must line up with testCaseLabels.
} }
var expectedByteSize = expectedSizeOfLabels // They are identical var expectedByteSize = expectedSizeOfLabels // They are identical
func TestScratchBuilderAdd_Strings(t *testing.T) {
t.Run("safe", func(t *testing.T) {
n := []byte("__name__")
v := []byte("metric1")
l := NewScratchBuilder(0)
l.Add(yoloString(n), yoloString(v))
ret := l.Labels()
// For slicelabels, in default mode strings are reused, so modifying the
// intput will cause `ret` labels to change too.
n[1] = byte('?')
v[2] = byte('?')
require.Empty(t, ret.Get("__name__"))
require.Equal(t, "me?ric1", ret.Get("_?name__"))
})
t.Run("unsafe", func(t *testing.T) {
n := []byte("__name__")
v := []byte("metric1")
l := NewScratchBuilder(0)
l.SetUnsafeAdd(true)
l.Add(yoloString(n), yoloString(v))
ret := l.Labels()
// Changing input strings should be now safe, because we marked adds as unsafe.
n[1] = byte('?')
v[2] = byte('?')
require.Equal(t, "metric1", ret.Get("__name__"))
})
}
/*
export bench=unsafe && go test -tags=slicelabels \
-run '^$' -bench '^BenchmarkScratchBuilderUnsafeAdd' \
-benchtime 5s -count 6 -cpu 2 -timeout 999m \
| tee ${bench}.txt
*/
func BenchmarkScratchBuilderUnsafeAdd(b *testing.B) {
l := NewScratchBuilder(0)
l.SetUnsafeAdd(true)
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
l.Add("__name__", "metric1")
l.add = l.add[:0] // Reset slice so add can be repeated without side effects.
}
}

View File

@ -618,12 +618,6 @@ func (b *ScratchBuilder) Add(name, value string) {
b.add = append(b.add, Label{Name: name, Value: value}) b.add = append(b.add, Label{Name: name, Value: value})
} }
// UnsafeAddBytes adds a name/value pair using []byte instead of string to reduce memory allocations.
// The values must remain live until Labels() is called.
func (b *ScratchBuilder) UnsafeAddBytes(name, value []byte) {
b.add = append(b.add, Label{Name: yoloString(name), Value: yoloString(value)})
}
// Sort the labels added so far by name. // Sort the labels added so far by name.
func (b *ScratchBuilder) Sort() { func (b *ScratchBuilder) Sort() {
slices.SortFunc(b.add, func(a, b Label) int { return strings.Compare(a.Name, b.Name) }) slices.SortFunc(b.add, func(a, b Label) int { return strings.Compare(a.Name, b.Name) })
@ -680,6 +674,13 @@ func (*ScratchBuilder) SetSymbolTable(*SymbolTable) {
// no-op // no-op
} }
// SetUnsafeAdd allows turning on/off the assumptions that added strings are unsafe
// for reuse. ScratchBuilder implementations that do reuse strings, must clone
// the strings.
//
// Stringlabels implementation, does not reuse added strings, so this operation is noop.
func (ScratchBuilder) SetUnsafeAdd(bool) {}
// SizeOfLabels returns the approximate space required for n copies of a label. // SizeOfLabels returns the approximate space required for n copies of a label.
func SizeOfLabels(name, value string, n uint64) uint64 { func SizeOfLabels(name, value string, n uint64) uint64 {
return uint64(labelSize(&Label{Name: name, Value: value})) * n return uint64(labelSize(&Label{Name: name, Value: value})) * n

View File

@ -94,10 +94,12 @@ type ProtobufParser struct {
// NewProtobufParser returns a parser for the payload in the byte slice. // NewProtobufParser returns a parser for the payload in the byte slice.
func NewProtobufParser(b []byte, parseClassicHistograms, convertClassicHistogramsToNHCB, enableTypeAndUnitLabels bool, st *labels.SymbolTable) Parser { func NewProtobufParser(b []byte, parseClassicHistograms, convertClassicHistogramsToNHCB, enableTypeAndUnitLabels bool, st *labels.SymbolTable) Parser {
builder := labels.NewScratchBuilderWithSymbolTable(st, 16)
builder.SetUnsafeAdd(true)
return &ProtobufParser{ return &ProtobufParser{
dec: dto.NewMetricStreamingDecoder(b), dec: dto.NewMetricStreamingDecoder(b),
entryBytes: &bytes.Buffer{}, entryBytes: &bytes.Buffer{},
builder: labels.NewScratchBuilderWithSymbolTable(st, 16), // TODO(bwplotka): Try base builder. builder: builder,
state: EntryInvalid, state: EntryInvalid,
parseClassicHistograms: parseClassicHistograms, parseClassicHistograms: parseClassicHistograms,
@ -622,10 +624,7 @@ func (p *ProtobufParser) onSeriesOrHistogramUpdate() error {
Unit: p.dec.GetUnit(), Unit: p.dec.GetUnit(),
} }
m.AddToLabels(&p.builder) m.AddToLabels(&p.builder)
if err := p.dec.Label(schema.IgnoreOverriddenMetadataLabelsScratchBuilder{ if err := p.dec.Label(m.NewIgnoreOverriddenMetadataLabelScratchBuilder(&p.builder)); err != nil {
Overwrite: m,
ScratchBuilder: &p.builder,
}); err != nil {
return err return err
} }
} else { } else {

View File

@ -16,10 +16,15 @@ package textparse
import ( import (
"bytes" "bytes"
"encoding/binary" "encoding/binary"
"errors"
"fmt" "fmt"
"io"
"math/rand"
"strings"
"testing" "testing"
"github.com/gogo/protobuf/proto" "github.com/gogo/protobuf/proto"
"github.com/gogo/protobuf/types"
"github.com/prometheus/common/model" "github.com/prometheus/common/model"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
@ -27,6 +32,7 @@ import (
"github.com/prometheus/prometheus/model/histogram" "github.com/prometheus/prometheus/model/histogram"
"github.com/prometheus/prometheus/model/labels" "github.com/prometheus/prometheus/model/labels"
dto "github.com/prometheus/prometheus/prompb/io/prometheus/client" dto "github.com/prometheus/prometheus/prompb/io/prometheus/client"
"github.com/prometheus/prometheus/util/pool"
) )
func metricFamiliesToProtobuf(t testing.TB, testMetricFamilies []string) *bytes.Buffer { func metricFamiliesToProtobuf(t testing.TB, testMetricFamilies []string) *bytes.Buffer {
@ -4349,3 +4355,187 @@ metric: <
}) })
} }
} }
func FuzzProtobufParser_Labels(f *testing.F) {
// Add to the "seed corpus" the values that are known to reproduce issues
// which this test has found in the past. These cases run during regular
// testing, as well as the first step of the fuzzing process.
f.Add(true, true, int64(123))
f.Add(true, false, int64(129))
f.Add(false, true, int64(159))
f.Add(false, true, int64(-127))
f.Fuzz(func(
t *testing.T,
parseClassicHistogram bool,
enableTypeAndUnitLabels bool,
randSeed int64,
) {
var (
r = rand.New(rand.NewSource(randSeed))
buffers = pool.New(1+r.Intn(128), 128+r.Intn(1024), 2, func(sz int) interface{} { return make([]byte, 0, sz) })
lastScrapeSize = 0
observedLabels []labels.Labels
st = labels.NewSymbolTable()
)
for i := 0; i < 20; i++ { // run multiple iterations to encounter memory corruptions
// Get buffer from pool like in scrape.go
b := buffers.Get(lastScrapeSize).([]byte)
buf := bytes.NewBuffer(b)
// Generate some scraped data to parse
mf := generateFuzzMetricFamily(r)
protoBuf, err := proto.Marshal(mf)
require.NoError(t, err)
sizeBuf := make([]byte, binary.MaxVarintLen32)
sizeBufSize := binary.PutUvarint(sizeBuf, uint64(len(protoBuf)))
buf.Write(sizeBuf[:sizeBufSize])
buf.Write(protoBuf)
// Use protobuf parser to parse like in real usage
b = buf.Bytes()
p := NewProtobufParser(b, parseClassicHistogram, false, enableTypeAndUnitLabels, st)
for {
entry, err := p.Next()
if errors.Is(err, io.EOF) {
break
}
require.NoError(t, err)
switch entry {
case EntryHelp:
name, help := p.Help()
require.Equal(t, mf.Name, string(name))
require.Equal(t, mf.Help, string(help))
case EntryType:
name, _ := p.Type()
require.Equal(t, mf.Name, string(name))
case EntryUnit:
name, unit := p.Unit()
require.Equal(t, mf.Name, string(name))
require.Equal(t, mf.Unit, string(unit))
case EntrySeries, EntryHistogram:
var lbs labels.Labels
p.Labels(&lbs)
observedLabels = append(observedLabels, lbs)
}
// Get labels from exemplars
for {
var e exemplar.Exemplar
if !p.Exemplar(&e) {
break
}
observedLabels = append(observedLabels, e.Labels)
}
}
// Validate all labels seen so far remain valid. This can find memory corruption issues.
for _, l := range observedLabels {
require.True(t, l.IsValid(model.LegacyValidation), "encountered corrupted labels: %v", l)
}
lastScrapeSize = len(b)
buffers.Put(b)
}
})
}
func generateFuzzMetricFamily(
r *rand.Rand,
) *dto.MetricFamily {
unit := generateValidLabelName(r)
metricName := fmt.Sprintf("%s_%s", generateValidMetricName(r), unit)
metricTypeProto := dto.MetricType(r.Intn(len(dto.MetricType_name)))
metricFamily := &dto.MetricFamily{
Name: metricName,
Help: generateHelp(r),
Type: metricTypeProto,
Unit: unit,
}
metricsCount := r.Intn(20)
for i := 0; i < metricsCount; i++ {
metric := dto.Metric{
Label: generateFuzzLabels(r),
}
switch metricTypeProto {
case dto.MetricType_GAUGE:
metric.Gauge = &dto.Gauge{Value: r.Float64()}
case dto.MetricType_COUNTER:
metric.Counter = &dto.Counter{Value: r.Float64()}
case dto.MetricType_SUMMARY:
metric.Summary = &dto.Summary{Quantile: []dto.Quantile{{Quantile: 0.5, Value: r.Float64()}}}
case dto.MetricType_HISTOGRAM:
metric.Histogram = &dto.Histogram{Exemplars: generateExemplars(r)}
}
metricFamily.Metric = append(metricFamily.Metric, metric)
}
return metricFamily
}
func generateExemplars(r *rand.Rand) []*dto.Exemplar {
exemplarsCount := r.Intn(5)
exemplars := make([]*dto.Exemplar, 0, exemplarsCount)
for i := 0; i < exemplarsCount; i++ {
exemplars = append(exemplars, &dto.Exemplar{
Label: generateFuzzLabels(r),
Value: r.Float64(),
Timestamp: &types.Timestamp{
Seconds: int64(r.Intn(1000000000)),
Nanos: int32(r.Intn(1000000000)),
},
})
}
return exemplars
}
func generateFuzzLabels(r *rand.Rand) []dto.LabelPair {
labelsCount := r.Intn(10)
ls := make([]dto.LabelPair, 0, labelsCount)
for i := 0; i < labelsCount; i++ {
ls = append(ls, dto.LabelPair{
Name: generateValidLabelName(r),
Value: generateValidLabelName(r),
})
}
return ls
}
func generateHelp(r *rand.Rand) string {
result := make([]string, 1+r.Intn(20))
for i := 0; i < len(result); i++ {
result[i] = generateValidLabelName(r)
}
return strings.Join(result, "_")
}
func generateValidLabelName(r *rand.Rand) string {
return generateString(r, validFirstRunes, validLabelNameRunes)
}
func generateValidMetricName(r *rand.Rand) string {
return generateString(r, validFirstRunes, validMetricNameRunes)
}
func generateString(r *rand.Rand, firstRunes, restRunes []rune) string {
result := make([]rune, 1+r.Intn(20))
for i := range result {
if i == 0 {
result[i] = firstRunes[r.Intn(len(firstRunes))]
} else {
result[i] = restRunes[r.Intn(len(restRunes))]
}
}
return string(result)
}
var (
validMetricNameRunes = []rune{
'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
'0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
'_', ':',
}
validLabelNameRunes = validMetricNameRunes[:len(validMetricNameRunes)-1] // skip the colon
validFirstRunes = validMetricNameRunes[:52] // only the letters
)

View File

@ -154,17 +154,22 @@ func (*MetricStreamingDecoder) GetLabel() {
panic("don't use GetLabel, use Label instead") panic("don't use GetLabel, use Label instead")
} }
type scratchBuilder interface { // unsafeLabelAdder is an interface that expects unsafe label adds.
// Typically, this means labels.ScratchBuilder with SetUnsafeAdd set to true.
type unsafeLabelAdder interface {
Add(name, value string) Add(name, value string)
UnsafeAddBytes(name, value []byte)
} }
// Label parses labels into labels scratch builder. Metric name is missing // Label parses labels into unsafeLabelAdder. Metric name is missing
// given the protobuf metric model and has to be deduced from the metric family name. // given the protobuf metric model and has to be deduced from the metric family name.
// TODO: The method name intentionally hide MetricStreamingDecoder.Metric.Label //
// TODO: The Label method name intentionally hide MetricStreamingDecoder.Metric.Label
// field to avoid direct use (it's not parsed). In future generator will generate // field to avoid direct use (it's not parsed). In future generator will generate
// structs tailored for streaming decoding. // structs tailored for streaming decoding.
func (m *MetricStreamingDecoder) Label(b scratchBuilder) error { //
// Unsafe in this context means that bytes and strings are reused across iterations.
// They are live only until the next NextMetric() or NextMetricFamily() call.
func (m *MetricStreamingDecoder) Label(b unsafeLabelAdder) error {
for _, l := range m.labels { for _, l := range m.labels {
if err := parseLabel(m.mData[l.start:l.end], b); err != nil { if err := parseLabel(m.mData[l.start:l.end], b); err != nil {
return err return err
@ -173,11 +178,9 @@ func (m *MetricStreamingDecoder) Label(b scratchBuilder) error {
return nil return nil
} }
// parseLabel is essentially LabelPair.Unmarshal but directly adding into scratch builder // parseLabel is essentially LabelPair.Unmarshal but directly adding into unsafeLabelAdder.
// via UnsafeAddBytes method to reuse strings. func parseLabel(dAtA []byte, b unsafeLabelAdder) error {
func parseLabel(dAtA []byte, b scratchBuilder) error { var unsafeName, unsafeValue string
var name, value []byte
var unsafeName string
l := len(dAtA) l := len(dAtA)
iNdEx := 0 iNdEx := 0
for iNdEx < l { for iNdEx < l {
@ -236,8 +239,7 @@ func parseLabel(dAtA []byte, b scratchBuilder) error {
if postIndex > l { if postIndex > l {
return io.ErrUnexpectedEOF return io.ErrUnexpectedEOF
} }
name = dAtA[iNdEx:postIndex] unsafeName = yoloString(dAtA[iNdEx:postIndex])
unsafeName = yoloString(name)
if !model.UTF8Validation.IsValidLabelName(unsafeName) { if !model.UTF8Validation.IsValidLabelName(unsafeName) {
return fmt.Errorf("invalid label name: %s", unsafeName) return fmt.Errorf("invalid label name: %s", unsafeName)
} }
@ -272,9 +274,9 @@ func parseLabel(dAtA []byte, b scratchBuilder) error {
if postIndex > l { if postIndex > l {
return io.ErrUnexpectedEOF return io.ErrUnexpectedEOF
} }
value = dAtA[iNdEx:postIndex] unsafeValue = yoloString(dAtA[iNdEx:postIndex])
if !utf8.ValidString(yoloString(value)) { if !utf8.ValidString(unsafeValue) {
return fmt.Errorf("invalid label value: %s", value) return fmt.Errorf("invalid label value: %s", unsafeValue)
} }
iNdEx = postIndex iNdEx = postIndex
default: default:
@ -295,7 +297,7 @@ func parseLabel(dAtA []byte, b scratchBuilder) error {
if iNdEx > l { if iNdEx > l {
return io.ErrUnexpectedEOF return io.ErrUnexpectedEOF
} }
b.UnsafeAddBytes(name, value) b.Add(unsafeName, unsafeValue)
return nil return nil
} }

View File

@ -165,6 +165,8 @@ func TestMetricStreamingDecoder(t *testing.T) {
require.Equal(t, 1.546544e+06, d.Metric.GetCounter().GetValue()) require.Equal(t, 1.546544e+06, d.Metric.GetCounter().GetValue())
b := labels.NewScratchBuilder(0) b := labels.NewScratchBuilder(0)
b.SetUnsafeAdd(true)
require.NoError(t, d.Label(&b)) require.NoError(t, d.Label(&b))
require.Equal(t, `{}`, b.Labels().String()) require.Equal(t, `{}`, b.Labels().String())
} }
@ -174,11 +176,14 @@ func TestMetricStreamingDecoder(t *testing.T) {
require.Equal(t, `{checksum="", path="github.com/prometheus/client_golang", version="(devel)"}`, firstMetricLset.String()) require.Equal(t, `{checksum="", path="github.com/prometheus/client_golang", version="(devel)"}`, firstMetricLset.String())
} }
// Regression test against https://github.com/prometheus/prometheus/pull/16946
func TestMetricStreamingDecoder_LabelsCorruption(t *testing.T) { func TestMetricStreamingDecoder_LabelsCorruption(t *testing.T) {
lastScrapeSize := 0 lastScrapeSize := 0
var allPreviousLabels []labels.Labels var allPreviousLabels []labels.Labels
buffers := pool.New(128, 1024, 2, func(sz int) any { return make([]byte, 0, sz) }) buffers := pool.New(128, 1024, 2, func(sz int) any { return make([]byte, 0, sz) })
builder := labels.NewScratchBuilder(0) builder := labels.NewScratchBuilder(0)
builder.SetUnsafeAdd(true)
for _, labelsCount := range []int{1, 2, 3, 5, 8, 5, 3, 2, 1} { for _, labelsCount := range []int{1, 2, 3, 5, 8, 5, 3, 2, 1} {
// Get buffer from pool like in scrape.go // Get buffer from pool like in scrape.go
b := buffers.Get(lastScrapeSize).([]byte) b := buffers.Get(lastScrapeSize).([]byte)
@ -201,9 +206,10 @@ func TestMetricStreamingDecoder_LabelsCorruption(t *testing.T) {
require.NoError(t, d.NextMetricFamily()) require.NoError(t, d.NextMetricFamily())
require.NoError(t, d.NextMetric()) require.NoError(t, d.NextMetric())
// Get the labels // Get the labels. Decode is reusing strings when adding labels. We
// test if scratchBuilder with unsafeAdd set to true handles that.
builder.Reset() builder.Reset()
require.NoError(t, d.Label(&builder)) // <- this uses unsafe strings to create labels require.NoError(t, d.Label(&builder))
lbs := builder.Labels() lbs := builder.Labels()
allPreviousLabels = append(allPreviousLabels, lbs) allPreviousLabels = append(allPreviousLabels, lbs)

View File

@ -139,18 +139,22 @@ func (m Metadata) SetToLabels(b *labels.Builder) {
b.Set(metricUnit, m.Unit) b.Set(metricUnit, m.Unit)
} }
// IgnoreOverriddenMetadataLabelsScratchBuilder is a wrapper over labels scratch builder // NewIgnoreOverriddenMetadataLabelScratchBuilder creates IgnoreOverriddenMetadataLabelScratchBuilder.
// that ignores label additions that would collide with non-empty Overwrite Metadata fields. func (m Metadata) NewIgnoreOverriddenMetadataLabelScratchBuilder(b *labels.ScratchBuilder) *IgnoreOverriddenMetadataLabelScratchBuilder {
type IgnoreOverriddenMetadataLabelsScratchBuilder struct { return &IgnoreOverriddenMetadataLabelScratchBuilder{ScratchBuilder: b, overwrite: m}
*labels.ScratchBuilder }
Overwrite Metadata // IgnoreOverriddenMetadataLabelScratchBuilder is a wrapper over labels.ScratchBuilder
// that ignores label additions that would collide with non-empty Overwrite Metadata fields.
type IgnoreOverriddenMetadataLabelScratchBuilder struct {
*labels.ScratchBuilder
overwrite Metadata
} }
// Add a name/value pair, unless it would collide with the non-empty Overwrite Metadata // Add a name/value pair, unless it would collide with the non-empty Overwrite Metadata
// field. Note if you Add the same name twice you will get a duplicate label, which is invalid. // field. Note if you Add the same name twice you will get a duplicate label, which is invalid.
func (b IgnoreOverriddenMetadataLabelsScratchBuilder) Add(name, value string) { func (b IgnoreOverriddenMetadataLabelScratchBuilder) Add(name, value string) {
if !b.Overwrite.IsEmptyFor(name) { if !b.overwrite.IsEmptyFor(name) {
return return
} }
b.ScratchBuilder.Add(name, value) b.ScratchBuilder.Add(name, value)

View File

@ -142,7 +142,7 @@ func TestIgnoreOverriddenMetadataLabelsScratchBuilder(t *testing.T) {
t.Run(fmt.Sprintf("meta=%#v", tcase.highPrioMeta), func(t *testing.T) { t.Run(fmt.Sprintf("meta=%#v", tcase.highPrioMeta), func(t *testing.T) {
lb := labels.NewScratchBuilder(0) lb := labels.NewScratchBuilder(0)
tcase.highPrioMeta.AddToLabels(&lb) tcase.highPrioMeta.AddToLabels(&lb)
wrapped := &IgnoreOverriddenMetadataLabelsScratchBuilder{ScratchBuilder: &lb, Overwrite: tcase.highPrioMeta} wrapped := tcase.highPrioMeta.NewIgnoreOverriddenMetadataLabelScratchBuilder(&lb)
incomingLabels.Range(func(l labels.Label) { incomingLabels.Range(func(l labels.Label) {
wrapped.Add(l.Name, l.Value) wrapped.Add(l.Name, l.Value)
}) })

View File

@ -19,6 +19,7 @@ import (
"errors" "errors"
"fmt" "fmt"
"math" "math"
"unsafe"
"github.com/prometheus/common/model" "github.com/prometheus/common/model"
@ -204,8 +205,10 @@ type Decoder struct {
builder labels.ScratchBuilder builder labels.ScratchBuilder
} }
func NewDecoder(*labels.SymbolTable) Decoder { // FIXME remove t func NewDecoder(*labels.SymbolTable) Decoder { // FIXME remove t (or use scratch builder with symbols)
return Decoder{builder: labels.NewScratchBuilder(0)} b := labels.NewScratchBuilder(0)
b.SetUnsafeAdd(true)
return Decoder{builder: b}
} }
// Type returns the type of the record. // Type returns the type of the record.
@ -289,6 +292,10 @@ func (*Decoder) Metadata(rec []byte, metadata []RefMetadata) ([]RefMetadata, err
return metadata, nil return metadata, nil
} }
func yoloString(b []byte) string {
return unsafe.String(unsafe.SliceData(b), len(b))
}
// DecodeLabels decodes one set of labels from buf. // DecodeLabels decodes one set of labels from buf.
func (d *Decoder) DecodeLabels(dec *encoding.Decbuf) labels.Labels { func (d *Decoder) DecodeLabels(dec *encoding.Decbuf) labels.Labels {
d.builder.Reset() d.builder.Reset()
@ -296,7 +303,7 @@ func (d *Decoder) DecodeLabels(dec *encoding.Decbuf) labels.Labels {
for range nLabels { for range nLabels {
lName := dec.UvarintBytes() lName := dec.UvarintBytes()
lValue := dec.UvarintBytes() lValue := dec.UvarintBytes()
d.builder.UnsafeAddBytes(lName, lValue) d.builder.Add(yoloString(lName), yoloString(lValue))
} }
return d.builder.Labels() return d.builder.Labels()
} }