chore: parallelize mapping symbolization to improve query performance (#4594)

* chore: parallelize mapping symbolization to improve query performance

Author: marcsanmi

pkg/symbolizer/symbolizer.go

@@ -15,6 +15,7 @@ import (
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/prometheus/client_golang/prometheus"
+	"golang.org/x/sync/errgroup"
 
 	googlev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
 	"github.com/grafana/pyroscope/lidia"
@@ -26,21 +27,35 @@ type DebuginfodClient interface {
 }
 
 type Config struct {
-	DebuginfodURL string `yaml:"debuginfod_url"`
+	DebuginfodURL            string `yaml:"debuginfod_url"`
+	MaxDebuginfodConcurrency int    `yaml:"max_debuginfod_concurrency" category:"advanced"`
 }
 
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.StringVar(&cfg.DebuginfodURL, "symbolizer.debuginfod-url", "https://debuginfod.elfutils.org", "URL of the debuginfod server")
+	f.IntVar(&cfg.MaxDebuginfodConcurrency, "symbolizer.max-debuginfod-concurrency", 10, "Maximum number of concurrent symbolization requests to debuginfod server.")
+}
+
+func (cfg *Config) Validate() error {
+	if cfg.MaxDebuginfodConcurrency < 1 {
+		return fmt.Errorf("invalid max-debuginfod-concurrency value, must be positive")
+	}
+	return nil
 }
 
 type Symbolizer struct {
 	logger  log.Logger
 	client  DebuginfodClient
 	bucket  objstore.Bucket
 	metrics *metrics
+	cfg     Config
 }
 
 func New(logger log.Logger, cfg Config, reg prometheus.Registerer, bucket objstore.Bucket) (*Symbolizer, error) {
+	if err := cfg.Validate(); err != nil {
+		return nil, err
+	}
+
 	metrics := newMetrics(reg)
 
 	client, err := NewDebuginfodClient(logger, cfg.DebuginfodURL, metrics)
@@ -53,6 +68,7 @@ func New(logger log.Logger, cfg Config, reg prometheus.Registerer, bucket objsto
 		client:  client,
 		bucket:  bucket,
 		metrics: metrics,
+		cfg:     cfg,
 	}, nil
 }
 
@@ -84,36 +100,103 @@ func (s *Symbolizer) SymbolizePprof(ctx context.Context, profile *googlev1.Profi
 		stringMap[str] = int64(i)
 	}
 
-	var allSymbolizedLocs []symbolizedLocation
+	allSymbolizedLocs, err := s.symbolizeMappingsConcurrently(ctx, profile, locationsByMapping)
+	if err != nil {
+		return fmt.Errorf("symbolizing mappings: %w", err)
+	}
+
+	s.updateAllSymbolsInProfile(profile, allSymbolizedLocs, stringMap)
+
+	return nil
+}
+
+// symbolizeMappingsConcurrently symbolizes multiple mappings concurrently with a concurrency limit.
+func (s *Symbolizer) symbolizeMappingsConcurrently(
+	ctx context.Context,
+	profile *googlev1.Profile,
+	locationsByMapping map[uint64][]*googlev1.Location,
+) ([]symbolizedLocation, error) {
+	maxConcurrency := s.cfg.MaxDebuginfodConcurrency
+	if maxConcurrency <= 0 {
+		maxConcurrency = 10
+	}
+
+	type mappingJob struct {
+		mappingID uint64
+		locations []*googlev1.Location
+	}
 
+	type mappingResult struct {
+		mappingID uint64
+		locations []symbolizedLocation
+	}
+
+	totalLocs := 0
+	jobs := make(chan mappingJob, len(locationsByMapping))
 	for mappingID, locations := range locationsByMapping {
-		mapping := profile.Mapping[mappingID-1]
+		totalLocs += len(locations)
+		jobs <- mappingJob{mappingID: mappingID, locations: locations}
+	}
+	close(jobs)
 
-		binaryName, err := s.extractBinaryName(profile, mapping)
-		if err != nil {
-			return fmt.Errorf("extract binary name: %w", err)
-		}
+	// Process jobs concurrently with errgroup for proper error handling
+	g, ctx := errgroup.WithContext(ctx)
+	g.SetLimit(maxConcurrency)
 
-		buildID, err := s.extractBuildID(profile, mapping)
-		if err != nil {
-			return fmt.Errorf("extract build ID: %w", err)
-		}
+	// Results channel with buffer to avoid blocking jobs
+	results := make(chan mappingResult, len(locationsByMapping))
 
-		req := s.createSymbolizationRequest(binaryName, buildID, locations)
-		s.symbolize(ctx, &req)
+	for job := range jobs {
+		job := job
+		g.Go(func() error {
+			mapping := profile.Mapping[job.mappingID-1]
 
-		for i, loc := range locations {
-			allSymbolizedLocs = append(allSymbolizedLocs, symbolizedLocation{
-				loc:     loc,
-				symLoc:  req.locations[i],
-				mapping: mapping,
-			})
-		}
+			binaryName, err := s.extractBinaryName(profile, mapping)
+			if err != nil {
+				return fmt.Errorf("extract binary name for mapping %d: %w", job.mappingID, err)
+			}
+
+			buildID, err := s.extractBuildID(profile, mapping)
+			if err != nil {
+				return fmt.Errorf("extract build ID for mapping %d: %w", job.mappingID, err)
+			}
+
+			req := s.createSymbolizationRequest(binaryName, buildID, job.locations)
+			s.symbolize(ctx, &req)
+
+			// Collect symbolized locations for this mapping
+			symbolizedLocs := make([]symbolizedLocation, len(job.locations))
+			for i, loc := range job.locations {
+				symbolizedLocs[i] = symbolizedLocation{
+					loc:     loc,
+					symLoc:  req.locations[i],
+					mapping: mapping,
+				}
+			}
+
+			select {
+			case results <- mappingResult{mappingID: job.mappingID, locations: symbolizedLocs}:
+			case <-ctx.Done():
+				return ctx.Err()
+			}
+
+			return nil
+		})
 	}
 
-	s.updateAllSymbolsInProfile(profile, allSymbolizedLocs, stringMap)
+	err := g.Wait()
+	close(results)
 
-	return nil
+	if err != nil {
+		return nil, err
+	}
+
+	allSymbolizedLocs := make([]symbolizedLocation, 0, totalLocs)
+	for result := range results {
+		allSymbolizedLocs = append(allSymbolizedLocs, result.locations...)
+	}
+
+	return allSymbolizedLocs, nil
 }
 
 // groupLocationsByMapping groups locations by their mapping ID

