fix: group by with enforced labels order (#3352)

Author: kolesnikovae

pkg/distributor/distributor.go

@@ -419,9 +419,7 @@ func (d *Distributor) sendRequests(ctx context.Context, req *distributormodel.Pu
 	enforceLabelsOrder := d.limits.EnforceLabelsOrder(tenantID)
 	for i, series := range profileSeries {
 		if enforceLabelsOrder {
-			labels := phlaremodel.Labels(series.Labels)
-			labels.Insert(phlaremodel.LabelNameOrder, phlaremodel.LabelOrderEnforced)
-			series.Labels = labels
+			series.Labels = phlaremodel.Labels(series.Labels).InsertSorted(phlaremodel.LabelNameOrder, phlaremodel.LabelOrderEnforced)
 		}
 		if err = validation.ValidateLabels(d.limits, tenantID, series.Labels); err != nil {
 			validation.DiscardedProfiles.WithLabelValues(string(validation.ReasonOf(err)), tenantID).Add(float64(req.TotalProfiles))

pkg/model/labels.go

@@ -5,17 +5,16 @@ import (
 	"encoding/binary"
 	"encoding/hex"
 	"fmt"
+	"slices"
 	"sort"
 	"strconv"
 	"strings"
 
 	"github.com/cespare/xxhash/v2"
 	pmodel "github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/labels"
-	"github.com/prometheus/prometheus/promql/parser"
 
 	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
-	"github.com/grafana/pyroscope/pkg/slices"
 	"github.com/grafana/pyroscope/pkg/util"
 )
 
@@ -102,52 +101,25 @@ func (ls Labels) Hash() uint64 {
 	return xxhash.Sum64(b)
 }
 
-// HashForLabels returns a hash value for the labels matching the provided names.
-// 'names' have to be sorted in ascending order.
-func (ls Labels) HashForLabels(b []byte, names ...string) (uint64, []byte) {
-	b = b[:0]
-	i, j := 0, 0
-	for i < len(ls) && j < len(names) {
-		if names[j] < ls[i].Name {
-			j++
-		} else if ls[i].Name < names[j] {
-			i++
-		} else {
-			b = append(b, ls[i].Name...)
-			b = append(b, seps[0])
-			b = append(b, ls[i].Value...)
-			b = append(b, seps[0])
-			i++
-			j++
-		}
-	}
-	return xxhash.Sum64(b), b
-}
-
 // BytesWithLabels is just as Bytes(), but only for labels matching names.
-// 'names' have to be sorted in ascending order.
 // It uses an byte invalid character as a separator and so should not be used for printing.
 func (ls Labels) BytesWithLabels(buf []byte, names ...string) []byte {
-	b := bytes.NewBuffer(buf[:0])
-	b.WriteByte(labelSep)
-	i, j := 0, 0
-	for i < len(ls) && j < len(names) {
-		if names[j] < ls[i].Name {
-			j++
-		} else if ls[i].Name < names[j] {
-			i++
-		} else {
-			if b.Len() > 1 {
-				b.WriteByte(seps[0])
+	buf = buf[:0]
+	buf = append(buf, labelSep)
+	for _, name := range names {
+		for _, l := range ls {
+			if l.Name == name {
+				if len(buf) > 1 {
+					buf = append(buf, seps[0])
+				}
+				buf = append(buf, l.Name...)
+				buf = append(buf, seps[0])
+				buf = append(buf, l.Value...)
+				break
 			}
-			b.WriteString(ls[i].Name)
-			b.WriteByte(seps[0])
-			b.WriteString(ls[i].Value)
-			i++
-			j++
 		}
 	}
-	return b.Bytes()
+	return buf
 }
 
 func (ls Labels) ToPrometheusLabels() labels.Labels {
@@ -180,22 +152,18 @@ func IsLabelAllowedForIngestion(name string) bool {
 	return allowed
 }
 
-// WithLabels returns a subset of Labels that matches match with the provided label names.
+// WithLabels returns a subset of Labels that match with the provided label names.
 func (ls Labels) WithLabels(names ...string) Labels {
-	matchedLabels := Labels{}
-
-	nameSet := make(map[string]struct{}, len(names))
-	for _, n := range names {
-		nameSet[n] = struct{}{}
-	}
-
-	for _, v := range ls {
-		if _, ok := nameSet[v.Name]; ok {
-			matchedLabels = append(matchedLabels, v)
+	matched := make(Labels, 0, len(names))
+	for _, name := range names {
+		for _, l := range ls {
+			if l.Name == name {
+				matched = append(matched, l)
+				break
+			}
 		}
 	}
-
-	return matchedLabels
+	return matched
 }
 
 // Get returns the value for the label with the given name.
@@ -221,37 +189,42 @@ func (ls Labels) GetLabel(name string) (*typesv1.LabelPair, bool) {
 
 // Delete removes the first label encountered with the name given in place.
 func (ls Labels) Delete(name string) Labels {
-	return slices.RemoveInPlace(ls, func(pair *typesv1.LabelPair, i int) bool {
-		return pair.Name == name
-	})
+	for i, l := range ls {
+		if l.Name == name {
+			return slices.Delete(ls, i, i+1)
+		}
+	}
+	return ls
 }
 
-// Insert adds the given label to the set of labels.
-// It assumes the labels are ordered
-func (ls *Labels) Insert(name, value string) {
-	// Find the index where the new label should be inserted
+// InsertSorted adds the given label to the set of labels.
+// It assumes the labels are sorted lexicographically.
+func (ls Labels) InsertSorted(name, value string) Labels {
+	// Find the index where the new label should be inserted.
+	// TODO: Use binary search on large label sets.
 	index := -1
-	for i, label := range *ls {
+	for i, label := range ls {
 		if label.Name > name {
 			index = i
 			break
 		}
 		if label.Name == name {
 			label.Value = value
-			return
+			return ls
 		}
 	}
-	// Insert the new label at the found index
+	// Insert the new label at the found index.
 	l := &typesv1.LabelPair{
 		Name:  name,
 		Value: value,
 	}
-	*ls = append(*ls, l)
+	c := append(ls, l)
 	if index == -1 {
-		return
+		return c
 	}
-	copy((*ls)[index+1:], (*ls)[index:])
-	(*ls)[index] = l
+	copy((c)[index+1:], (c)[index:])
+	(c)[index] = l
+	return c
 }
 
 func (ls Labels) Clone() Labels {
@@ -301,22 +274,6 @@ func LabelPairsString(lbs []*typesv1.LabelPair) string {
 	return b.String()
 }
 
-// StringToLabelsPairs converts a string representation of label pairs to a slice of label pairs.
-func StringToLabelsPairs(s string) ([]*typesv1.LabelPair, error) {
-	matchers, err := parser.ParseMetricSelector(s)
-	if err != nil {
-		return nil, err
-	}
-	result := make([]*typesv1.LabelPair, len(matchers))
-	for i := range matchers {
-		result[i] = &typesv1.LabelPair{
-			Name:  matchers[i].Name,
-			Value: matchers[i].Value,
-		}
-	}
-	return result, nil
-}
-
 // LabelsFromStrings creates new labels from pairs of strings.
 func LabelsFromStrings(ss ...string) Labels {
 	if len(ss)%2 != 0 {
@@ -454,33 +411,6 @@ Outer:
 	return append(res, b.add...)
 }
 
-// StableHash is a labels hashing implementation which is guaranteed to not change over time.
-// This function should be used whenever labels hashing backward compatibility must be guaranteed.
-func StableHash(ls labels.Labels) uint64 {
-	// Use xxhash.Sum64(b) for fast path as it's faster.
-	b := make([]byte, 0, 1024)
-	for i, v := range ls {
-		if len(b)+len(v.Name)+len(v.Value)+2 >= cap(b) {
-			// If labels entry is 1KB+ do not allocate whole entry.
-			h := xxhash.New()
-			_, _ = h.Write(b)
-			for _, v := range ls[i:] {
-				_, _ = h.WriteString(v.Name)
-				_, _ = h.Write(seps)
-				_, _ = h.WriteString(v.Value)
-				_, _ = h.Write(seps)
-			}
-			return h.Sum64()
-		}
-
-		b = append(b, v.Name...)
-		b = append(b, seps[0])
-		b = append(b, v.Value...)
-		b = append(b, seps[0])
-	}
-	return xxhash.Sum64(b)
-}
-
 type SessionID uint64
 
 func (s SessionID) String() string {

pkg/model/labels_test.go

@@ -158,6 +158,27 @@ func TestLabels_LabelsEnforcedOrder(t *testing.T) {
 	})
 }
 
+func TestLabels_LabelsEnforcedOrder_BytesWithLabels(t *testing.T) {
+	labels := Labels{
+		{Name: LabelNameProfileType, Value: "cpu"},
+		{Name: LabelNameServiceNamePrivate, Value: "service"},
+		{Name: "__request_id__", Value: "mess"},
+		{Name: "A_label", Value: "bad"},
+		{Name: "foo", Value: "bar"},
+	}
+	sort.Sort(LabelsEnforcedOrder(labels))
+
+	assert.NotEqual(t,
+		labels.BytesWithLabels(nil, "A_label"),
+		labels.BytesWithLabels(nil, "not_a_label"),
+	)
+
+	assert.Equal(t,
+		labels.BytesWithLabels(nil, "A_label"),
+		Labels{{Name: "A_label", Value: "bad"}}.BytesWithLabels(nil, "A_label"),
+	)
+}
+
 func permute[T any](s []T, f func([]T)) {
 	n := len(s)
 	stack := make([]int, n)
@@ -256,8 +277,7 @@ func TestInsert(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			test.labels.Insert(test.insertName, test.insertValue)
-			assert.Equal(t, test.expected, test.labels)
+			assert.Equal(t, test.expected, test.labels.InsertSorted(test.insertName, test.insertValue))
 		})
 	}
 }