optimization: avoid merging single profiles (#4421)

* Skip merge on resolve, for a single partition

* Skip extra merges in report aggregation

* Do not skip merges for specific PGO requests

* Add report aggregator test

* Improve comment

* Fix lint error

Author: aleks-p

pkg/phlaredb/symdb/resolver.go

@@ -257,6 +257,24 @@ func (r *Resolver) Tree() (*model.Tree, error) {
 func (r *Resolver) Pprof() (*googlev1.Profile, error) {
 	span, ctx := opentracing.StartSpanFromContext(r.ctx, "Resolver.Pprof")
 	defer span.Finish()
+
+	if r.canSkipProfileMerge() {
+		// this is the same as the block below, without the profile merge
+		var p *googlev1.Profile
+		err := r.withSymbols(ctx, func(symbols *Symbols, appender *SampleAppender) error {
+			resolved, err := symbols.Pprof(ctx, appender, r.maxNodes, SelectStackTraces(symbols, r.sts))
+			if err != nil {
+				return err
+			}
+			p = resolved
+			return nil
+		})
+		if err != nil {
+			return nil, err
+		}
+		return p, nil
+	}
+
 	var p pprof.ProfileMerge
 	err := r.withSymbols(ctx, func(symbols *Symbols, appender *SampleAppender) error {
 		resolved, err := symbols.Pprof(ctx, appender, r.maxNodes, SelectStackTraces(symbols, r.sts))
@@ -286,6 +304,18 @@ func (r *Resolver) withSymbols(ctx context.Context, fn func(*Symbols, *SampleApp
 	return g.Wait()
 }
 
+func (r *Resolver) canSkipProfileMerge() bool {
+	if len(r.p) > 1 {
+		return false
+	}
+	if r.sts != nil && r.sts.GoPgo != nil && r.sts.GoPgo.AggregateCallees {
+		// we rely on merges to implement GoPgo.AggregateCallees
+		return false
+	}
+
+	return true
+}
+
 func (r *Symbols) Pprof(
 	ctx context.Context,
 	appender *SampleAppender,

pkg/querybackend/report_aggregator.go

@@ -130,22 +130,15 @@ func (ra *reportAggregator) aggregateReportNoCheck(report *queryv1.Report) (err
 	return a.aggregate(report)
 }
 
-func (ra *reportAggregator) aggregateStaged() error {
-	for _, r := range ra.staged {
-		if r != nil {
-			if err := ra.aggregateReportNoCheck(r); err != nil {
-				return err
-			}
-		}
-	}
-	return nil
-}
-
 func (ra *reportAggregator) response() (*queryv1.InvokeResponse, error) {
-	if err := ra.aggregateStaged(); err != nil {
-		return nil, err
-	}
+	// if there are staged reports, we can just add them, no need to aggregate because there is one per type
 	reports := make([]*queryv1.Report, 0, len(ra.staged))
+	for _, st := range ra.staged {
+		if st != nil {
+			reports = append(reports, st)
+		}
+	}
+	// build and add reports from already performed aggregations
 	for t, a := range ra.aggregators {
 		r := a.build()
 		r.ReportType = t

pkg/querybackend/report_aggregator_test.go

@@ -0,0 +1,300 @@
+package querybackend
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
+)
+
+type mockAggregator struct {
+	reports []*queryv1.Report
+	mu      sync.Mutex
+}
+
+func (m *mockAggregator) aggregate(r *queryv1.Report) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.reports = append(m.reports, r)
+	return nil
+}
+
+func (m *mockAggregator) build() *queryv1.Report {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if len(m.reports) == 0 {
+		return &queryv1.Report{}
+	}
+
+	result := &queryv1.Report{
+		ReportType: m.reports[0].ReportType,
+	}
+	return result
+}
+
+func (m *mockAggregator) getReportCount() int {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return len(m.reports)
+}
+
+func mockAggregatorProvider(req *queryv1.InvokeRequest) aggregator {
+	return &mockAggregator{
+		reports: make([]*queryv1.Report, 0),
+	}
+}
+
+func TestReportAggregator_SingleReport(t *testing.T) {
+	reportType := queryv1.ReportType(999) // use a high number that won't conflict with other registrations
+	registerAggregator(reportType, mockAggregatorProvider)
+	defer func() {
+		aggregatorMutex.Lock()
+		delete(aggregators, reportType)
+		aggregatorMutex.Unlock()
+	}()
+
+	request := &queryv1.InvokeRequest{}
+	ra := newAggregator(request)
+
+	report := &queryv1.Report{ReportType: reportType}
+	err := ra.aggregateReport(report)
+	require.NoError(t, err)
+
+	// a single report should be staged and no aggregators should be created
+	assert.Len(t, ra.staged, 1)
+	assert.Len(t, ra.aggregators, 0)
+	assert.Equal(t, report, ra.staged[reportType])
+
+	// the response should contain the single report
+	resp, err := ra.response()
+	require.NoError(t, err)
+	require.Len(t, resp.Reports, 1)
+	assert.Equal(t, report, resp.Reports[0])
+}
+
+func TestReportAggregator_TwoReports(t *testing.T) {
+	reportType := queryv1.ReportType(999)
+	registerAggregator(reportType, mockAggregatorProvider)
+	defer func() {
+		aggregatorMutex.Lock()
+		delete(aggregators, reportType)
+		aggregatorMutex.Unlock()
+	}()
+
+	request := &queryv1.InvokeRequest{}
+	ra := newAggregator(request)
+
+	// the first report should be staged
+	report1 := &queryv1.Report{ReportType: reportType}
+	err := ra.aggregateReport(report1)
+	require.NoError(t, err)
+	assert.Len(t, ra.staged, 1)
+	assert.Len(t, ra.aggregators, 0)
+
+	// the second report should trigger aggregation
+	report2 := &queryv1.Report{ReportType: reportType}
+	err = ra.aggregateReport(report2)
+	require.NoError(t, err)
+	assert.Len(t, ra.aggregators, 1)
+	assert.Nil(t, ra.staged[reportType]) // staged entry should be nil after aggregation
+	agg := ra.aggregators[reportType].(*mockAggregator)
+	assert.Equal(t, 2, agg.getReportCount())
+
+	// the response should contain the aggregated result
+	resp, err := ra.response()
+	require.NoError(t, err)
+	require.Len(t, resp.Reports, 1)
+	assert.Equal(t, reportType, resp.Reports[0].ReportType)
+}
+
+func TestReportAggregator_MultipleTypes(t *testing.T) {
+	type1 := queryv1.ReportType(999)
+	type2 := queryv1.ReportType(998)
+
+	registerAggregator(type1, mockAggregatorProvider)
+	registerAggregator(type2, mockAggregatorProvider)
+	defer func() {
+		aggregatorMutex.Lock()
+		delete(aggregators, type1)
+		delete(aggregators, type2)
+		aggregatorMutex.Unlock()
+	}()
+
+	request := &queryv1.InvokeRequest{}
+	ra := newAggregator(request)
+
+	report1Type1 := &queryv1.Report{ReportType: type1}
+	report2Type2 := &queryv1.Report{ReportType: type2}
+	report3Type1 := &queryv1.Report{ReportType: type1}
+
+	err := ra.aggregateReport(report1Type1)
+	require.NoError(t, err)
+	err = ra.aggregateReport(report2Type2)
+	require.NoError(t, err)
+	err = ra.aggregateReport(report3Type1)
+	require.NoError(t, err)
+
+	// should have one staged report and one aggregator
+	assert.Equal(t, report2Type2, ra.staged[type2])
+	assert.Nil(t, ra.staged[type1])
+	assert.Len(t, ra.aggregators, 1)
+
+	resp, err := ra.response()
+	require.NoError(t, err)
+	require.Len(t, resp.Reports, 2)
+
+	reportTypes := make(map[queryv1.ReportType]bool)
+	for _, r := range resp.Reports {
+		reportTypes[r.ReportType] = true
+	}
+	assert.True(t, reportTypes[type1])
+	assert.True(t, reportTypes[type2])
+}
+
+func TestReportAggregator_NilReport(t *testing.T) {
+	request := &queryv1.InvokeRequest{}
+	ra := newAggregator(request)
+
+	err := ra.aggregateReport(nil)
+	require.NoError(t, err)
+	assert.Len(t, ra.staged, 0)
+	assert.Len(t, ra.aggregators, 0)
+}
+
+func TestReportAggregator_AggregateResponse(t *testing.T) {
+	reportType := queryv1.ReportType(999)
+	registerAggregator(reportType, mockAggregatorProvider)
+	defer func() {
+		aggregatorMutex.Lock()
+		delete(aggregators, reportType)
+		aggregatorMutex.Unlock()
+	}()
+
+	request := &queryv1.InvokeRequest{}
+	ra := newAggregator(request)
+
+	resp := &queryv1.InvokeResponse{
+		Reports: []*queryv1.Report{
+			{ReportType: reportType},
+			{ReportType: reportType},
+		},
+	}
+
+	err := ra.aggregateResponse(resp, nil)
+	require.NoError(t, err)
+
+	assert.Len(t, ra.aggregators, 1)
+	agg := ra.aggregators[reportType].(*mockAggregator)
+	assert.Equal(t, 2, agg.getReportCount())
+}
+
+func TestReportAggregator_ConcurrentAccess(t *testing.T) {
+	reportType := queryv1.ReportType(999)
+	registerAggregator(reportType, mockAggregatorProvider)
+	defer func() {
+		aggregatorMutex.Lock()
+		delete(aggregators, reportType)
+		aggregatorMutex.Unlock()
+	}()
+
+	request := &queryv1.InvokeRequest{}
+	ra := newAggregator(request)
+
+	const numGoroutines = 10
+	const reportsPerGoroutine = 5
+
+	var wg sync.WaitGroup
+	wg.Add(numGoroutines)
+
+	for i := 0; i < numGoroutines; i++ {
+		go func() {
+			defer wg.Done()
+			for j := 0; j < reportsPerGoroutine; j++ {
+				report := &queryv1.Report{ReportType: reportType}
+				err := ra.aggregateReport(report)
+				assert.NoError(t, err)
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	resp, err := ra.response()
+	require.NoError(t, err)
+	assert.Len(t, resp.Reports, 1)
+}
+
+func TestGetAggregator(t *testing.T) {
+	reportType := queryv1.ReportType(999)
+	registerAggregator(reportType, mockAggregatorProvider)
+	defer func() {
+		aggregatorMutex.Lock()
+		delete(aggregators, reportType)
+		aggregatorMutex.Unlock()
+	}()
+
+	request := &queryv1.InvokeRequest{}
+	report := &queryv1.Report{ReportType: reportType}
+
+	agg, err := getAggregator(request, report)
+	require.NoError(t, err)
+	assert.NotNil(t, agg)
+}
+
+func TestGetAggregator_UnknownReportType(t *testing.T) {
+	request := &queryv1.InvokeRequest{}
+	unknownReport := &queryv1.Report{ReportType: queryv1.ReportType(996)}
+	_, err := getAggregator(request, unknownReport)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "unknown build type")
+}
+
+func TestRegisterAggregator_Duplicate(t *testing.T) {
+	reportType := queryv1.ReportType(999)
+
+	registerAggregator(reportType, mockAggregatorProvider)
+	assert.Panics(t, func() {
+		registerAggregator(reportType, mockAggregatorProvider)
+	})
+
+	aggregatorMutex.Lock()
+	delete(aggregators, reportType)
+	aggregatorMutex.Unlock()
+}
+
+func TestQueryReportType(t *testing.T) {
+	queryType := queryv1.QueryType(999)
+	reportType := queryv1.ReportType(999)
+
+	registerQueryReportType(queryType, reportType)
+	defer func() {
+		aggregatorMutex.Lock()
+		delete(queryReportType, queryType)
+		aggregatorMutex.Unlock()
+	}()
+
+	result := QueryReportType(queryType)
+	assert.Equal(t, reportType, result)
+
+	assert.Panics(t, func() {
+		QueryReportType(queryv1.QueryType(889)) // Use an unregistered query type
+	})
+}
+
+func TestRegisterQueryReportType_Duplicate(t *testing.T) {
+	queryType := queryv1.QueryType(999)
+	reportType := queryv1.ReportType(999)
+
+	registerQueryReportType(queryType, reportType)
+	assert.Panics(t, func() {
+		registerQueryReportType(queryType, queryv1.ReportType_REPORT_PPROF)
+	})
+
+	aggregatorMutex.Lock()
+	delete(queryReportType, queryType)
+	aggregatorMutex.Unlock()
+}