Add -tagroot and -tagleaf options

These add synthetic stack frames to samples.

Example usage:

$ pprof -tagroot thread pprof.proto

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

Closes google#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,125 @@
+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.
+	locsByName := map[string]*profile.Location{}
+	functionsByName := map[string]*profile.Function{}
+
+	internFunction := func(functionName string) *profile.Function {
+		function, found := functionsByName[functionName]
+		if found {
+			return function
+		}
+		function = &profile.Function{
+			ID:   nextFuncID,
+			Name: functionName,
+		}
+		nextFuncID++
+		p.Function = append(p.Function, function)
+		functionsByName[functionName] = function
+		return function
+	}
+
+	internLocation := func(functionName string, function *profile.Function) *profile.Location {
+		location, found := locsByName[functionName]
+		if found {
+			return location
+		}
+		location = &profile.Location{
+			ID: nextLocID,
+			Line: []profile.Line{
+				{
+					Function: function,
+				},
+			},
+		}
+		nextLocID++
+		p.Location = append(p.Location, location)
+		locsByName[functionName] = location
+		return location
+	}
+
+	makeLabelLocations := func(s *profile.Sample, keys []string) ([]*profile.Location, bool) {
+		var locations []*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
+			}
+			functionName := k + ":" + strings.Join(values, ",")
+			location := internLocation(functionName, internFunction(functionName))
+			locations = append(locations, location)
+		}
+		return locations, match
+	}
+
+	for _, s := range p.Sample {
+		rootsToAdd, sampleMatchedRoot := makeLabelLocations(s, rootKeys)
+		if sampleMatchedRoot {
+			rootm = true
+		}
+		leavesToAdd, sampleMatchedLeaf := makeLabelLocations(s, leafKeys)
+		if sampleMatchedLeaf {
+			leafm = true
+		}
+
+		var newLocations []*profile.Location
+		newLocations = append(newLocations, leavesToAdd...)
+		newLocations = append(newLocations, s.Location...)
+		newLocations = append(newLocations, rootsToAdd...)
+		s.Location = newLocations
+	}
+	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
+}