Add symbolization inside the read path

Author: marcsanmi

cmd/symbolization/main.go

@@ -7,7 +7,7 @@ import (
 
 	pprof "github.com/google/pprof/profile"
 
-	"github.com/grafana/pyroscope/pkg/experiment/symbolization"
+	"github.com/grafana/pyroscope/pkg/experiment/symbolizer"
 )
 
 const (
@@ -16,12 +16,12 @@ const (
 )
 
 func main() {
-	client := symbolization.NewDebuginfodClient(debuginfodBaseURL)
+	client := symbolizer.NewDebuginfodClient(debuginfodBaseURL)
 
 	// Alternatively, use a local debug info file:
 	//client := &localDebuginfodClient{debugFilePath: "/path/to/your/debug/file"}
 
-	symbolizer := symbolization.NewSymbolizer(client)
+	s := symbolizer.NewSymbolizer(client)
 	ctx := context.Background()
 
 	_, err := client.FetchDebuginfo(buildID)
@@ -31,11 +31,11 @@ func main() {
 	//defer os.Remove(debugFilePath)
 
 	// Create a request to symbolize specific addresses
-	req := symbolization.Request{
+	req := symbolizer.Request{
 		BuildID: buildID,
-		Mappings: []symbolization.RequestMapping{
+		Mappings: []symbolizer.RequestMapping{
 			{
-				Locations: []*symbolization.Location{
+				Locations: []*symbolizer.Location{
 					{
 						Address: 0x1500,
 						Mapping: &pprof.Mapping{},
@@ -53,7 +53,7 @@ func main() {
 		},
 	}
 
-	if err := symbolizer.Symbolize(ctx, req); err != nil {
+	if err := s.Symbolize(ctx, req); err != nil {
 		log.Fatalf("Failed to symbolize: %v", err)
 	}
 

pkg/experiment/query_backend/backend.go

@@ -14,16 +14,19 @@ import (
 
 	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
 	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
+	"github.com/grafana/pyroscope/pkg/experiment/symbolizer"
 	"github.com/grafana/pyroscope/pkg/util"
 )
 
 type Config struct {
 	Address          string            `yaml:"address"`
 	GRPCClientConfig grpcclient.Config `yaml:"grpc_client_config" doc:"description=Configures the gRPC client used to communicate between the query-frontends and the query-schedulers."`
+	DebuginfodURL    string            `yaml:"debuginfod_url"`
 }
 
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.StringVar(&cfg.Address, "query-backend.address", "localhost:9095", "")
+	f.StringVar(&cfg.DebuginfodURL, "query-backend.debuginfod-url", "https://debuginfod.elfutils.org", "URL of the debuginfod server")
 	cfg.GRPCClientConfig.RegisterFlagsWithPrefix("query-backend.grpc-client-config", f)
 }
 
@@ -48,6 +51,8 @@ type QueryBackend struct {
 
 	backendClient QueryHandler
 	blockReader   QueryHandler
+
+	symbolizer *symbolizer.Symbolizer
 }
 
 func New(
@@ -57,13 +62,27 @@ func New(
 	backendClient QueryHandler,
 	blockReader QueryHandler,
 ) (*QueryBackend, error) {
+	var sym *symbolizer.Symbolizer
+	if config.DebuginfodURL != "" {
+		sym = symbolizer.NewSymbolizer(
+			symbolizer.NewDebuginfodClient(config.DebuginfodURL),
+		)
+	}
+
 	q := QueryBackend{
 		config:        config,
 		logger:        logger,
 		reg:           reg,
 		backendClient: backendClient,
 		blockReader:   blockReader,
+		symbolizer:    sym,
 	}
+
+	// Pass symbolizer to BlockReader if it's the right type
+	if br, ok := blockReader.(*BlockReader); ok {
+		br.symbolizer = sym
+	}
+
 	q.service = services.NewIdleService(q.starting, q.stopping)
 	return &q, nil
 }

pkg/experiment/query_backend/query_tree.go

@@ -56,7 +56,9 @@ func queryTree(q *queryContext, query *queryv1.Query) (*queryv1.Report, error) {
 	defer runutil.CloseWithErrCapture(&err, profiles, "failed to close profile stream")
 
 	resolver := symdb.NewResolver(q.ctx, q.ds.Symbols(),
-		symdb.WithResolverMaxNodes(query.Tree.GetMaxNodes()))
+		symdb.WithResolverMaxNodes(query.Tree.GetMaxNodes()),
+		symdb.WithSymbolizer(q.symbolizer))
+
 	defer resolver.Release()
 
 	if len(spanSelector) > 0 {

pkg/experiment/symbolization/addrmapper.go renamed to pkg/experiment/symbolizer/addrmapper.go

@@ -1,4 +1,4 @@
-package symbolization
+package symbolizer
 
 import (
 	"debug/elf"

pkg/experiment/symbolization/debuginfod_client.go renamed to pkg/experiment/symbolizer/debuginfod_client.go

@@ -1,11 +1,12 @@
-package symbolization
+package symbolizer
 
 import (
 	"fmt"
 	"io"
 	"net/http"
 	"os"
 	"path/filepath"
+	"regexp"
 )
 
 type DebuginfodClient interface {
@@ -24,7 +25,12 @@ func NewDebuginfodClient(baseURL string) DebuginfodClient {
 
 // FetchDebuginfo fetches the debuginfo file for a specific build ID.
 func (c *debuginfodClient) FetchDebuginfo(buildID string) (string, error) {
-	url := fmt.Sprintf("%s/buildid/%s/debuginfo", c.baseURL, buildID)
+	sanitizedBuildID, err := sanitizeBuildID(buildID)
+	if err != nil {
+		return "", err
+	}
+
+	url := fmt.Sprintf("%s/buildid/%s/debuginfo", c.baseURL, sanitizedBuildID)
 
 	resp, err := http.Get(url)
 	if err != nil {
@@ -36,9 +42,10 @@ func (c *debuginfodClient) FetchDebuginfo(buildID string) (string, error) {
 		return "", fmt.Errorf("unexpected HTTP status: %s", resp.Status)
 	}
 
+	// TODO: Avoid file operations and handle debuginfo in memory.
 	// Save the debuginfo to a temporary file
 	tempDir := os.TempDir()
-	filePath := filepath.Join(tempDir, fmt.Sprintf("%s.elf", buildID))
+	filePath := filepath.Join(tempDir, fmt.Sprintf("%s.elf", sanitizedBuildID))
 	outFile, err := os.Create(filePath)
 	if err != nil {
 		return "", fmt.Errorf("failed to create temp file: %w", err)
@@ -52,3 +59,13 @@ func (c *debuginfodClient) FetchDebuginfo(buildID string) (string, error) {
 
 	return filePath, nil
 }
+
+// sanitizeBuildID ensures that the buildID is a safe and valid string for use in file paths.
+func sanitizeBuildID(buildID string) (string, error) {
+	// Allow only alphanumeric characters, dashes, and underscores.
+	validBuildID := regexp.MustCompile(`^[a-zA-Z0-9_-]+$`)
+	if !validBuildID.MatchString(buildID) {
+		return "", fmt.Errorf("invalid build ID: %s", buildID)
+	}
+	return buildID, nil
+}

pkg/experiment/symbolization/dwarf.go renamed to pkg/experiment/symbolizer/dwarf.go

@@ -1,4 +1,4 @@
-package symbolization
+package symbolizer
 
 import (
 	"context"

pkg/experiment/symbolization/symbolizer.go renamed to pkg/experiment/symbolizer/symbolizer.go

@@ -1,4 +1,4 @@
-package symbolization
+package symbolizer
 
 import (
 	"context"
@@ -92,7 +92,7 @@ func (s *Symbolizer) Symbolize(ctx context.Context, req Request) error {
 				continue // Skip errors for individual addresses
 			}
 
-			// Update the location directly (this is why Parca modifies the request - it's updating the shared locations)
+			// Update the location directly
 			loc.Lines = lines
 		}
 	}

pkg/experiment/symbolization/types.go renamed to pkg/experiment/symbolizer/types.go

@@ -1,4 +1,4 @@
-package symbolization
+package symbolizer
 
 import (
 	"context"

pkg/phlaredb/symdb/resolver.go

@@ -11,6 +11,7 @@ import (
 
 	googlev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
 	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
+	"github.com/grafana/pyroscope/pkg/experiment/symbolizer"
 	"github.com/grafana/pyroscope/pkg/model"
 	schemav1 "github.com/grafana/pyroscope/pkg/phlaredb/schemas/v1"
 	"github.com/grafana/pyroscope/pkg/pprof"
@@ -37,6 +38,8 @@ type Resolver struct {
 
 	maxNodes int64
 	sts      *typesv1.StackTraceSelector
+
+	symbolizer *symbolizer.Symbolizer
 }
 
 type ResolverOption func(*Resolver)
@@ -57,6 +60,12 @@ func WithResolverMaxNodes(n int64) ResolverOption {
 	}
 }
 
+func WithSymbolizer(s *symbolizer.Symbolizer) ResolverOption {
+	return func(r *Resolver) {
+		r.symbolizer = s
+	}
+}
+
 // WithResolverStackTraceSelector specifies the stack trace selector.
 // Only stack traces that belong to the callSite (have the prefix provided)
 // will be selected. If empty, the filter is ignored.
@@ -273,7 +282,9 @@ func (r *Resolver) withSymbols(ctx context.Context, fn func(*Symbols, *SampleApp
 			if err := p.fetch(ctx); err != nil {
 				return err
 			}
-			return fn(p.reader.Symbols(), p.samples)
+			symbols := p.reader.Symbols()
+			symbols.SetSymbolizer(r.symbolizer)
+			return fn(symbols, p.samples)
 		}))
 	}
 	return g.Wait()
@@ -295,3 +306,18 @@ func (r *Symbols) Tree(
 ) (*model.Tree, error) {
 	return buildTree(ctx, r, appender, maxNodes)
 }
+
+func (r *Symbols) SetSymbolizer(sym *symbolizer.Symbolizer) {
+	r.Symbolizer = sym
+}
+
+func (r *Symbols) needsDebuginfodSymbolization(loc *schemav1.InMemoryLocation, mapping *schemav1.InMemoryMapping) bool {
+	if r.Symbolizer == nil {
+		return false
+	}
+	if len(loc.Line) == 0 {
+		// Must have mapping with build ID
+		return mapping != nil && mapping.BuildId != 0
+	}
+	return false
+}

pkg/phlaredb/symdb/resolver_tree.go

@@ -4,8 +4,10 @@ import (
 	"context"
 	"sync"
 
+	pprof "github.com/google/pprof/profile"
 	"golang.org/x/sync/errgroup"
 
+	"github.com/grafana/pyroscope/pkg/experiment/symbolizer"
 	"github.com/grafana/pyroscope/pkg/iter"
 	"github.com/grafana/pyroscope/pkg/model"
 	schemav1 "github.com/grafana/pyroscope/pkg/phlaredb/schemas/v1"
@@ -19,6 +21,13 @@ func buildTree(
 	appender *SampleAppender,
 	maxNodes int64,
 ) (*model.Tree, error) {
+	// Try debuginfod symbolization first
+	if symbols != nil && symbols.Symbolizer != nil {
+		if err := symbolizeLocations(ctx, symbols); err != nil {
+			// TODO: Log error but continue? partial symbolization is better than none
+		}
+	}
+
 	// If the number of samples is large (> 128K) and the StacktraceResolver
 	// implements the range iterator, we will be building the tree based on
 	// the parent pointer tree of the partition (a copy of). The only exception
@@ -239,3 +248,85 @@ func minValue(nodes []Node, maxNodes int64) int64 {
 	}
 	return h[0]
 }
+
+func symbolizeLocations(ctx context.Context, symbols *Symbols) error {
+	type locToSymbolize struct {
+		idx     int32
+		loc     *schemav1.InMemoryLocation
+		mapping *schemav1.InMemoryMapping
+	}
+	locsByBuildId := make(map[string][]locToSymbolize)
+
+	// Find all locations needing symbolization
+	for i, loc := range symbols.Locations {
+		if mapping := &symbols.Mappings[loc.MappingId]; symbols.needsDebuginfodSymbolization(&loc, mapping) {
+			buildIDStr := symbols.Strings[mapping.BuildId]
+			locsByBuildId[buildIDStr] = append(locsByBuildId[buildIDStr], locToSymbolize{
+				idx:     int32(i),
+				loc:     &loc,
+				mapping: mapping,
+			})
+		}
+	}
+
+	for buildID, locs := range locsByBuildId {
+		req := symbolizer.Request{
+			BuildID: buildID,
+			Mappings: []symbolizer.RequestMapping{{
+				Locations: make([]*symbolizer.Location, len(locs)),
+			}},
+		}
+
+		for i, loc := range locs {
+			req.Mappings[0].Locations[i] = &symbolizer.Location{
+				Address: loc.loc.Address,
+				Mapping: &pprof.Mapping{
+					Start:   loc.mapping.MemoryStart,
+					Limit:   loc.mapping.MemoryLimit,
+					Offset:  loc.mapping.FileOffset,
+					BuildID: buildID,
+				},
+			}
+		}
+
+		if err := symbols.Symbolizer.Symbolize(ctx, req); err != nil {
+			// TODO: log/process errors but continue with other build IDs
+			continue
+		}
+
+		// Store symbolization results back
+		for i, symLoc := range req.Mappings[0].Locations {
+			if len(symLoc.Lines) > 0 {
+				// Get the original location we're updating
+				locIdx := locs[i].idx
+
+				// Clear the existing lines for the location
+				symbols.Locations[locIdx].Line = nil
+
+				for _, line := range symLoc.Lines {
+					// Create string entries first
+					nameIdx := uint32(len(symbols.Strings))
+					symbols.Strings = append(symbols.Strings, line.Function.Name)
+
+					filenameIdx := uint32(len(symbols.Strings))
+					symbols.Strings = append(symbols.Strings, line.Function.Filename)
+
+					// Create function entry
+					funcId := uint32(len(symbols.Functions))
+					symbols.Functions = append(symbols.Functions, schemav1.InMemoryFunction{
+						Id:        uint64(funcId),
+						Name:      nameIdx,
+						Filename:  filenameIdx,
+						StartLine: uint32(line.Function.StartLine),
+					})
+
+					symbols.Locations[locIdx].Line = append(symbols.Locations[locIdx].Line, schemav1.InMemoryLine{
+						FunctionId: funcId,
+						Line:       int32(line.Line),
+					})
+				}
+			}
+		}
+	}
+	return nil
+}