fix issue in pprof split when using relabel.LabelDrop (#4365)

* fix broken pprof split when using relabel.LabelDrop

* Add benchmark

* Add test

Author: aleks-p

pkg/model/pprofsplit/pprof_split.go

@@ -174,13 +174,13 @@ func (g *lazyGroup) addSampleGroup(stringTable []string, sg pprof.SampleGroup) {
 }
 
 type groupsWithFingerprints struct {
-	m     map[uint64][]lazyGroup
+	m     map[uint64][]*lazyGroup
 	order []uint64
 }
 
 func newGroupsWithFingerprints() *groupsWithFingerprints {
 	return &groupsWithFingerprints{
-		m: make(map[uint64][]lazyGroup),
+		m: make(map[uint64][]*lazyGroup),
 	}
 }
 
@@ -201,7 +201,7 @@ func (g *groupsWithFingerprints) add(stringTable []string, lbls phlaremodel.Labe
 	}
 
 	// add the labels to the list
-	g.m[fp] = append(g.m[fp], lazyGroup{
+	g.m[fp] = append(g.m[fp], &lazyGroup{
 		sampleGroup: group,
 		labels:      lbls,
 	})

pkg/model/pprofsplit/pprof_split_test.go

@@ -1,6 +1,8 @@
 package pprofsplit
 
 import (
+	"fmt"
+	"math/rand"
 	"testing"
 
 	"github.com/prometheus/common/model"
@@ -350,22 +352,72 @@ func Test_VisitSampleSeries(t *testing.T) {
 				},
 			},
 		},
+		{
+			description: "does not drop samples when a label is dropped",
+			rules: []*relabel.Config{
+				{
+					Action: relabel.LabelDrop,
+					Regex:  relabel.MustNewRegexp("^label_to_drop$"),
+				},
+			},
+			labels: []*typesv1.LabelPair{},
+			profile: &profilev1.Profile{
+				StringTable: []string{"", "label_to_drop", "value_1", "value_2"},
+				Sample: []*profilev1.Sample{
+					{
+						LocationId: []uint64{1, 2},
+						Value:      []int64{2},
+						Label:      []*profilev1.Label{{Key: 1, Str: 2}},
+					},
+					{
+						LocationId: []uint64{1, 3},
+						Value:      []int64{2},
+						Label:      []*profilev1.Label{{Key: 1, Str: 2}},
+					},
+					{
+						LocationId: []uint64{1, 3},
+						Value:      []int64{2},
+						Label:      []*profilev1.Label{{Key: 1, Str: 3}}, // will get merged with the previous one
+					},
+					{
+						LocationId: []uint64{1, 4},
+						Value:      []int64{2},
+						Label:      []*profilev1.Label{{Key: 1, Str: 3}},
+					},
+				},
+			},
+			expectProfilesDropped: 0,
+			expectBytesDropped:    0,
+			expected: []sampleSeries{
+				{
+					labels: []*typesv1.LabelPair{},
+					samples: []*profilev1.Sample{
+						{
+							LocationId: []uint64{1, 2},
+							Label:      []*profilev1.Label{},
+							Value:      []int64{2},
+						},
+						{
+							LocationId: []uint64{1, 3},
+							Label:      []*profilev1.Label{},
+							Value:      []int64{4},
+						},
+						{
+							LocationId: []uint64{1, 4},
+							Label:      []*profilev1.Label{},
+							Value:      []int64{2},
+						},
+					},
+				},
+			},
+		},
 		{
 			description: "ensure only samples of same stacktraces get grouped",
 			labels: []*typesv1.LabelPair{
 				{Name: "__name__", Value: "profile"},
 			},
 			profile: &profilev1.Profile{
 				StringTable: []string{"", "foo", "bar", "binary", "span_id", "aaaabbbbccccdddd", "__name__"},
-				Location: []*profilev1.Location{
-					{Id: 1, MappingId: 1, Line: []*profilev1.Line{{FunctionId: 1}}},
-					{Id: 2, MappingId: 1, Line: []*profilev1.Line{{FunctionId: 2}}},
-				},
-				Mapping: []*profilev1.Mapping{{}, {Id: 1, Filename: 3}},
-				Function: []*profilev1.Function{
-					{Id: 1, Name: 1},
-					{Id: 2, Name: 2},
-				},
 				Sample: []*profilev1.Sample{
 					{
 						LocationId: []uint64{1, 2},
@@ -456,3 +508,51 @@ func Test_VisitSampleSeries(t *testing.T) {
 		})
 	}
 }
+
+func Benchmark_VisitSampleSeries_HighCardinality(b *testing.B) {
+	defaultRelabelConfigs := validation.MockDefaultOverrides().IngestionRelabelingRules("")
+	defaultRelabelConfigs = append(defaultRelabelConfigs, &relabel.Config{
+		Action: relabel.LabelDrop,
+		Regex:  relabel.MustNewRegexp("^high_cardinality_label$"),
+	})
+
+	stringTable := []string{"", "foo", "bar", "binary", "span_id", "aaaabbbbccccdddd", "high_cardinality_label"}
+	highCardinalityOffset := int64(len(stringTable))
+	for i := 0; i < 10000; i++ {
+		stringTable = append(stringTable, fmt.Sprintf("value_%d", i))
+	}
+
+	profile := &profilev1.Profile{
+		StringTable: stringTable,
+		Location:    []*profilev1.Location{{Id: 1, MappingId: 1, Line: []*profilev1.Line{{FunctionId: 1}}}},
+		Mapping:     []*profilev1.Mapping{{}, {Id: 1, Filename: 3}},
+		Function:    []*profilev1.Function{{Id: 1, Name: 1}},
+	}
+
+	for i := 0; i < 30000; i++ {
+		labelValue := highCardinalityOffset + int64(i/10)
+		if rand.Float64() < 0.3 {
+			labelValue = highCardinalityOffset - 2 // lower the cardinality to create large groups
+		}
+		labels := []*profilev1.Label{
+			{Key: highCardinalityOffset - 1, Str: labelValue},
+		}
+		profile.Sample = append(profile.Sample, &profilev1.Sample{
+			LocationId: []uint64{uint64(i + 1)},
+			Value:      []int64{2},
+			Label:      labels,
+		})
+	}
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		visitor := new(mockVisitor)
+		err := VisitSampleSeries(profile, []*typesv1.LabelPair{
+			{Name: "__name__", Value: "profile"},
+			{Name: "foo", Value: "bar"},
+		}, defaultRelabelConfigs, visitor)
+		require.NoError(b, err)
+	}
+}