fix: sanitize pprof references

pkg/distributor/distributor.go

@@ -230,14 +230,18 @@ func (d *Distributor) Push(ctx context.Context, grpcReq *connect.Request[pushv1.
 }
 
 func (d *Distributor) GetProfileLanguage(series *distributormodel.ProfileSeries) string {
+	if series.Language != "" {
+		return series.Language
+	}
 	if len(series.Samples) == 0 {
 		return "unknown"
 	}
 	lang := series.GetLanguage()
-	if lang != "" {
-		return lang
+	if lang == "" {
+		lang = pprof.GetLanguage(series.Samples[0].Profile, d.logger)
 	}
-	return pprof.GetLanguage(series.Samples[0].Profile, d.logger)
+	series.Language = lang
+	return series.Language
 }
 
 func (d *Distributor) PushParsed(ctx context.Context, req *distributormodel.PushRequest) (resp *connect.Response[pushv1.PushResponse], err error) {
@@ -280,9 +284,6 @@ func (d *Distributor) PushParsed(ctx context.Context, req *distributormodel.Push
 				d.metrics.receivedCompressedBytes.WithLabelValues(profName, tenantID).Observe(float64(len(raw.RawProfile)))
 			}
 			p := raw.Profile
-			if profLanguage == "go" {
-				p.Profile = pprof.FixGoProfile(p.Profile)
-			}
 			decompressedSize := p.SizeVT()
 			d.metrics.receivedDecompressedBytes.WithLabelValues(profName, tenantID).Observe(float64(decompressedSize))
 			d.metrics.receivedSamples.WithLabelValues(profName, tenantID).Observe(float64(len(p.Sample)))
@@ -309,6 +310,9 @@ func (d *Distributor) PushParsed(ctx context.Context, req *distributormodel.Push
 	// therefore it should be done after the rate limit check.
 	for _, series := range req.Series {
 		for _, sample := range series.Samples {
+			if series.Language == "go" {
+				sample.Profile.Profile = pprof.FixGoProfile(sample.Profile.Profile)
+			}
 			sample.Profile.Normalize()
 		}
 	}

pkg/distributor/model/push.go

@@ -28,8 +28,9 @@ type ProfileSample struct {
 }
 
 type ProfileSeries struct {
-	Labels  []*v1.LabelPair
-	Samples []*ProfileSample
+	Labels   []*v1.LabelPair
+	Samples  []*ProfileSample
+	Language string
 }
 
 func (p *ProfileSeries) GetLanguage() string {

pkg/pprof/merge.go

@@ -36,7 +36,7 @@ func (m *ProfileMerge) merge(p *profilev1.Profile, clone bool) error {
 	if p == nil || len(p.StringTable) < 2 {
 		return nil
 	}
-	ConvertIDsToIndices(p)
+	sanitizeReferences(p)
 	var initial bool
 	if m.profile == nil {
 		m.init(p, clone)
@@ -431,88 +431,3 @@ func (t *RewriteTable[K, V, M]) Append(values []V) {
 }
 
 func (t *RewriteTable[K, V, M]) Values() []M { return t.s }
-
-func ConvertIDsToIndices(p *profilev1.Profile) {
-	denseMappings := hasDenseMappings(p)
-	denseLocations := hasDenseLocations(p)
-	denseFunctions := hasDenseFunctions(p)
-	if denseMappings && denseLocations && denseFunctions {
-		// In most cases IDs are dense (do match the element index),
-		// therefore the function does not change anything.
-		return
-	}
-	// NOTE(kolesnikovae):
-	// In some cases IDs is a non-monotonically increasing sequence,
-	// therefore the same map can be reused to avoid re-allocations.
-	t := make(map[uint64]uint64, len(p.Location))
-	if !denseMappings {
-		for i, x := range p.Mapping {
-			idx := uint64(i + 1)
-			x.Id, t[x.Id] = idx, idx
-		}
-		RewriteMappingsWithMap(p, t)
-	}
-	if !denseLocations {
-		for i, x := range p.Location {
-			idx := uint64(i + 1)
-			x.Id, t[x.Id] = idx, idx
-		}
-		RewriteLocationsWithMap(p, t)
-	}
-	if !denseFunctions {
-		for i, x := range p.Function {
-			idx := uint64(i + 1)
-			x.Id, t[x.Id] = idx, idx
-		}
-		RewriteFunctionsWithMap(p, t)
-	}
-}
-
-func hasDenseFunctions(p *profilev1.Profile) bool {
-	for i, f := range p.Function {
-		if f.Id != uint64(i+1) {
-			return false
-		}
-	}
-	return true
-}
-
-func hasDenseLocations(p *profilev1.Profile) bool {
-	for i, loc := range p.Location {
-		if loc.Id != uint64(i+1) {
-			return false
-		}
-	}
-	return true
-}
-
-func hasDenseMappings(p *profilev1.Profile) bool {
-	for i, m := range p.Mapping {
-		if m.Id != uint64(i+1) {
-			return false
-		}
-	}
-	return true
-}
-
-func RewriteFunctionsWithMap(p *profilev1.Profile, n map[uint64]uint64) {
-	for _, loc := range p.Location {
-		for _, line := range loc.Line {
-			line.FunctionId = n[line.FunctionId]
-		}
-	}
-}
-
-func RewriteMappingsWithMap(p *profilev1.Profile, n map[uint64]uint64) {
-	for _, loc := range p.Location {
-		loc.MappingId = n[loc.MappingId]
-	}
-}
-
-func RewriteLocationsWithMap(p *profilev1.Profile, n map[uint64]uint64) {
-	for _, s := range p.Sample {
-		for i, loc := range s.LocationId {
-			s.LocationId[i] = n[loc]
-		}
-	}
-}

pkg/pprof/merge_test.go

@@ -1,8 +1,10 @@
 package pprof
 
 import (
+	"os"
 	"testing"
 
+	"github.com/google/pprof/profile"
 	"github.com/stretchr/testify/require"
 
 	profilev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
@@ -33,51 +35,6 @@ func Test_Merge_Self(t *testing.T) {
 	testhelper.EqualProto(t, p.Profile, m.Profile())
 }
 
-func Test_Merge_ZeroReferences(t *testing.T) {
-	p, err := OpenFile("testdata/go.cpu.labels.pprof")
-	require.NoError(t, err)
-
-	t.Run("mappingID=0", func(t *testing.T) {
-		before := p.Location[10]
-		p.Location[10].MappingId = 0
-		defer func() {
-			p.Location[10] = before
-		}()
-
-		var m ProfileMerge
-		require.NoError(t, m.Merge(p.Profile))
-
-		testhelper.EqualProto(t, p.Profile, m.Profile())
-	})
-
-	t.Run("locationID=0", func(t *testing.T) {
-		before := p.Sample[10].LocationId[0]
-		p.Sample[10].LocationId[0] = 0
-		defer func() {
-			p.Sample[10].LocationId[0] = before
-		}()
-
-		var m ProfileMerge
-		require.NoError(t, m.Merge(p.Profile))
-
-		testhelper.EqualProto(t, p.Profile, m.Profile())
-	})
-
-	t.Run("functionID=0", func(t *testing.T) {
-		before := p.Location[10].Line[0].FunctionId
-		p.Location[10].Line[0].FunctionId = 0
-		defer func() {
-			p.Location[10].Line[0].FunctionId = before
-		}()
-
-		var m ProfileMerge
-		require.NoError(t, m.Merge(p.Profile))
-
-		testhelper.EqualProto(t, p.Profile, m.Profile())
-	})
-
-}
-
 func Test_Merge_Halves(t *testing.T) {
 	p, err := OpenFile("testdata/go.cpu.labels.pprof")
 	require.NoError(t, err)
@@ -470,3 +427,41 @@ func TestMergeEmpty(t *testing.T) {
 	})
 	require.NoError(t, err)
 }
+
+// Benchmark_Merge_self/pprof.MergeNoClone-10         	    4174	    290190 ns/op
+// Benchmark_Merge_self/pprof.Merge-10                	    2722	    421419 ns/op
+// Benchmark_Merge_self/profile.Merge-10              	     802	   1417907 ns/op
+func Benchmark_Merge_self(b *testing.B) {
+	d, err := os.ReadFile("testdata/go.cpu.labels.pprof")
+	require.NoError(b, err)
+
+	b.Run("pprof.MergeNoClone", func(b *testing.B) {
+		p, err := RawFromBytes(d)
+		require.NoError(b, err)
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			var m ProfileMerge
+			require.NoError(b, m.MergeNoClone(p.Profile.CloneVT()))
+		}
+	})
+
+	b.Run("pprof.Merge", func(b *testing.B) {
+		p, err := RawFromBytes(d)
+		require.NoError(b, err)
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			var m ProfileMerge
+			require.NoError(b, m.Merge(p.Profile.CloneVT()))
+		}
+	})
+
+	b.Run("profile.Merge", func(b *testing.B) {
+		p, err := profile.ParseData(d)
+		require.NoError(b, err)
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_, err = profile.Merge([]*profile.Profile{p.Copy()})
+			require.NoError(b, err)
+		}
+	})
+}