pkg/symbolizer/symbolizer_test.go

@@ -245,6 +245,7 @@ func TestSymbolizePprof(t *testing.T) {
 				client:  mockClient,
 				bucket:  mockBucket,
 				metrics: newMetrics(nil),
+				cfg:     Config{MaxDebuginfodConcurrency: 1},
 			}
 
 			err := s.SymbolizePprof(context.Background(), tt.profile)
@@ -284,6 +285,7 @@ func TestSymbolizationKeepsSequentialFunctionIDs(t *testing.T) {
 		client:  mockClient,
 		bucket:  mockBucket,
 		metrics: newMetrics(nil),
+		cfg:     Config{MaxDebuginfodConcurrency: 1},
 	}
 
 	err := s.SymbolizePprof(context.Background(), profile)
@@ -320,6 +322,7 @@ func TestSymbolizationWithLidiaData(t *testing.T) {
 		client:  nil,
 		bucket:  mockBucket,
 		metrics: newMetrics(prometheus.NewRegistry()),
+		cfg:     Config{MaxDebuginfodConcurrency: 1},
 	}
 
 	req := &request{
@@ -364,6 +367,7 @@ func TestSymbolizeWithObjectStore(t *testing.T) {
 		client:  mockClient,
 		bucket:  mockBucket,
 		metrics: newMetrics(prometheus.NewRegistry()),
+		cfg:     Config{MaxDebuginfodConcurrency: 1},
 	}
 
 	elfTestFile := openTestFile(t)
@@ -621,3 +625,40 @@ func createRequest(t *testing.T, buildID string, address uint64) *request {
 		},
 	}
 }
+
+func TestConfigValidate(t *testing.T) {
+	tests := []struct {
+		name    string
+		setup   func(cfg *Config)
+		wantErr bool
+	}{
+		{
+			name:    "valid config with positive concurrency",
+			setup:   func(cfg *Config) { cfg.MaxDebuginfodConcurrency = 10 },
+			wantErr: false,
+		},
+		{
+			name:    "invalid config with zero concurrency",
+			setup:   func(cfg *Config) { cfg.MaxDebuginfodConcurrency = 0 },
+			wantErr: true,
+		},
+		{
+			name:    "invalid config with negative concurrency",
+			setup:   func(cfg *Config) { cfg.MaxDebuginfodConcurrency = -1 },
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cfg := Config{}
+			tt.setup(&cfg)
+			err := cfg.Validate()
+			if tt.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}