add -diff flag for better profile comparision (#369)

doc/README.md

 * **-weblist= _regex_:** Generates a source/assembly combined annotated listing for
   functions matching *regex*, and starts a web browser to display it.
 
+## Comparing profiles
+
+pprof can subtract one profile from another, provided the profiles are of
+compatible types (i.e. two heap profiles). pprof has two options which can be
+used to specify the filename or URL for a profile to be subtracted from the
+source profile:
+
+* **-diff_base= _profile_:** useful for comparing two profiles. Percentages in
+the output are relative to the total of samples in the diff base profile.
+
+* **-base= _profile_:** useful for subtracting a cumulative profile, like a
+[golang block profile](https://golang.org/doc/diagnostics.html#profiling),
+from another cumulative profile collected from the same program at a later time.
+When comparing cumulative profiles collected on the same program, percentages in
+the output are relative to the difference between the total for the source
+profile and the total for the base profile.
+
+The **-normalize** flag can be used when a base profile is specified with either 
+the `-diff_base` or the `-base` option. This flag scales the source profile so
+that the total of samples in the source profile is equal to the total of samples
+in the base profile prior to subtracting the base profile from the source
+profile. Useful for determining the relative differences between profiles, for
+example, which profile has a larger percentage of CPU time used in a particular
+function.
+
+When using the **-diff_base** option, some report entries may have negative
+values. If the merged profile is output as a protocol buffer, all samples in the
+diff base profile will have a label with the key "pprof::base" and a value of
+"true". If pprof is then used to look at the merged profile, it will behave as
+if separate source and base profiles were passed in.
+
+When using the **-base** option to subtract one cumulative profile from another
+collected on the same program at a later time, percentages will be relative to
+the difference between the total for the source profile and the total for
+the base profile, and all values will be positive. In the general case, some 
+report entries may have negative values and percentages will be relative to the
+total of the absolute value of all samples when aggregated at the address level.
+
 # Fetching profiles
 
 pprof can read profiles from a file or directly from a URL over http. Its native
 distributed job. The profiles may be from different programs but must be
 compatible (for example, CPU profiles cannot be combined with heap profiles).
 
-pprof can subtract a profile from another in order to compare them. For that,
-use the **-base= _profile_** option, where *profile* is the filename or URL for the
-profile to be subtracted. This may result on some report entries having negative
-values.
-
 ## Symbolization
 
 pprof can add symbol information to a profile that was collected only with
 
 * **-symbolize=demangle=templates:** Demangle, and trim function parameters, but
   not template parameters.
-

internal/driver/cli.go

 package driver
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"strings"
 	ExecName  string
 	BuildID   string
 	Base      []string
+	DiffBase  bool
 	Normalize bool
 
 	Seconds      int
 func parseFlags(o *plugin.Options) (*source, []string, error) {
 	flag := o.Flagset
 	// Comparisons.
-	flagBase := flag.StringList("base", "", "Source for base profile for comparison")
+	flagBase := flag.StringList("base", "", "Source for base profile for profile subtraction")
+	flagDiffBase := flag.StringList("diff_base", "", "Source for diff base profile for comparison")
 	// Source options.
 	flagSymbolize := flag.String("symbolize", "", "Options for profile symbolization")
 	flagBuildID := flag.String("buildid", "", "Override build id for first mapping")
 			usageMsgVars)
 	})
 	if len(args) == 0 {
-		return nil, nil, fmt.Errorf("no profile source specified")
+		return nil, nil, errors.New("no profile source specified")
 	}
 
 	var execName string
 		return nil, nil, err
 	}
 	if cmd != nil && *flagHTTP != "" {
-		return nil, nil, fmt.Errorf("-http is not compatible with an output format on the command line")
+		return nil, nil, errors.New("-http is not compatible with an output format on the command line")
 	}
 
 	si := pprofVariables["sample_index"].value
 		Comment:      *flagAddComment,
 	}
 
