feat: Support line numbers in pprof

fgprof.go

@@ -4,10 +4,26 @@
 package fgprof
 
 import (
+	"fmt"
 	"io"
 	"runtime"
+	"sort"
 	"strings"
 	"time"
+
+	"github.com/google/pprof/profile"
+)
+
+// Format decides how the output is rendered to the user.
+type Format string
+
+const (
+	// FormatFolded is used by Brendan Gregg's FlameGraph utility, see
+	// https://github.com/brendangregg/FlameGraph#2-fold-stacks.
+	FormatFolded Format = "folded"
+	// FormatPprof is used by Google's pprof utility, see
+	// https://github.com/google/pprof/blob/master/proto/README.md.
+	FormatPprof Format = "pprof"
 )
 
 // Start begins profiling the goroutines of the program and returns a function
@@ -23,7 +39,7 @@ func Start(w io.Writer, format Format) func() error {
 	stopCh := make(chan struct{})
 
 	prof := &profiler{}
-	stackCounts := stackCounter{}
+	profile := newWallclockProfile()
 
 	go func() {
 		defer ticker.Stop()
@@ -32,7 +48,7 @@ func Start(w io.Writer, format Format) func() error {
 			select {
 			case <-ticker.C:
 				stacks := prof.GoroutineProfile()
-				stackCounts.Update(stacks)
+				profile.Add(stacks)
 			case <-stopCh:
 				return
 			}
@@ -42,14 +58,8 @@ func Start(w io.Writer, format Format) func() error {
 	return func() error {
 		stopCh <- struct{}{}
 		endTime := time.Now()
-		return writeFormat(
-			w,
-			stackCounts.HumanMap(prof.SelfFrame()),
-			format,
-			hz,
-			startTime,
-			endTime,
-		)
+		profile.Ignore(prof.SelfFrames()...)
+		return profile.Export(w, format, hz, startTime, endTime)
 	}
 }
 
@@ -95,58 +105,203 @@ func (p *profiler) GoroutineProfile() []runtime.StackRecord {
 	}
 }
 
-func (p *profiler) SelfFrame() *runtime.Frame {
-	return p.selfFrame
+// SelfFrames returns frames that belong to the profiler so that we can ignore
+// them when exporting the final profile.
+func (p *profiler) SelfFrames() []*runtime.Frame {
+	if p.selfFrame != nil {
+		return []*runtime.Frame{p.selfFrame}
+	}
+	return nil
 }
 
-type stringStackCounter map[string]int
+func newWallclockProfile() *wallclockProfile {
+	return &wallclockProfile{stacks: map[[32]uintptr]*wallclockStack{}}
+}
 
-func (s stringStackCounter) Update(p []runtime.StackRecord) {
-	for _, pp := range p {
-		frames := runtime.CallersFrames(pp.Stack())
+// wallclockProfile holds a wallclock profile that can be exported in different
+// formats.
+type wallclockProfile struct {
+	stacks map[[32]uintptr]*wallclockStack
+	ignore []*runtime.Frame
+}
 
-		var stack []string
-		for {
-			frame, more := frames.Next()
-			stack = append([]string{frame.Function}, stack...)
-			if !more {
-				break
+// wallclockStack holds the symbolized frames of a stack trace and the number
+// of times it has been seen.
+type wallclockStack struct {
+	frames []*runtime.Frame
+	count  int
+}
+
+// Ignore sets a list of frames that should be ignored when exporting the
+// profile.
+func (p *wallclockProfile) Ignore(frames ...*runtime.Frame) {
+	p.ignore = frames
+}
+
+// Add adds the given stack traces to the profile.
+func (p *wallclockProfile) Add(stackRecords []runtime.StackRecord) {
+	for _, stackRecord := range stackRecords {
+		if _, ok := p.stacks[stackRecord.Stack0]; !ok {
+			ws := &wallclockStack{}
+			// symbolize pcs into frames
+			frames := runtime.CallersFrames(stackRecord.Stack())
+			for {
+				frame, more := frames.Next()
+				ws.frames = append(ws.frames, &frame)
+				if !more {
+					break
+				}
 			}
+			p.stacks[stackRecord.Stack0] = ws
 		}
-		key := strings.Join(stack, ";")
-		s[key]++
+		p.stacks[stackRecord.Stack0].count++
 	}
 }
 
-type stackCounter map[[32]uintptr]int
+func (p *wallclockProfile) Export(w io.Writer, f Format, hz int, startTime, endTime time.Time) error {
+	switch f {
+	case FormatFolded:
+		return p.exportFolded(w)
+	case FormatPprof:
+		return p.exportPprof(hz, startTime, endTime).Write(w)
+	default:
+		return fmt.Errorf("unknown format: %q", f)
+	}
+}
 
-func (s stackCounter) Update(p []runtime.StackRecord) {
-	for _, pp := range p {
-		s[pp.Stack0]++
+// exportStacks returns the stacks in this profile except those that have been
+// set to Ignore().
+func (p *wallclockProfile) exportStacks() []*wallclockStack {
+	stacks := make([]*wallclockStack, 0, len(p.stacks))
+nextStack:
+	for _, ws := range p.stacks {
+		for _, f := range ws.frames {
+			for _, igf := range p.ignore {
+				if f.Entry == igf.Entry {
+					continue nextStack
+				}
+			}
+		}
+		stacks = append(stacks, ws)
 	}
+	return stacks
 }
 
-// @TODO(fg) create a better interface that avoids the pprof output having to
-// split the stacks using the `;` separator.
-func (s stackCounter) HumanMap(exclude *runtime.Frame) map[string]int {
-	m := map[string]int{}
+func (p *wallclockProfile) exportFolded(w io.Writer) error {
+	var lines []string
+	stacks := p.exportStacks()
+	for _, ws := range stacks {
+		var foldedStack []string
+		for _, f := range ws.frames {
+			foldedStack = append(foldedStack, f.Function)
+		}
+		line := fmt.Sprintf("%s %d", strings.Join(foldedStack, ";"), ws.count)
+		lines = append(lines, line)
+	}
+	sort.Strings(lines)
+	_, err := io.WriteString(w, strings.Join(lines, "\n")+"\n")
+	return err
+}
+
+func (p *wallclockProfile) exportPprof(hz int, startTime, endTime time.Time) *profile.Profile {
+	prof := &profile.Profile{}
+	m := &profile.Mapping{ID: 1, HasFunctions: true}
+	prof.Period = int64(1e9 / hz) // Number of nanoseconds between samples.
+	prof.TimeNanos = startTime.UnixNano()
+	prof.DurationNanos = int64(endTime.Sub(startTime))
+	prof.Mapping = []*profile.Mapping{m}
+	prof.SampleType = []*profile.ValueType{
+		{
+			Type: "samples",
+			Unit: "count",
+		},
+		{
+			Type: "time",
+			Unit: "nanoseconds",
+		},
+	}
+	prof.PeriodType = &profile.ValueType{
+		Type: "wallclock",
+		Unit: "nanoseconds",
+	}
+
+	type functionKey struct {
+		Name     string
+		Filename string
+	}
+	funcIdx := map[functionKey]*profile.Function{}
+
+	type locationKey struct {
+		Function functionKey
+		Line     int
+	}
+	locationIdx := map[locationKey]*profile.Location{}
+	for _, ws := range p.exportStacks() {
+		sample := &profile.Sample{
+			Value: []int64{
+				int64(ws.count),
+				int64(1000 * 1000 * 1000 / hz * ws.count),
+			},
+		}
+
+		for _, frame := range ws.frames {
+			fnKey := functionKey{Name: frame.Function, Filename: frame.File}
+			function, ok := funcIdx[fnKey]
+			if !ok {
+				function = &profile.Function{
+					ID:         uint64(len(prof.Function)) + 1,
+					Name:       frame.Function,
+					SystemName: frame.Function,
+					Filename:   frame.File,
+				}
+				funcIdx[fnKey] = function
+				prof.Function = append(prof.Function, function)
+			}
+
+			locKey := locationKey{Function: fnKey, Line: frame.Line}
+			location, ok := locationIdx[locKey]
+			if !ok {
+				location = &profile.Location{
+					ID:      uint64(len(prof.Location)) + 1,
+					Mapping: m,
+					Line: []profile.Line{{
+						Function: function,
+						Line:     int64(frame.Line),
+					}},
+				}
+				locationIdx[locKey] = location
+				prof.Location = append(prof.Location, location)
+			}
+			sample.Location = append(sample.Location, location)
+		}
+		prof.Sample = append(prof.Sample, sample)
+	}
+	return prof
+}
+
+type symbolizedStacks map[[32]uintptr][]frameCount
+
+func (w wallclockProfile) Symbolize(exclude *runtime.Frame) symbolizedStacks {
+	m := make(symbolizedStacks)
 outer:
-	for stack0, count := range s {
+	for stack0, ws := range w.stacks {
 		frames := runtime.CallersFrames((&runtime.StackRecord{Stack0: stack0}).Stack())
 
-		var stack []string
 		for {
 			frame, more := frames.Next()
 			if frame.Entry == exclude.Entry {
 				continue outer
 			}
-			stack = append([]string{frame.Function}, stack...)
+			m[stack0] = append(m[stack0], frameCount{Frame: &frame, Count: ws.count})
 			if !more {
 				break
 			}
 		}
-		key := strings.Join(stack, ";")
-		m[key] = count
 	}
 	return m
 }
+
+type frameCount struct {
+	*runtime.Frame
+	Count int
+}