refactor(distributor): extract profile stripping into sampling package

pkg/distributor/distributor.go

@@ -114,6 +114,7 @@ type Distributor struct {
 	asyncRequests          sync.WaitGroup
 	ingestionLimitsSampler *ingestlimits.Sampler
 	usageGroupEvaluator    *validation.UsageGroupEvaluator
+	stripper               *sampling.ProfileStripper
 
 	subservices        *services.Manager
 	subservicesWatcher *services.FailureWatcher
@@ -195,6 +196,7 @@ func New(
 		profileReceivedStats:    usagestats.NewMultiCounter("distributor_profiles_received", "lang"),
 		profileScopeStats:       usagestats.NewMultiCounter("distributor_profiles_received_by_scope", "scope"),
 		profileSizeStats:        usagestats.NewMultiStatistics("distributor_profile_sizes", "lang"),
+		stripper:                sampling.NewProfileStripper(),
 	}
 
 	ingesterRoute := writepath.IngesterFunc(d.sendRequestsToIngester)
@@ -609,7 +611,7 @@ func (d *Distributor) pushSeries(ctx context.Context, req *distributormodel.Prof
 		// Language detection reads the string table, which is about to be
 		// stripped; the result is cached in the request.
 		d.GetProfileLanguage(req)
-		stripProfileToTotals(req.Profile.Profile)
+		d.stripper.StripToTotals(req.Profile.Profile)
 		req.Labels = phlaremodel.Labels(req.Labels).InsertSorted(phlaremodel.LabelNameSampled, "true")
 
 		// The stripped part of the profile is discarded, and from here on
@@ -953,70 +955,6 @@ func (d *Distributor) sendRequestsToSegmentWriter(ctx context.Context, req *dist
 	return connect.NewResponse(&pushv1.PushResponse{}), nil
 }
 
-// stripProfileToTotals reduces the profile to a single sample holding the
-// sum of all sample values: stacktraces, symbols, and sample labels are
-// dropped, only the series totals are kept.
-func stripProfileToTotals(p *profilev1.Profile) {
-	var total *profilev1.Sample
-	for _, s := range p.Sample {
-		// Samples that Normalize would drop (value length mismatch,
-		// negative values) must not contribute to the totals.
-		if len(s.Value) != len(p.SampleType) || hasNegativeValue(s) {
-			continue
-		}
-		if total == nil {
-			total = &profilev1.Sample{Value: make([]int64, len(p.SampleType))}
-		}
-		for i, v := range s.Value {
-			total.Value[i] += v
-		}
-	}
-	p.Sample = nil
-	if total != nil {
-		p.Sample = []*profilev1.Sample{total}
-	}
-	p.Location = nil
-	p.Function = nil
-	p.Mapping = nil
-
-	oldStrings := p.StringTable
-	newStrings := []string{""}
-	remap := map[int64]int64{0: 0}
-	intern := func(old int64) int64 {
-		if n, ok := remap[old]; ok {
-			return n
-		}
-		n := int64(len(newStrings))
-		newStrings = append(newStrings, oldStrings[old])
-		remap[old] = n
-		return n
-	}
-	p.DropFrames = intern(p.DropFrames)
-	p.KeepFrames = intern(p.KeepFrames)
-	p.DefaultSampleType = intern(p.DefaultSampleType)
-	for _, vt := range p.SampleType {
-		vt.Type = intern(vt.Type)
-		vt.Unit = intern(vt.Unit)
-	}
-	if p.PeriodType != nil {
-		p.PeriodType.Type = intern(p.PeriodType.Type)
-		p.PeriodType.Unit = intern(p.PeriodType.Unit)
-	}
-	for i, c := range p.Comment {
-		p.Comment[i] = intern(c)
-	}
-	p.StringTable = newStrings
-}
-
-func hasNegativeValue(s *profilev1.Sample) bool {
-	for _, v := range s.Value {
-		if v < 0 {
-			return true
-		}
-	}
-	return false
-}
-
 // profileSizeBytes returns the size of symbols and samples in bytes from a given fullSize.
 func profileSizeBytes(p *profilev1.Profile, fullSize int64) (symbols, samples int64) {
 	// remove samples

pkg/distributor/distributor_test.go

@@ -2696,43 +2696,6 @@ func strippedTestProfile() *profilev1.Profile {
 	}
 }
 
-func TestStripProfileToTotals(t *testing.T) {
-	p := stripTestProfile()
-	stripProfileToTotals(p)
-
-	want := strippedTestProfile()
-	require.Len(t, p.Sample, 1)
-	assert.Equal(t, want.Sample[0].Value, p.Sample[0].Value)
-	assert.Empty(t, p.Sample[0].LocationId)
-	assert.Empty(t, p.Sample[0].Label)
-	assert.Empty(t, p.Location)
-	assert.Empty(t, p.Function)
-	assert.Empty(t, p.Mapping)
-	assert.Equal(t, want.StringTable, p.StringTable)
-	assert.Equal(t, want.SampleType, p.SampleType)
-	assert.Equal(t, want.PeriodType, p.PeriodType)
-	assert.Equal(t, want.SizeVT(), p.SizeVT())
-}
-
-func TestStripProfileToTotals_NoSamples(t *testing.T) {
-	p := stripTestProfile()
-	p.Sample = nil
-	stripProfileToTotals(p)
-	assert.Empty(t, p.Sample)
-	assert.Empty(t, p.Location)
-}
-
-func TestStripProfileToTotals_AllSamplesInvalid(t *testing.T) {
-	p := stripTestProfile()
-	p.Sample = []*profilev1.Sample{
-		{LocationId: []uint64{1}, Value: []int64{-1, 100}},
-		{LocationId: []uint64{1}, Value: []int64{7}},
-	}
-	stripProfileToTotals(p)
-	assert.Empty(t, p.Sample)
-	assert.Empty(t, p.Location)
-}
-
 func newStripTestDistributor(t *testing.T, keepStrippedProfiles bool) (*Distributor, *fakeIngester) {
 	t.Helper()
 	overrides := validation.MockOverrides(func(defaults *validation.Limits, tenantLimits map[string]*validation.Limits) {
@@ -2840,27 +2803,3 @@ func TestPushSeries_DropSampledOutProfiles(t *testing.T) {
 	defer ing.mtx.Unlock()
 	assert.Empty(t, ing.requests)
 }
-
-func TestStripProfileToTotals_SampleLabels(t *testing.T) {
-	p := stripTestProfile()
-	// Indices into the extended string table: span_id=9, abc=10, def=11.
-	p.StringTable = append(p.StringTable, "span_id", "abc", "def")
-	p.Sample = []*profilev1.Sample{
-		{LocationId: []uint64{1}, Value: []int64{1, 10}, Label: []*profilev1.Label{{Key: 9, Str: 10}}},
-		{LocationId: []uint64{2}, Value: []int64{2, 20}, Label: []*profilev1.Label{{Key: 9, Str: 10}}},
-		{LocationId: []uint64{1}, Value: []int64{4, 40}, Label: []*profilev1.Label{{Key: 9, Str: 11}}},
-		{LocationId: []uint64{2}, Value: []int64{8, 80}},
-		{LocationId: []uint64{1}, Value: []int64{16, 160}, Label: []*profilev1.Label{{Key: 9, Num: 42}}},
-	}
-
-	stripProfileToTotals(p)
-
-	require.Len(t, p.Sample, 1)
-	assert.Equal(t, []int64{31, 310}, p.Sample[0].Value)
-	assert.Empty(t, p.Sample[0].Label)
-	assert.Empty(t, p.Sample[0].LocationId)
-	assert.Empty(t, p.Location)
-	assert.Empty(t, p.Function)
-	assert.Empty(t, p.Mapping)
-	assert.Equal(t, []string{"", "samples", "count", "cpu", "nanoseconds"}, p.StringTable)
-}

pkg/distributor/sampling/profile_stripper.go

@@ -0,0 +1,77 @@
+package sampling
+
+import (
+	profilev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
+)
+
+// ProfileStripper reduces profiles that were sampled out but should be kept
+// in a minimal form.
+type ProfileStripper struct{}
+
+func NewProfileStripper() *ProfileStripper {
+	return &ProfileStripper{}
+}
+
+// StripToTotals reduces the profile to a single sample holding the
+// sum of all sample values: stacktraces, symbols, and sample labels are
+// dropped, only the series totals are kept.
+func (*ProfileStripper) StripToTotals(p *profilev1.Profile) {
+	var total *profilev1.Sample
+	for _, s := range p.Sample {
+		// Samples that Normalize would drop (value length mismatch,
+		// negative values) must not contribute to the totals.
+		if len(s.Value) != len(p.SampleType) || hasNegativeValue(s) {
+			continue
+		}
+		if total == nil {
+			total = &profilev1.Sample{Value: make([]int64, len(p.SampleType))}
+		}
+		for i, v := range s.Value {
+			total.Value[i] += v
+		}
+	}
+	p.Sample = nil
+	if total != nil {
+		p.Sample = []*profilev1.Sample{total}
+	}
+	p.Location = nil
+	p.Function = nil
+	p.Mapping = nil
+
+	oldStrings := p.StringTable
+	newStrings := []string{""}
+	remap := map[int64]int64{0: 0}
+	intern := func(old int64) int64 {
+		if n, ok := remap[old]; ok {
+			return n
+		}
+		n := int64(len(newStrings))
+		newStrings = append(newStrings, oldStrings[old])
+		remap[old] = n
+		return n
+	}
+	p.DropFrames = intern(p.DropFrames)
+	p.KeepFrames = intern(p.KeepFrames)
+	p.DefaultSampleType = intern(p.DefaultSampleType)
+	for _, vt := range p.SampleType {
+		vt.Type = intern(vt.Type)
+		vt.Unit = intern(vt.Unit)
+	}
+	if p.PeriodType != nil {
+		p.PeriodType.Type = intern(p.PeriodType.Type)
+		p.PeriodType.Unit = intern(p.PeriodType.Unit)
+	}
+	for i, c := range p.Comment {
+		p.Comment[i] = intern(c)
+	}
+	p.StringTable = newStrings
+}
+
+func hasNegativeValue(s *profilev1.Sample) bool {
+	for _, v := range s.Value {
+		if v < 0 {
+			return true
+		}
+	}
+	return false
+}

pkg/distributor/sampling/profile_stripper_test.go

@@ -0,0 +1,132 @@
+package sampling
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	profilev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
+)
+
+// stripTestProfile returns a deterministic profile without sample labels,
+// so it maps to a single series.
+func stripTestProfile() *profilev1.Profile {
+	return &profilev1.Profile{
+		SampleType: []*profilev1.ValueType{
+			{Type: 1, Unit: 2},
+			{Type: 3, Unit: 4},
+		},
+		Sample: []*profilev1.Sample{
+			{LocationId: []uint64{1, 2}, Value: []int64{10, 1000}},
+			{LocationId: []uint64{2}, Value: []int64{5, 500}},
+			// Dropped by Normalize (negative value): must not count towards totals.
+			{LocationId: []uint64{1}, Value: []int64{-1, 100}},
+			// Dropped by sanitization (value length mismatch): must not count either.
+			{LocationId: []uint64{1}, Value: []int64{7}},
+		},
+		Mapping: []*profilev1.Mapping{{Id: 1, HasFunctions: true}},
+		Location: []*profilev1.Location{
+			{Id: 1, MappingId: 1, Line: []*profilev1.Line{{FunctionId: 1}}},
+			{Id: 2, MappingId: 1, Line: []*profilev1.Line{{FunctionId: 2}}},
+		},
+		Function: []*profilev1.Function{
+			{Id: 1, Name: 5, SystemName: 5, Filename: 6},
+			{Id: 2, Name: 7, SystemName: 7, Filename: 8},
+		},
+		StringTable: []string{
+			"",
+			"samples",
+			"count",
+			"cpu",
+			"nanoseconds",
+			"func-a",
+			"app.py",
+			"func-b",
+			"path-b",
+		},
+		TimeNanos:     1000000000,
+		DurationNanos: 10000000000,
+		PeriodType:    &profilev1.ValueType{Type: 3, Unit: 4},
+		Period:        10000000,
+	}
+}
+
+// strippedTestProfile is what stripTestProfile must be reduced to when it
+// is sampled out: a single sample with the totals and a string table that
+// holds only the sample type units.
+func strippedTestProfile() *profilev1.Profile {
+	return &profilev1.Profile{
+		SampleType: []*profilev1.ValueType{
+			{Type: 1, Unit: 2},
+			{Type: 3, Unit: 4},
+		},
+		Sample:        []*profilev1.Sample{{Value: []int64{15, 1500}}},
+		StringTable:   []string{"", "samples", "count", "cpu", "nanoseconds"},
+		TimeNanos:     1000000000,
+		DurationNanos: 10000000000,
+		PeriodType:    &profilev1.ValueType{Type: 3, Unit: 4},
+		Period:        10000000,
+	}
+}
+
+func TestProfileStripper_StripToTotals(t *testing.T) {
+	p := stripTestProfile()
+	NewProfileStripper().StripToTotals(p)
+
+	want := strippedTestProfile()
+	require.Len(t, p.Sample, 1)
+	assert.Equal(t, want.Sample[0].Value, p.Sample[0].Value)
+	assert.Empty(t, p.Sample[0].LocationId)
+	assert.Empty(t, p.Sample[0].Label)
+	assert.Empty(t, p.Location)
+	assert.Empty(t, p.Function)
+	assert.Empty(t, p.Mapping)
+	assert.Equal(t, want.StringTable, p.StringTable)
+	assert.Equal(t, want.SampleType, p.SampleType)
+	assert.Equal(t, want.PeriodType, p.PeriodType)
+	assert.Equal(t, want.SizeVT(), p.SizeVT())
+}
+
+func TestProfileStripper_StripToTotals_NoSamples(t *testing.T) {
+	p := stripTestProfile()
+	p.Sample = nil
+	NewProfileStripper().StripToTotals(p)
+	assert.Empty(t, p.Sample)
+	assert.Empty(t, p.Location)
+}
+
+func TestProfileStripper_StripToTotals_AllSamplesInvalid(t *testing.T) {
+	p := stripTestProfile()
+	p.Sample = []*profilev1.Sample{
+		{LocationId: []uint64{1}, Value: []int64{-1, 100}},
+		{LocationId: []uint64{1}, Value: []int64{7}},
+	}
+	NewProfileStripper().StripToTotals(p)
+	assert.Empty(t, p.Sample)
+	assert.Empty(t, p.Location)
+}
+
+func TestProfileStripper_StripToTotals_SampleLabels(t *testing.T) {
+	p := stripTestProfile()
+	// Indices into the extended string table: span_id=9, abc=10, def=11.
+	p.StringTable = append(p.StringTable, "span_id", "abc", "def")
+	p.Sample = []*profilev1.Sample{
+		{LocationId: []uint64{1}, Value: []int64{1, 10}, Label: []*profilev1.Label{{Key: 9, Str: 10}}},
+		{LocationId: []uint64{2}, Value: []int64{2, 20}, Label: []*profilev1.Label{{Key: 9, Str: 10}}},
+		{LocationId: []uint64{1}, Value: []int64{4, 40}, Label: []*profilev1.Label{{Key: 9, Str: 11}}},
+		{LocationId: []uint64{2}, Value: []int64{8, 80}},
+		{LocationId: []uint64{1}, Value: []int64{16, 160}, Label: []*profilev1.Label{{Key: 9, Num: 42}}},
+	}
+
+	NewProfileStripper().StripToTotals(p)
+
+	require.Len(t, p.Sample, 1)
+	assert.Equal(t, []int64{31, 310}, p.Sample[0].Value)
+	assert.Empty(t, p.Sample[0].Label)
+	assert.Empty(t, p.Sample[0].LocationId)
+	assert.Empty(t, p.Location)
+	assert.Empty(t, p.Function)
+	assert.Empty(t, p.Mapping)
+	assert.Equal(t, []string{"", "samples", "count", "cpu", "nanoseconds"}, p.StringTable)
+}