Add -tagroot and -tagleaf options (#649)

These add synthetic stack frames to samples.

Example usage:

$ pprof -tagroot thread pprof.proto

Will add synthetic stack frames at the root like "UIThread" and
"BackgroundPool-1".

The filename of the added frames will be the label key.

Closes #558

internal/driver/commands.go

@@ -189,6 +189,14 @@ var configHelp = map[string]string{
 		"Drops functions above the highest matched frame.",
 		"If set, all frames above the highest match are dropped from every sample.",
 		"Matching includes the function name, filename or object name."),
+	"tagroot": helpText(
+		"Adds pseudo stack frames for labels key/value pairs at the callstack root.",
+		"A comma-separated list of label keys.",
+		"The first key creates frames at the new root."),
+	"tagleaf": helpText(
+		"Adds pseudo stack frames for labels key/value pairs at the callstack leaf.",
+		"A comma-separated list of label keys.",
+		"The last key creates frames at the new leaf."),
 	"tagfocus": helpText(
 		"Restricts to samples with tags in range or matched by regexp",
 		"Use name=value syntax to limit the matching to a specific tag.",

internal/driver/config.go

@@ -30,6 +30,10 @@ type config struct {
 	Normalize           bool    `json:"normalize,omitempty"`
 	Sort                string  `json:"sort,omitempty"`
 
+	// Label pseudo stack frame generation options
+	TagRoot string `json:"tagroot,omitempty"`
+	TagLeaf string `json:"tagleaf,omitempty"`
+
 	// Filtering options
 	DropNegative bool    `json:"drop_negative,omitempty"`
 	NodeCount    int     `json:"nodecount,omitempty"`

internal/driver/driver.go

@@ -73,6 +73,10 @@ func generateRawReport(p *profile.Profile, cmd []string, cfg config, o *plugin.O
 
 	cfg = applyCommandOverrides(cmd[0], c.format, cfg)
 
+	// Create label pseudo nodes before filtering, in case the filters use
+	// the generated nodes.
+	generateTagRootsLeaves(p, cfg, o.UI)
+
 	// Delay focus after configuring report to get percentages on all samples.
 	relative := cfg.RelativePercentages
 	if relative {
@@ -208,6 +212,25 @@ func applyCommandOverrides(cmd string, outputFormat int, cfg config) config {
 	return cfg
 }
 
+// generateTagRootsLeaves generates extra nodes from the tagroot and tagleaf options.
+func generateTagRootsLeaves(prof *profile.Profile, cfg config, ui plugin.UI) {
+	tagRootLabelKeys := dropEmptyStrings(strings.Split(cfg.TagRoot, ","))
+	tagLeafLabelKeys := dropEmptyStrings(strings.Split(cfg.TagLeaf, ","))
+	rootm, leafm := addLabelNodes(prof, tagRootLabelKeys, tagLeafLabelKeys, cfg.Unit)
+	warnNoMatches(cfg.TagRoot == "" || rootm, "TagRoot", ui)
+	warnNoMatches(cfg.TagLeaf == "" || leafm, "TagLeaf", ui)
+}
+
+// dropEmptyStrings filters a slice to only non-empty strings
+func dropEmptyStrings(in []string) (out []string) {
+	for _, s := range in {
+		if s != "" {
+			out = append(out, s)
+		}
+	}
+	return
+}
+
 func aggregate(prof *profile.Profile, cfg config) error {
 	var function, filename, linenumber, address bool
 	inlines := !cfg.NoInlines

internal/driver/tagroot.go

@@ -0,0 +1,124 @@
+package driver
+
+import (
+	"strings"
+
+	"github.com/google/pprof/internal/measurement"
+	"github.com/google/pprof/profile"
+)
+
+// addLabelNodes adds pseudo stack frames "label:value" to each Sample with
+// labels matching the supplied keys.
+//
+// rootKeys adds frames at the root of the callgraph (first key becomes new root).
+// leafKeys adds frames at the leaf of the callgraph (last key becomes new leaf).
+//
+// Returns whether there were matches found for the label keys.
+func addLabelNodes(p *profile.Profile, rootKeys, leafKeys []string, outputUnit string) (rootm, leafm bool) {
+	// Find where to insert the new locations and functions at the end of
+	// their ID spaces.
+	var maxLocID uint64
+	var maxFunctionID uint64
+	for _, loc := range p.Location {
+		if loc.ID > maxLocID {
+			maxLocID = loc.ID
+		}
+	}
+	for _, f := range p.Function {
+		if f.ID > maxFunctionID {
+			maxFunctionID = f.ID
+		}
+	}
+	nextLocID := maxLocID + 1
+	nextFuncID := maxFunctionID + 1
+
+	// Intern the new locations and functions we are generating.
+	type locKey struct {
+		functionName, fileName string
+	}
+	locs := map[locKey]*profile.Location{}
+
+	internLoc := func(locKey locKey) *profile.Location {
+		loc, found := locs[locKey]
+		if found {
+			return loc
+		}
+
+		function := &profile.Function{
+			ID:       nextFuncID,
+			Name:     locKey.functionName,
+			Filename: locKey.fileName,
+		}
+		nextFuncID++
+		p.Function = append(p.Function, function)
+
+		loc = &profile.Location{
+			ID: nextLocID,
+			Line: []profile.Line{
+				{
+					Function: function,
+				},
+			},
+		}
+		nextLocID++
+		p.Location = append(p.Location, loc)
+		locs[locKey] = loc
+		return loc
+	}
+
+	makeLabelLocs := func(s *profile.Sample, keys []string) ([]*profile.Location, bool) {
+		var locs []*profile.Location
+		var match bool
+		for i := range keys {
+			// Loop backwards, ensuring the first tag is closest to the root,
+			// and the last tag is closest to the leaves.
+			k := keys[len(keys)-1-i]
+			values := formatLabelValues(s, k, outputUnit)
+			if len(values) > 0 {
+				match = true
+			}
+			locKey := locKey{
+				functionName: strings.Join(values, ","),
+				fileName:     k,
+			}
+			loc := internLoc(locKey)
+			locs = append(locs, loc)
+		}
+		return locs, match
+	}
+
+	for _, s := range p.Sample {
+		rootsToAdd, sampleMatchedRoot := makeLabelLocs(s, rootKeys)
+		if sampleMatchedRoot {
+			rootm = true
+		}
+		leavesToAdd, sampleMatchedLeaf := makeLabelLocs(s, leafKeys)
+		if sampleMatchedLeaf {
+			leafm = true
+		}
+
+		var newLocs []*profile.Location
+		newLocs = append(newLocs, leavesToAdd...)
+		newLocs = append(newLocs, s.Location...)
+		newLocs = append(newLocs, rootsToAdd...)
+		s.Location = newLocs
+	}
+	return
+}
+
+// formatLabelValues returns all the string and numeric labels in Sample, with
+// the numeric labels formatted according to outputUnit.
+func formatLabelValues(s *profile.Sample, k string, outputUnit string) []string {
+	var values []string
+	values = append(values, s.Label[k]...)
+	numLabels := s.NumLabel[k]
+	numUnits := s.NumUnit[k]
+	if len(numLabels) != len(numUnits) {
+		return values
+	}
+	for i, numLabel := range numLabels {
+		unit := numUnits[i]
+		values = append(values, measurement.ScaledLabel(numLabel, unit, outputUnit))
+	}
+	return values
+}