From ca93c310eec2a8cf83be8a5ee3cc0aa374e6d975 Mon Sep 17 00:00:00 2001
From: Tolya Korniltsev <korniltsev.anatoly@gmail.com>
Date: Tue, 6 Sep 2022 10:45:43 +0700
Subject: [PATCH] fix: pprof parser formatting for rbspy (#1454)

* fix: pprof parser formatting for rbspy

* revert some changes

* linter

* fix build

* pass formatter explicitly to parser constructor
---
 pkg/convert/pprof/format.go  | 40 +++++++++++++++++++++++++++++
 pkg/convert/pprof/parser.go  | 50 ++++++++++++++++++------------------
 pkg/convert/pprof/profile.go | 11 ++++----
 3 files changed, 71 insertions(+), 30 deletions(-)
 create mode 100644 pkg/convert/pprof/format.go

diff --git a/pkg/convert/pprof/format.go b/pkg/convert/pprof/format.go
new file mode 100644
index 0000000000..7c7effeae2
--- /dev/null
+++ b/pkg/convert/pprof/format.go
@@ -0,0 +1,40 @@
+package pprof
+
+import (
+	"fmt"
+	"github.com/pyroscope-io/pyroscope/pkg/storage/tree"
+	"reflect"
+	"unsafe"
+)
+
+type StackFrameFormatter interface {
+	format(x *tree.Profile, fn *tree.Function, line *tree.Line) []byte
+}
+
+func unsafeStrToSlice(s string) []byte {
+	return (*[0x7fff0000]byte)(unsafe.Pointer((*reflect.StringHeader)(unsafe.Pointer(&s)).Data))[:len(s):len(s)]
+}
+
+type UnsafeFunctionNameFormatter struct {
+}
+
+func (UnsafeFunctionNameFormatter) format(x *tree.Profile, fn *tree.Function, _ *tree.Line) []byte {
+	return unsafeStrToSlice(x.StringTable[fn.Name])
+}
+
+type RbspyFormatter struct {
+}
+
+func (RbspyFormatter) format(x *tree.Profile, fn *tree.Function, line *tree.Line) []byte {
+	return []byte(fmt.Sprintf("%s:%d - %s",
+		x.StringTable[fn.Filename],
+		line.Line,
+		x.StringTable[fn.Name]))
+}
+
+func StackFrameFormatterForSpyName(spyName string) StackFrameFormatter {
+	if spyName == "rbspy" {
+		return RbspyFormatter{}
+	}
+	return UnsafeFunctionNameFormatter{}
+}
diff --git a/pkg/convert/pprof/parser.go b/pkg/convert/pprof/parser.go
index 46ba2c075f..10804af934 100644
--- a/pkg/convert/pprof/parser.go
+++ b/pkg/convert/pprof/parser.go
@@ -3,43 +3,46 @@ package pprof
 import (
 	"context"
 	"fmt"
-	"io"
-	"reflect"
-	"time"
-	"unsafe"
-
 	"github.com/pyroscope-io/pyroscope/pkg/storage"
 	"github.com/pyroscope-io/pyroscope/pkg/storage/metadata"
 	"github.com/pyroscope-io/pyroscope/pkg/storage/segment"
 	"github.com/pyroscope-io/pyroscope/pkg/storage/tree"
+	"io"
+	"time"
 )
 
 type Parser struct {
-	putter        storage.Putter
-	spyName       string
-	labels        map[string]string
-	skipExemplars bool
-	sampleTypes   map[string]*tree.SampleTypeConfig
+	putter              storage.Putter
+	spyName             string
+	labels              map[string]string
+	skipExemplars       bool
+	sampleTypes         map[string]*tree.SampleTypeConfig
+	stackFrameFormatter StackFrameFormatter
 
 	cache             tree.LabelsCache
 	sampleTypesFilter func(string) bool
 }
 
 type ParserConfig struct {
-	Putter        storage.Putter
-	SpyName       string
-	Labels        map[string]string
-	SkipExemplars bool
-	SampleTypes   map[string]*tree.SampleTypeConfig
+	Putter              storage.Putter
+	SpyName             string
+	Labels              map[string]string
+	SkipExemplars       bool
+	SampleTypes         map[string]*tree.SampleTypeConfig
+	StackFrameFormatter StackFrameFormatter
 }
 
 func NewParser(config ParserConfig) *Parser {
+	if config.StackFrameFormatter == nil {
+		config.StackFrameFormatter = &UnsafeFunctionNameFormatter{}
+	}
 	return &Parser{
-		putter:        config.Putter,
-		spyName:       config.SpyName,
-		labels:        config.Labels,
-		sampleTypes:   config.SampleTypes,
-		skipExemplars: config.SkipExemplars,
+		putter:              config.Putter,
+		spyName:             config.SpyName,
+		labels:              config.Labels,
+		sampleTypes:         config.SampleTypes,
+		skipExemplars:       config.SkipExemplars,
+		stackFrameFormatter: config.StackFrameFormatter,
 
 		cache:             make(tree.LabelsCache),
 		sampleTypesFilter: filterKnownSamples(config.SampleTypes),
@@ -197,7 +200,8 @@ func (p *Parser) readTrees(x *tree.Profile, c tree.LabelsCache, f tree.Finder) {
 				if !ok || x.StringTable[fn.Name] == "" {
 					continue
 				}
-				stack = append(stack, unsafeStrToSlice(x.StringTable[fn.Name]))
+				sf := p.stackFrameFormatter.format(x, fn, loc.Line[j])
+				stack = append(stack, sf)
 			}
 		}
 		// Insert tree nodes.
@@ -221,10 +225,6 @@ func (p *Parser) readTrees(x *tree.Profile, c tree.LabelsCache, f tree.Finder) {
 	}
 }
 
-func unsafeStrToSlice(s string) []byte {
-	return (*[0x7fff0000]byte)(unsafe.Pointer((*reflect.StringHeader)(unsafe.Pointer(&s)).Data))[:len(s):len(s)]
-}
-
 func labelIndex(p *tree.Profile, labels tree.Labels, key string) int {
 	for i, label := range labels {
 		if n, ok := p.ResolveLabelName(label); ok && n == key {
diff --git a/pkg/convert/pprof/profile.go b/pkg/convert/pprof/profile.go
index 2ea75b3dad..ac3d368ec8 100644
--- a/pkg/convert/pprof/profile.go
+++ b/pkg/convert/pprof/profile.go
@@ -139,11 +139,12 @@ func (p *RawProfile) Parse(ctx context.Context, putter storage.Putter, _ storage
 			sampleTypes = p.SampleTypeConfig
 		}
 		p.parser = NewParser(ParserConfig{
-			SpyName:       md.SpyName,
-			Labels:        md.Key.Labels(),
-			Putter:        putter,
-			SampleTypes:   sampleTypes,
-			SkipExemplars: p.SkipExemplars,
+			SpyName:             md.SpyName,
+			Labels:              md.Key.Labels(),
+			Putter:              putter,
+			SampleTypes:         sampleTypes,
+			SkipExemplars:       p.SkipExemplars,
+			StackFrameFormatter: StackFrameFormatterForSpyName(md.SpyName),
 		})
 
 		if p.PreviousProfile != nil {