Repository: google/pprof (go)

Issue:
Support visualizing tags as nodes
The `pprof` CLI has some support for visualizing profiles with tags (the `-tags` output format allows one to see all tags in the profile, [tag filtering](https://github.com/google/pprof/tree/master/doc#tag-filtering) allows one to view only or exclude portions of the profile without tags, and the graph view(s) show some tags as nodeletes), but it's not always possible to visualize which tags are associated with which stacks when using `pprof`.

On a few occasions now, internal users have hacked `pprof` to show certain tags as nodes at the root of the call stack, so it's something that has been useful. For example, if samples are tagged such that they can be associated with particular chunks of work (for example, work done for a particular URL endpoint), it can be useful to add a pseudo-node for a particular tag toward the root of the call stack to group stacks associated with a particular chunk of work in the flame graph view.

On possibility for doing this would be to introduce the following options to `pprof`:
* `-tagroot`: takes in a comma-separated list of label keys and adds the associated key/value pairs  as pseudo-nodes on the stack for samples which have the listed labels. For each sample and label key, if the sample has a label with the specified key, a pseudo-node will be added. Pseudo-nodes will be added at the top of the call stack (closest to the root), with pseudo-node for the label that is first in the list being highest in the callstack and the pseudo-node for the label that is last in the list appearing just before frames in the original profile.
* `-tagleaf`: takes in a comma-separated list of label keys and adds the associated labels as pseudo-nodes on the stack. For each sample and label key, if the sample has a label with the specified key, a pseudo-node will be added. Pseudo-nodes will be added at the bottom of the call stack (after leaf nodes), with pseudo-node for the label that is first in the list being the child of what was the leaf in the original profile and the pseudo-node for the label that is last in the list appearing as the new leaf node.

A key-value pair with key “<key>” and value “value” could be represented as a Function with a function name “<key>:<value>”, and a filename of “<key>”.

Pull request:
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 #558

Patch:
diff --git a/internal/driver/commands.go b/internal/driver/commands.go
index 4397e25..c9edf10 100644
--- a/internal/driver/commands.go
+++ b/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.",
diff --git a/internal/driver/config.go b/internal/driver/config.go
index b3f82f2..9fcdd45 100644
--- a/internal/driver/config.go
+++ b/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"`
diff --git a/internal/driver/driver.go b/internal/driver/driver.go
index 3967a12..6a1e64c 100644
--- a/internal/driver/driver.go
+++ b/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
diff --git a/internal/driver/tagroot.go b/internal/driver/tagroot.go
new file mode 100644
index 0000000..c2cdfa4
--- /dev/null
+++ b/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
+}


Test patch:
diff --git a/internal/driver/tagroot_test.go b/internal/driver/tagroot_test.go
new file mode 100644
index 0000000..d70a901
--- /dev/null
+++ b/internal/driver/tagroot_test.go
@@ -0,0 +1,368 @@
+package driver
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/google/pprof/internal/proftest"
+	"github.com/google/pprof/profile"
+)
+
+const mainBinary = "/bin/main"
+
+var cpuF = []*profile.Function{
+	{ID: 1, Name: "main", SystemName: "main", Filename: "main.c"},
+	{ID: 2, Name: "foo", SystemName: "foo", Filename: "foo.c"},
+	{ID: 3, Name: "foo_caller", SystemName: "foo_caller", Filename: "foo.c"},
+}
+
+var cpuM = []*profile.Mapping{
+	{
+		ID:              1,
+		Start:           0x10000,
+		Limit:           0x40000,
+		File:            mainBinary,
+		HasFunctions:    true,
+		HasFilenames:    true,
+		HasLineNumbers:  true,
+		HasInlineFrames: true,
+	},
+	{
+		ID:              2,
+		Start:           0x1000,
+		Limit:           0x4000,
+		File:            "/lib/lib.so",
+		HasFunctions:    true,
+		HasFilenames:    true,
+		HasLineNumbers:  true,
+		HasInlineFrames: true,
+	},
+	{
+		ID:              3,
+		Start:           0x4000,
+		Limit:           0x5000,
+		File:            "/lib/lib2_c.so.6",
+		HasFunctions:    true,
+		HasFilenames:    true,
+		HasLineNumbers:  true,
+		HasInlineFrames: true,
+	},
+	{
+		ID:              4,
+		Start:           0x5000,
+		Limit:           0x9000,
+		File:            "/lib/lib.so_6 (deleted)",
+		HasFunctions:    true,
+		HasFilenames:    true,
+		HasLineNumbers:  true,
+		HasInlineFrames: true,
+	},
+}
+
+var cpuL = []*profile.Location{
+	{
+		ID:      1000,
+		Mapping: cpuM[1],
+		Address: 0x1000,
+		Line: []profile.Line{
+			{Function: cpuF[0], Line: 1},
+		},
+	},
+	{
+		ID:      2000,
+		Mapping: cpuM[0],
+		Address: 0x2000,
+		Line: []profile.Line{
+			{Function: cpuF[1], Line: 2},
+			{Function: cpuF[2], Line: 1},
+		},
+	},
+	{
+		ID:      3000,
+		Mapping: cpuM[0],
+		Address: 0x3000,
+		Line: []profile.Line{
+			{Function: cpuF[1], Line: 2},
+			{Function: cpuF[2], Line: 1},
+		},
+	},
+	{
+		ID:      3001,
+		Mapping: cpuM[0],
+		Address: 0x3001,
+		Line: []profile.Line{
+			{Function: cpuF[2], Line: 2},
+		},
+	},
+	{
+		ID:      3002,
+		Mapping: cpuM[0],
+		Address: 0x3002,
+		Line: []profile.Line{
+			{Function: cpuF[2], Line: 3},
+		},
+	},
+}
+
+var testProfile1 = &profile.Profile{
+	TimeNanos:     10000,
+	PeriodType:    &profile.ValueType{Type: "cpu", Unit: "milliseconds"},
+	Period:        1,
+	DurationNanos: 10e9,
+	SampleType: []*profile.ValueType{
+		{Type: "samples", Unit: "count"},
+		{Type: "cpu", Unit: "milliseconds"},
+	},
+	Sample: []*profile.Sample{
+		{
+			Location: []*profile.Location{cpuL[0]},
+			Value:    []int64{1000, 1000},
+			Label: map[string][]string{
+				"key1": {"tag1"},
+				"key2": {"tag1"},
+			},
+		},
+		{
+			Location: []*profile.Location{cpuL[1], cpuL[0]},
+			Value:    []int64{100, 100},
+			Label: map[string][]string{
+				"key1": {"tag2"},
+				"key3": {"tag2"},
+			},
+		},
+		{
+			Location: []*profile.Location{cpuL[2], cpuL[0]},
+			Value:    []int64{10, 10},
+			Label: map[string][]string{
+				"key1": {"tag3"},
+				"key2": {"tag2"},
+			},
+			NumLabel: map[string][]int64{
+				"allocations": {1024},
+			},
+			NumUnit: map[string][]string{
+				"allocations": {""},
+			},
+		},
+		{
+			Location: []*profile.Location{cpuL[3], cpuL[0]},
+			Value:    []int64{10000, 10000},
+			Label: map[string][]string{
+				"key1": {"tag4"},
+				"key2": {"tag1"},
+			},
+			NumLabel: map[string][]int64{
+				"allocations": {1024, 2048},
+			},
+			NumUnit: map[string][]string{
+				"allocations": {"bytes", "b"},
+			},
+		},
+		{
+			Location: []*profile.Location{cpuL[4], cpuL[0]},
+			Value:    []int64{1, 1},
+			Label: map[string][]string{
+				"key1": {"tag4"},
+				"key2": {"tag1", "tag5"},
+			},
+			NumLabel: map[string][]int64{
+				"allocations": {1024, 1},
+			},
+			NumUnit: map[string][]string{
+				"allocations": {"byte", "kilobyte"},
+			},
+		},
+	},
+	Location: cpuL,
+	Function: cpuF,
+	Mapping:  cpuM,
+}
+
+func TestAddLabelNodesMatchBooleans(t *testing.T) {
+	type addLabelNodesTestcase struct {
+		name             string
+		tagroot, tagleaf []string
+		outputUnit       string
+		rootm, leafm     bool
+		// wantSampleFuncs contains expected stack functions and sample value after
+		// adding nodes, in the same order as in the profile. The format is as
+		// returned by stackCollapse function, which is "callee caller: <num>".
+		wantSampleFuncs []string
+	}
+	for _, tc := range []addLabelNodesTestcase{
+		{
+			name: "Without tagroot or tagleaf, add no extra nodes, and should not match",
+			wantSampleFuncs: []string{
+				"main(main.c) 1000",
+				"main(main.c);foo(foo.c);foo_caller(foo.c) 100",
+				"main(main.c);foo(foo.c);foo_caller(foo.c) 10",
+				"main(main.c);foo_caller(foo.c) 10000",
+				"main(main.c);foo_caller(foo.c) 1",
+			},
+		},
+		{
+			name:    "Keys that aren't found add empty nodes, and should not match",
+			tagroot: []string{"key404"},
+			tagleaf: []string{"key404"},
+			wantSampleFuncs: []string{
+				"(key404);main(main.c);(key404) 1000",
+				"(key404);main(main.c);foo(foo.c);foo_caller(foo.c);(key404) 100",
+				"(key404);main(main.c);foo(foo.c);foo_caller(foo.c);(key404) 10",
+				"(key404);main(main.c);foo_caller(foo.c);(key404) 10000",
+				"(key404);main(main.c);foo_caller(foo.c);(key404) 1",
+			},
+		},
+		{
+			name:    "tagroot adds nodes for key1 and reports a match",
+			tagroot: []string{"key1"},
+			rootm:   true,
+			wantSampleFuncs: []string{
+				"tag1(key1);main(main.c) 1000",
+				"tag2(key1);main(main.c);foo(foo.c);foo_caller(foo.c) 100",
+				"tag3(key1);main(main.c);foo(foo.c);foo_caller(foo.c) 10",
+				"tag4(key1);main(main.c);foo_caller(foo.c) 10000",
+				"tag4(key1);main(main.c);foo_caller(foo.c) 1",
+			},
+		},
+		{
+			name:    "tagroot adds nodes for key2 and reports a match",
+			tagroot: []string{"key2"},
+			rootm:   true,
+			wantSampleFuncs: []string{
+				"tag1(key2);main(main.c) 1000",
+				"(key2);main(main.c);foo(foo.c);foo_caller(foo.c) 100",
+				"tag2(key2);main(main.c);foo(foo.c);foo_caller(foo.c) 10",
+				"tag1(key2);main(main.c);foo_caller(foo.c) 10000",
+				"tag1,tag5(key2);main(main.c);foo_caller(foo.c) 1",
+			},
+		},
+		{
+			name:    "tagleaf adds nodes for key1 and reports a match",
+			tagleaf: []string{"key1"},
+			leafm:   true,
+			wantSampleFuncs: []string{
+				"main(main.c);tag1(key1) 1000",
+				"main(main.c);foo(foo.c);foo_caller(foo.c);tag2(key1) 100",
+				"main(main.c);foo(foo.c);foo_caller(foo.c);tag3(key1) 10",
+				"main(main.c);foo_caller(foo.c);tag4(key1) 10000",
+				"main(main.c);foo_caller(foo.c);tag4(key1) 1",
+			},
+		},
+		{
+			name:    "tagleaf adds nodes for key3 and reports a match",
+			tagleaf: []string{"key3"},
+			leafm:   true,
+			wantSampleFuncs: []string{
+				"main(main.c);(key3) 1000",
+				"main(main.c);foo(foo.c);foo_caller(foo.c);tag2(key3) 100",
+				"main(main.c);foo(foo.c);foo_caller(foo.c);(key3) 10",
+				"main(main.c);foo_caller(foo.c);(key3) 10000",
+				"main(main.c);foo_caller(foo.c);(key3) 1",
+			},
+		},
+		{
+			name:    "tagroot adds nodes for key1,key2 in order and reports a match",
+			tagroot: []string{"key1", "key2"},
+			rootm:   true,
+			wantSampleFuncs: []string{
+				"tag1(key1);tag1(key2);main(main.c) 1000",
+				"tag2(key1);(key2);main(main.c);foo(foo.c);foo_caller(foo.c) 100",
+				"tag3(key1);tag2(key2);main(main.c);foo(foo.c);foo_caller(foo.c) 10",
+				"tag4(key1);tag1(key2);main(main.c);foo_caller(foo.c) 10000",
+				"tag4(key1);tag1,tag5(key2);main(main.c);foo_caller(foo.c) 1",
+			},
+		},
+		{
+			name:    "tagleaf adds nodes for key1,key2 in order and reports a match",
+			tagleaf: []string{"key1", "key2"},
+			leafm:   true,
+			wantSampleFuncs: []string{
+				"main(main.c);tag1(key1);tag1(key2) 1000",
+				"main(main.c);foo(foo.c);foo_caller(foo.c);tag2(key1);(key2) 100",
+				"main(main.c);foo(foo.c);foo_caller(foo.c);tag3(key1);tag2(key2) 10",
+				"main(main.c);foo_caller(foo.c);tag4(key1);tag1(key2) 10000",
+				"main(main.c);foo_caller(foo.c);tag4(key1);tag1,tag5(key2) 1",
+			},
+		},
+		{
+			name:    "Numeric units are added with units with tagleaf",
+			tagleaf: []string{"allocations"},
+			leafm:   true,
+			wantSampleFuncs: []string{
+				"main(main.c);(allocations) 1000",
+				"main(main.c);foo(foo.c);foo_caller(foo.c);(allocations) 100",
+				// TODO: bug #651, this line should say 1024(allocations)
+				"main(main.c);foo(foo.c);foo_caller(foo.c);(allocations) 10",
+				"main(main.c);foo_caller(foo.c);1024B,2048B(allocations) 10000",
+				"main(main.c);foo_caller(foo.c);1024B,1024B(allocations) 1",
+			},
+		},
+		{
+			name:    "Numeric units are added with units with tagroot",
+			tagroot: []string{"allocations"},
+			rootm:   true,
+			wantSampleFuncs: []string{
+				"(allocations);main(main.c) 1000",
+				"(allocations);main(main.c);foo(foo.c);foo_caller(foo.c) 100",
+				// TODO: bug #651, this line should say 1024(allocations)
+				"(allocations);main(main.c);foo(foo.c);foo_caller(foo.c) 10",
+				"1024B,2048B(allocations);main(main.c);foo_caller(foo.c) 10000",
+				"1024B,1024B(allocations);main(main.c);foo_caller(foo.c) 1",
+			},
+		},
+		{
+			name:       "Numeric labels are formatted according to outputUnit",
+			outputUnit: "kB",
+			tagleaf:    []string{"allocations"},
+			leafm:      true,
+			wantSampleFuncs: []string{
+				"main(main.c);(allocations) 1000",
+				"main(main.c);foo(foo.c);foo_caller(foo.c);(allocations) 100",
+				// TODO: bug, this line should say 1024(allocations)
+				"main(main.c);foo(foo.c);foo_caller(foo.c);(allocations) 10",
+				"main(main.c);foo_caller(foo.c);1kB,2kB(allocations) 10000",
+				"main(main.c);foo_caller(foo.c);1kB,1kB(allocations) 1",
+			},
+		},
+	} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			p := testProfile1.Copy()
+			rootm, leafm := addLabelNodes(p, tc.tagroot, tc.tagleaf, tc.outputUnit)
+			if rootm != tc.rootm {
+				t.Errorf("Got rootm=%v, want=%v", rootm, tc.rootm)
+			}
+			if leafm != tc.leafm {
+				t.Errorf("Got leafm=%v, want=%v", leafm, tc.leafm)
+			}
+			if got, want := strings.Join(stackCollapse(p), "\n")+"\n", strings.Join(tc.wantSampleFuncs, "\n")+"\n"; got != want {
+				diff, err := proftest.Diff([]byte(want), []byte(got))
+				if err != nil {
+					t.Fatalf("Failed to get diff: %v", err)
+				}
+				t.Errorf("Profile samples got diff(want->got):\n%s", diff)
+			}
+		})
+	}
+}
+
+// stackCollapse returns a slice of strings where each string represents one
+// profile sample in Brendan Gregg's "Folded Stacks" format:
+// "<root_fn>(filename);<fun2>(filename);<leaf_fn>(filename) <value>". This
+// allows the expected values for test cases to be specified in human-readable
+// strings.
+func stackCollapse(p *profile.Profile) []string {
+	var ret []string
+	for _, s := range p.Sample {
+		var funcs []string
+		for i := range s.Location {
+			loc := s.Location[len(s.Location)-1-i]
+			for _, line := range loc.Line {
+				funcs = append(funcs, fmt.Sprintf("%s(%s)", line.Function.Name, line.Function.Filename))
+			}
+		}
+		ret = append(ret, fmt.Sprintf("%s %d", strings.Join(funcs, ";"), s.Value[0]))
+	}
+	return ret
+}