1package driver
2
3import (
4	"strings"
5
6	"github.com/google/pprof/internal/measurement"
7	"github.com/google/pprof/profile"
8)
9
10// addLabelNodes adds pseudo stack frames "label:value" to each Sample with
11// labels matching the supplied keys.
12//
13// rootKeys adds frames at the root of the callgraph (first key becomes new root).
14// leafKeys adds frames at the leaf of the callgraph (last key becomes new leaf).
15//
16// Returns whether there were matches found for the label keys.
17func addLabelNodes(p *profile.Profile, rootKeys, leafKeys []string, outputUnit string) (rootm, leafm bool) {
18	// Find where to insert the new locations and functions at the end of
19	// their ID spaces.
20	var maxLocID uint64
21	var maxFunctionID uint64
22	for _, loc := range p.Location {
23		if loc.ID > maxLocID {
24			maxLocID = loc.ID
25		}
26	}
27	for _, f := range p.Function {
28		if f.ID > maxFunctionID {
29			maxFunctionID = f.ID
30		}
31	}
32	nextLocID := maxLocID + 1
33	nextFuncID := maxFunctionID + 1
34
35	// Intern the new locations and functions we are generating.
36	type locKey struct {
37		functionName, fileName string
38	}
39	locs := map[locKey]*profile.Location{}
40
41	internLoc := func(locKey locKey) *profile.Location {
42		loc, found := locs[locKey]
43		if found {
44			return loc
45		}
46
47		function := &profile.Function{
48			ID:       nextFuncID,
49			Name:     locKey.functionName,
50			Filename: locKey.fileName,
51		}
52		nextFuncID++
53		p.Function = append(p.Function, function)
54
55		loc = &profile.Location{
56			ID: nextLocID,
57			Line: []profile.Line{
58				{
59					Function: function,
60				},
61			},
62		}
63		nextLocID++
64		p.Location = append(p.Location, loc)
65		locs[locKey] = loc
66		return loc
67	}
68
69	makeLabelLocs := func(s *profile.Sample, keys []string) ([]*profile.Location, bool) {
70		var locs []*profile.Location
71		var match bool
72		for i := range keys {
73			// Loop backwards, ensuring the first tag is closest to the root,
74			// and the last tag is closest to the leaves.
75			k := keys[len(keys)-1-i]
76			values := formatLabelValues(s, k, outputUnit)
77			if len(values) > 0 {
78				match = true
79			}
80			locKey := locKey{
81				functionName: strings.Join(values, ","),
82				fileName:     k,
83			}
84			loc := internLoc(locKey)
85			locs = append(locs, loc)
86		}
87		return locs, match
88	}
89
90	for _, s := range p.Sample {
91		rootsToAdd, sampleMatchedRoot := makeLabelLocs(s, rootKeys)
92		if sampleMatchedRoot {
93			rootm = true
94		}
95		leavesToAdd, sampleMatchedLeaf := makeLabelLocs(s, leafKeys)
96		if sampleMatchedLeaf {
97			leafm = true
98		}
99
100		var newLocs []*profile.Location
101		newLocs = append(newLocs, leavesToAdd...)
102		newLocs = append(newLocs, s.Location...)
103		newLocs = append(newLocs, rootsToAdd...)
104		s.Location = newLocs
105	}
106	return
107}
108
109// formatLabelValues returns all the string and numeric labels in Sample, with
110// the numeric labels formatted according to outputUnit.
111func formatLabelValues(s *profile.Sample, k string, outputUnit string) []string {
112	var values []string
113	values = append(values, s.Label[k]...)
114	numLabels := s.NumLabel[k]
115	numUnits := s.NumUnit[k]
116	if len(numLabels) != len(numUnits) {
117		return values
118	}
119	for i, numLabel := range numLabels {
120		unit := numUnits[i]
121		values = append(values, measurement.ScaledLabel(numLabel, unit, outputUnit))
122	}
123	return values
124}
125