-	for _, s := range *flagBase {
-		if *s != "" {
-			source.Base = append(source.Base, *s)
-		}
+	if err := source.addBaseProfiles(*flagBase, *flagDiffBase); err != nil {
+		return nil, nil, err
 	}
 
 	normalize := pprofVariables["normalize"].boolValue()
 	if normalize && len(source.Base) == 0 {
-		return nil, nil, fmt.Errorf("Must have base profile to normalize by")
+		return nil, nil, errors.New("must have base profile to normalize by")
 	}
 	source.Normalize = normalize
 
 	return source, cmd, nil
 }
 
+// addBaseProfiles adds the list of base profiles or diff base profiles to
+// the source. This function will return an error if both base and diff base
+// profiles are specified.
+func (source *source) addBaseProfiles(flagBase, flagDiffBase []*string) error {
+	base, diffBase := dropEmpty(flagBase), dropEmpty(flagDiffBase)
+	if len(base) > 0 && len(diffBase) > 0 {
+		return errors.New("-base and -diff_base flags cannot both be specified")
+	}
+
+	source.Base = base
+	if len(diffBase) > 0 {
+		source.Base, source.DiffBase = diffBase, true
+	}
+	return nil
+}
+
+// dropEmpty list takes a slice of string pointers, and outputs a slice of
+// non-empty strings associated with the flag.
+func dropEmpty(list []*string) []string {
+	var l []string
+	for _, s := range list {
+		if *s != "" {
+			l = append(l, *s)
+		}
+	}
+	return l
+}
+
 // installFlags creates command line flags for pprof variables.
 func installFlags(flag plugin.FlagSet) flagsInstalled {
 	f := flagsInstalled{
 	for n, b := range bcmd {
 		if *b {
 			if cmd != nil {
-				return nil, fmt.Errorf("must set at most one output format")
+				return nil, errors.New("must set at most one output format")
 			}
 			cmd = []string{n}
 		}
 	for n, s := range acmd {
 		if *s != "" {
 			if cmd != nil {
-				return nil, fmt.Errorf("must set at most one output format")
+				return nil, errors.New("must set at most one output format")
 			}
 			cmd = []string{n, *s}
 		}

internal/driver/driver_test.go

 	floats      map[string]float64
 	strings     map[string]string
 	args        []string
-	stringLists map[string][]*string
+	stringLists map[string][]string
 }
 
 func (testFlags) ExtraUsage() string { return "" }
 
 func (f testFlags) StringList(s, d, c string) *[]*string {
 	if t, ok := f.stringLists[s]; ok {
-		return &t
+		// convert slice of strings to slice of string pointers before returning.
+		tp := make([]*string, len(t))
+		for i, v := range t {
+			tp[i] = &v
+		}
+		return &tp
 	}
 	return &[]*string{}
 }

internal/driver/fetch.go

 	}
 
 	if pbase != nil {
+		if s.DiffBase {
+			pbase.SetLabel("pprof::base", []string{"true"})
+		}
 		if s.Normalize {
 			err := p.Normalize(pbase)
 			if err != nil {

internal/driver/fetch_test.go

 	baseVars := pprofVariables
 	defer func() { pprofVariables = baseVars }()
 
+	type WantSample struct {
+		values []int64
+		labels map[string][]string
+	}
+
 	const path = "testdata/"
 	type testcase struct {
-		desc            string
-		sources         []string
-		bases           []string
-		normalize       bool
-		expectedSamples [][]int64
+		desc         string
+		sources      []string
+		bases        []string
+		diffBases    []string
+		normalize    bool
+		wantSamples  []WantSample
+		wantErrorMsg string
 	}
 
 	testcases := []testcase{
 			"not normalized base is same as source",
 			[]string{path + "cppbench.contention"},
 			[]string{path + "cppbench.contention"},
+			nil,
+			false,
+			nil,
+			"",
+		},
+		{
+			"not normalized base is same as source",
+			[]string{path + "cppbench.contention"},
+			[]string{path + "cppbench.contention"},
+			nil,
 			false,
-			[][]int64{},
+			nil,
+			"",
 		},
 		{
 			"not normalized single source, multiple base (all profiles same)",
 			[]string{path + "cppbench.contention"},
 			[]string{path + "cppbench.contention", path + "cppbench.contention"},
+			nil,
 			false,
-			[][]int64{{-2700, -608881724}, {-100, -23992}, {-200, -179943}, {-100, -17778444}, {-100, -75976}, {-300, -63568134}},
+			[]WantSample{
+				{
+					values: []int64{-2700, -608881724},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{-100, -23992},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{-200, -179943},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{-100, -17778444},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{-100, -75976},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{-300, -63568134},
+					labels: map[string][]string{},
+				},
+			},
+			"",
 		},
 		{
 			"not normalized, different base and source",
 			[]string{path + "cppbench.contention"},
 			[]string{path + "cppbench.small.contention"},
+			nil,
 			false,
-			[][]int64{{1700, 608878600}, {100, 23992}, {200, 179943}, {100, 17778444}, {100, 75976}, {300, 63568134}},
+			[]WantSample{
+				{
+					values: []int64{1700, 608878600},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{100, 23992},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{200, 179943},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{100, 17778444},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{100, 75976},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{300, 63568134},
+					labels: map[string][]string{},
+				},
+			},
+			"",
 		},
 		{
 			"normalized base is same as source",
 			[]string{path + "cppbench.contention"},
 			[]string{path + "cppbench.contention"},
+			nil,
 			true,
-			[][]int64{},
+			nil,
+			"",
 		},
 		{
 			"normalized single source, multiple base (all profiles same)",
 			[]string{path + "cppbench.contention"},
 			[]string{path + "cppbench.contention", path + "cppbench.contention"},
+			nil,
 			true,
-			[][]int64{},
+			nil,
+			"",
 		},
 		{
 			"normalized different base and source",
 			[]string{path + "cppbench.contention"},
 			[]string{path + "cppbench.small.contention"},
+			nil,
 			true,
-			[][]int64{{-229, -370}, {28, 0}, {57, 0}, {28, 80}, {28, 0}, {85, 287}},
+			[]WantSample{
+				{
+					values: []int64{-229, -370},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{28, 0},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{57, 0},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{28, 80},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{28, 0},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{85, 287},
+					labels: map[string][]string{},
+				},
+			},
+			"",
+		},
+		{
+			"not normalized diff base is same as source",
+			[]string{path + "cppbench.contention"},
+			nil,
+			[]string{path + "cppbench.contention"},
+			false,
+			[]WantSample{
+				{
+					values: []int64{2700, 608881724},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{100, 23992},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{200, 179943},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{100, 17778444},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{100, 75976},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{300, 63568134},
+					labels: map[string][]string{},
+				},
+				{
+					values: []int64{-2700, -608881724},
+					labels: map[string][]string{"pprof::base": {"true"}},
+				},
+				{
+					values: []int64{-100, -23992},
+					labels: map[string][]string{"pprof::base": {"true"}},
+				},
+				{
+					values: []int64{-200, -179943},
+					labels: map[string][]string{"pprof::base": {"true"}},
+				},
+				{
+					values: []int64{-100, -17778444},
+					labels: map[string][]string{"pprof::base": {"true"}},
+				},
+				{
+					values: []int64{-100, -75976},
+					labels: map[string][]string{"pprof::base": {"true"}},
+				},
+				{
+					values: []int64{-300, -63568134},
+					labels: map[string][]string{"pprof::base": {"true"}},
+				},
+			},
+			"",
+		},
+		{
+			"diff_base and base both specified",
+			[]string{path + "cppbench.contention"},
+			[]string{path + "cppbench.contention"},
+			[]string{path + "cppbench.contention"},
+			false,
+			nil,
+			"-base and -diff_base flags cannot both be specified",
 		},
 	}
 
 	for _, tc := range testcases {
 		t.Run(tc.desc, func(t *testing.T) {
 			pprofVariables = baseVars.makeCopy()
-
-			base := make([]*string, len(tc.bases))
-			for i, s := range tc.bases {
-				base[i] = &s
-			}
-
 			f := testFlags{
-				stringLists: map[string][]*string{
-					"base": base,
+				stringLists: map[string][]string{
+					"base":      tc.bases,
+					"diff_base": tc.diffBases,
 				},
 				bools: map[string]bool{
 					"normalize": tc.normalize,
 			})
 			src, _, err := parseFlags(o)
 
+			if tc.wantErrorMsg != "" {
+				if err == nil {
+					t.Fatalf("got nil, want error %q", tc.wantErrorMsg)
+				}
+
+				if gotErrMsg := err.Error(); gotErrMsg != tc.wantErrorMsg {
+					t.Fatalf("got error %q, want error %q", gotErrMsg, tc.wantErrorMsg)
+				}
+				return
+			}
+
 			if err != nil {
-				t.Fatalf("%s: %v", tc.desc, err)
+				t.Fatalf("got error %q, want no error", err)
 			}
 
 			p, err := fetchProfiles(src, o)
-			pprofVariables = baseVars
+
 			if err != nil {
-				t.Fatal(err)
+				t.Fatalf("got error %q, want no error", err)
 			}
 
-			if want, got := len(tc.expectedSamples), len(p.Sample); want != got {
-				t.Fatalf("want %d samples got %d", want, got)
+			if got, want := len(p.Sample), len(tc.wantSamples); got != want {
+				t.Fatalf("got %d samples want %d", got, want)
 			}
 
-			if len(p.Sample) > 0 {
-				for i, sample := range p.Sample {
-					if want, got := len(tc.expectedSamples[i]), len(sample.Value); want != got {
-						t.Errorf("want %d values for sample %d, got %d", want, i, got)
-					}
-					for j, value := range sample.Value {
-						if want, got := tc.expectedSamples[i][j], value; want != got {
-							t.Errorf("want value of %d for value %d of sample %d, got %d", want, j, i, got)
-						}
-					}
+			for i, sample := range p.Sample {
+				if !reflect.DeepEqual(tc.wantSamples[i].values, sample.Value) {
+					t.Errorf("for sample %d got values %v, want %v", i, sample.Value, tc.wantSamples[i])
+				}
+				if !reflect.DeepEqual(tc.wantSamples[i].labels, sample.Label) {
+					t.Errorf("for sample %d got labels %v, want %v", i, sample.Label, tc.wantSamples[i].labels)
 				}
 			}
 		})

internal/report/report.go

 		s.NumUnit = numUnits
 	}
 
+	// Remove label marking samples from the base profiles, so it does not appear
+	// as a nodelet in the graph view.
+	prof.RemoveLabel("pprof::base")
+
 	formatTag := func(v int64, key string) string {
 		return measurement.ScaledLabel(v, key, o.OutputUnit)
 	}
 	return New(prof, o)
 }
 
-// computeTotal computes the sum of all sample values. This will be
-// used to compute percentages.
+// computeTotal computes the sum of the absolute value of all sample values.
+// If any samples have the label "pprof::base" with value "true", then the total
+// will only include samples with that label.
 func computeTotal(prof *profile.Profile, value, meanDiv func(v []int64) int64) int64 {
-	var div, ret int64
+	var div, total, diffDiv, diffTotal int64
 	for _, sample := range prof.Sample {
 		var d, v int64
 		v = value(sample.Value)
 		if v < 0 {
 			v = -v
 		}
-		ret += v
+		total += v
 		div += d
+		if sample.HasLabel("pprof::base", "true") {
+			diffTotal += v
+			diffDiv += d
+		}
+	}
+	if diffTotal > 0 {
+		total = diffTotal
+		div = diffDiv
 	}
 	if div != 0 {
-		return ret / div
+		return total / div
 	}
-	return ret
+	return total
 }
 
 // Report contains the data and associated routines to extract a

internal/report/report_test.go

 		}
 	}
 }
+
+func TestComputeTotal(t *testing.T) {
+	p1 := testProfile.Copy()
+	p1.Sample = []*profile.Sample{
+		{
+			Location: []*profile.Location{testL[0]},
+			Value:    []int64{1, 1},
+		},
+		{
+			Location: []*profile.Location{testL[2], testL[1], testL[0]},
+			Value:    []int64{1, 10},
+		},
+		{
+			Location: []*profile.Location{testL[4], testL[2], testL[0]},
+			Value:    []int64{1, 100},
+		},
+	}
+
+	p2 := testProfile.Copy()
+	p2.Sample = []*profile.Sample{
+		{
+			Location: []*profile.Location{testL[0]},
+			Value:    []int64{1, 1},
+		},
+		{
+			Location: []*profile.Location{testL[2], testL[1], testL[0]},
+			Value:    []int64{1, -10},
+		},
+		{
+			Location: []*profile.Location{testL[4], testL[2], testL[0]},
+			Value:    []int64{1, 100},
+		},
+	}
+
+	p3 := testProfile.Copy()
+	p3.Sample = []*profile.Sample{
+		{
+			Location: []*profile.Location{testL[0]},
+			Value:    []int64{10000, 1},
+		},
+		{
+			Location: []*profile.Location{testL[2], testL[1], testL[0]},
+			Value:    []int64{-10, 3},
+			Label:    map[string][]string{"pprof::base": {"true"}},
+		},
+		{
+			Location: []*profile.Location{testL[2], testL[1], testL[0]},
+			Value:    []int64{1000, -10},
+		},
+		{
+			Location: []*profile.Location{testL[2], testL[1], testL[0]},
+			Value:    []int64{-9000, 3},
+			Label:    map[string][]string{"pprof::base": {"true"}},
+		},
+		{
+			Location: []*profile.Location{testL[2], testL[1], testL[0]},
+			Value:    []int64{-1, 3},
+			Label:    map[string][]string{"pprof::base": {"true"}},
+		},
+		{
+			Location: []*profile.Location{testL[4], testL[2], testL[0]},
+			Value:    []int64{100, 100},
+		},
+		{
+			Location: []*profile.Location{testL[2], testL[1], testL[0]},
+			Value:    []int64{100, 3},
+			Label:    map[string][]string{"pprof::base": {"true"}},
+		},
+	}
+
+	testcases := []struct {
+		desc           string
+		prof           *profile.Profile
+		value, meanDiv func(v []int64) int64
+		wantTotal      int64
+	}{
+		{
+			desc: "no diff base, all positive values, index 1",
+			prof: p1,
+			value: func(v []int64) int64 {
+				return v[0]
+			},
+			wantTotal: 3,
+		},
+		{
+			desc: "no diff base, all positive values, index 2",
+			prof: p1,
+			value: func(v []int64) int64 {
+				return v[1]
+			},
+			wantTotal: 111,
+		},
+		{
+			desc: "no diff base, some negative values",
+			prof: p2,
+			value: func(v []int64) int64 {
+				return v[1]
+			},
+			wantTotal: 111,
+		},
+		{
+			desc: "diff base, some negative values",
+			prof: p3,
+			value: func(v []int64) int64 {
+				return v[0]
+			},
+			wantTotal: 9111,
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.desc, func(t *testing.T) {
+			if gotTotal := computeTotal(tc.prof, tc.value, tc.meanDiv); gotTotal != tc.wantTotal {
+				t.Errorf("got total %d, want %v", gotTotal, tc.wantTotal)
+			}
+		})
+	}
+}

profile/profile.go

 	return strings.Join(ls, " ")
 }
 
+// SetLabel sets the specified key to the specified value for all samples in the
+// profile.
+func (p *Profile) SetLabel(key string, value []string) {
+	for _, sample := range p.Sample {
+		if sample.Label == nil {
+			sample.Label = map[string][]string{key: value}
+		} else {
+			sample.Label[key] = value
+		}
+	}
+}
+
+// RemoveLabel removes all labels associated with the specified key for all
+// samples in the profile.
+func (p *Profile) RemoveLabel(key string) {
+	for _, sample := range p.Sample {
+		delete(sample.Label, key)
+	}
+}
+
+// HasLabel returns true if a sample has a label with indicated key and value.
+func (s *Sample) HasLabel(key, value string) bool {
+	for _, v := range s.Label[key] {
+		if v == value {
+			return true
+		}
+	}
+	return false
+}
+
 // Scale multiplies all sample values in a profile by a constant.
 func (p *Profile) Scale(ratio float64) {
 	if ratio == 1 {

profile/profile_test.go

 		wantNumUnits  []map[string][]string
 	}{
 		{
-			name:  "different tag units not merged",
+			name:  "different label units not merged",
 			profs: []*Profile{testProfile4.Copy(), testProfile5.Copy()},
 			wantNumLabels: []map[string][]int64{
 				{
 	return tb
 }
 
+func TestHasLabel(t *testing.T) {
+	var testcases = []struct {
+		desc         string
+		labels       map[string][]string
+		key          string
+		value        string
+		wantHasLabel bool
+	}{
+		{
+			desc:         "empty label does not have label",
+			labels:       map[string][]string{},
+			key:          "key",
+			value:        "value",
+			wantHasLabel: false,
+		},
+		{
+			desc:         "label with one key and value has label",
+			labels:       map[string][]string{"key": {"value"}},
+			key:          "key",
+			value:        "value",
+			wantHasLabel: true,
+		},
+		{
+			desc:         "label with one key and value does not have label",
+			labels:       map[string][]string{"key": {"value"}},
+			key:          "key1",
+			value:        "value1",
+			wantHasLabel: false,
+		},
+		{
+			desc: "label with many keys and values has label",
+			labels: map[string][]string{
+				"key1": {"value2", "value1"},
+				"key2": {"value1", "value2", "value2"},
+				"key3": {"value1", "value2", "value2"},
+			},
+			key:          "key1",
+			value:        "value1",
+			wantHasLabel: true,
+		},
+		{
+			desc: "label with many keys and values does not have label",
+			labels: map[string][]string{
+				"key1": {"value2", "value1"},
+				"key2": {"value1", "value2", "value2"},
+				"key3": {"value1", "value2", "value2"},
+			},
+			key:          "key5",
+			value:        "value5",
+			wantHasLabel: false,
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.desc, func(t *testing.T) {
+			sample := &Sample{
+				Label: tc.labels,
+			}
+			if gotHasLabel := sample.HasLabel(tc.key, tc.value); gotHasLabel != tc.wantHasLabel {
+				t.Errorf("sample.HasLabel(%q, %q) got %v, want %v", tc.key, tc.value, gotHasLabel, tc.wantHasLabel)
+			}
+		})
+	}
+}
+
+func TestRemove(t *testing.T) {
+	var testcases = []struct {
+		desc       string
+		samples    []*Sample
+		removeKey  string
+		wantLabels []map[string][]string
+	}{
+		{
+			desc: "some samples have label already",
+			samples: []*Sample{
+				{
+					Location: []*Location{cpuL[0]},
+					Value:    []int64{1000},
+				},
+				{
+					Location: []*Location{cpuL[0]},
+					Value:    []int64{1000},
+					Label: map[string][]string{
+						"key1": {"value1", "value2", "value3"},
+						"key2": {"value1"},
+					},
+				},
+				{
+					Location: []*Location{cpuL[0]},
+					Value:    []int64{1000},
+					Label: map[string][]string{
+						"key1": {"value2"},
+					},
+				},
+			},
+			removeKey: "key1",
+			wantLabels: []map[string][]string{
+				{},
+				{"key2": {"value1"}},
+				{},
+			},
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.desc, func(t *testing.T) {
+			profile := testProfile1.Copy()
+			profile.Sample = tc.samples
+			profile.RemoveLabel(tc.removeKey)
+			if got, want := len(profile.Sample), len(tc.wantLabels); got != want {
+				t.Fatalf("got %v samples, want %v samples", got, want)
+			}
+			for i, sample := range profile.Sample {
+				wantLabels := tc.wantLabels[i]
+				if got, want := len(sample.Label), len(wantLabels); got != want {
+					t.Errorf("got %v label keys for sample %v, want %v", got, i, want)
+					continue
+				}
+				for wantKey, wantValues := range wantLabels {
+					if gotValues, ok := sample.Label[wantKey]; ok {
+						if !reflect.DeepEqual(gotValues, wantValues) {
+							t.Errorf("for key %s, got values %v, want values %v", wantKey, gotValues, wantValues)
+						}
+					} else {
+						t.Errorf("for key %s got no values, want %v", wantKey, wantValues)
+					}
+				}
+			}
+		})
+	}
+}
+
+func TestSetLabel(t *testing.T) {
+	var testcases = []struct {
+		desc       string
+		samples    []*Sample
+		setKey     string
+		setVal     []string
+		wantLabels []map[string][]string
+	}{
+		{
+			desc: "some samples have label already",
+			samples: []*Sample{
+				{
+					Location: []*Location{cpuL[0]},
+					Value:    []int64{1000},
+				},
+				{
+					Location: []*Location{cpuL[0]},
+					Value:    []int64{1000},
+					Label: map[string][]string{
+						"key1": {"value1", "value2", "value3"},
+						"key2": {"value1"},
+					},
+				},
+				{
+					Location: []*Location{cpuL[0]},
+					Value:    []int64{1000},
+					Label: map[string][]string{
+						"key1": {"value2"},
+					},
+				},
+			},
+			setKey: "key1",
+			setVal: []string{"value1"},
+			wantLabels: []map[string][]string{
+				{"key1": {"value1"}},
+				{"key1": {"value1"}, "key2": {"value1"}},
+				{"key1": {"value1"}},
+			},
+		},
+		{
+			desc: "no samples have labels",
+			samples: []*Sample{
+				{
+					Location: []*Location{cpuL[0]},
+					Value:    []int64{1000},
+				},
+			},
+			setKey: "key1",
+			setVal: []string{"value1"},
+			wantLabels: []map[string][]string{
+				{"key1": {"value1"}},
+			},
+		},
+		{
+			desc: "all samples have some labels, but not key being added",
+			samples: []*Sample{
+				{
+					Location: []*Location{cpuL[0]},
+					Value:    []int64{1000},
+					Label: map[string][]string{
+						"key2": {"value2"},
+					},
+				},
+				{
+					Location: []*Location{cpuL[0]},
+					Value:    []int64{1000},
+					Label: map[string][]string{
+						"key3": {"value3"},
+					},
+				},
+			},
+			setKey: "key1",
+			setVal: []string{"value1"},
+			wantLabels: []map[string][]string{
+				{"key1": {"value1"}, "key2": {"value2"}},
+				{"key1": {"value1"}, "key3": {"value3"}},
+			},
+		},
+		{
+			desc: "all samples have key being added",
+			samples: []*Sample{
+				{
+					Location: []*Location{cpuL[0]},
+					Value:    []int64{1000},
+					Label: map[string][]string{
+						"key1": {"value1"},
+					},
+				},
+				{
+					Location: []*Location{cpuL[0]},
+					Value:    []int64{1000},
+					Label: map[string][]string{
+						"key1": {"value1"},
+					},
+				},
+			},
+			setKey: "key1",
+			setVal: []string{"value1"},
+			wantLabels: []map[string][]string{
+				{"key1": {"value1"}},
+				{"key1": {"value1"}},
+			},
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.desc, func(t *testing.T) {
+			profile := testProfile1.Copy()
+			profile.Sample = tc.samples
+			profile.SetLabel(tc.setKey, tc.setVal)
+			if got, want := len(profile.Sample), len(tc.wantLabels); got != want {
+				t.Fatalf("got %v samples, want %v samples", got, want)
+			}
+			for i, sample := range profile.Sample {
+				wantLabels := tc.wantLabels[i]
+				if got, want := len(sample.Label), len(wantLabels); got != want {
+					t.Errorf("got %v label keys for sample %v, want %v", got, i, want)
+					continue
+				}
+				for wantKey, wantValues := range wantLabels {
+					if gotValues, ok := sample.Label[wantKey]; ok {
+						if !reflect.DeepEqual(gotValues, wantValues) {
+							t.Errorf("for key %s, got values %v, want values %v", wantKey, gotValues, wantValues)
+						}
+					} else {
+						t.Errorf("for key %s got no values, want %v", wantKey, wantValues)
+					}
+				}
+			}
+		})
+	}
+}
+
 func TestNumLabelUnits(t *testing.T) {
 	var tagFilterTests = []struct {
 		desc             string