Refactor pprof profile traversal

Author: kolesnikovae

pkg/convert/parser_test.go

@@ -26,7 +26,7 @@ var _ = Describe("convert", func() {
 		It("parses data correctly", func() {
 			result := []string{}
 
-			b, err := ioutil.ReadFile("fixtures/cpu.pprof")
+			b, err := ioutil.ReadFile("testdata/cpu.pprof")
 			Expect(err).ToNot(HaveOccurred())
 			r := bytes.NewReader(b)
 			g, err := gzip.NewReader(r)

pkg/convert/profile-extensions.go

@@ -0,0 +1,73 @@
+package convert
+
+// These functions are kept separately as profile.pb.go is a generated file
+
+import (
+	"sort"
+
+	"github.com/valyala/bytebufferpool"
+)
+
+func (x *Profile) Get(sampleType string, cb func(name []byte, val int)) error {
+	valueIndex := 0
+	if sampleType != "" {
+		for i, v := range x.SampleType {
+			if x.StringTable[v.Type] == sampleType {
+				valueIndex = i
+				break
+			}
+		}
+	}
+
+	b := bytebufferpool.Get()
+	for _, s := range x.Sample {
+		for i := len(s.LocationId) - 1; i >= 0; i-- {
+			name, ok := x.findFunctionName(s.LocationId[i])
+			if !ok {
+				continue
+			}
+			if b.Len() > 0 {
+				_ = b.WriteByte(';')
+			}
+			_, _ = b.WriteString(name)
+		}
+		cb(b.Bytes(), int(s.Value[valueIndex]))
+		b.Reset()
+	}
+
+	bytebufferpool.Put(b)
+	return nil
+}
+
+func (x *Profile) findFunctionName(locId uint64) (string, bool) {
+	if loc, ok := x.findLocation(locId); ok {
+		if fn, ok := x.findFunction(loc.Line[0].FunctionId); ok {
+			return x.StringTable[fn.Name], true
+		}
+	}
+	return "", false
+}
+
+func (x *Profile) findLocation(lid uint64) (*Location, bool) {
+	idx := sort.Search(len(x.Location), func(i int) bool {
+		return x.Location[i].Id >= lid
+	})
+	if idx < len(x.Location) {
+		if l := x.Location[idx]; l.Id == lid {
+			return l, true
+		}
+	}
+	return nil, false
+}
+
+func (x *Profile) findFunction(fid uint64) (*Function, bool) {
+	idx := sort.Search(len(x.Function), func(i int) bool {
+		return x.Function[i].Id >= fid
+	})
+	if idx < len(x.Function) {
+		if f := x.Function[idx]; f.Id == fid {
+			return f, true
+		}
+	}
+	return nil, false
+}

pkg/convert/profile_extra.go

@@ -0,0 +1,22 @@
+package convert
+
+import (
+	"bytes"
+	"compress/gzip"
+	"io/ioutil"
+	"testing"
+)
+
+func BenchmarkProfile_Get(b *testing.B) {
+	buf, _ := ioutil.ReadFile("testdata/cpu.pprof")
+	g, _ := gzip.NewReader(bytes.NewReader(buf))
+	p, _ := ParsePprof(g)
+	noop := func(name []byte, val int) {}
+	b.ResetTimer()
+
+	b.Run("ByteBufferPool", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			_ = p.Get("samples", noop)
+		}
+	})
+}