diff --git a/src/go/build/deps_test.go b/src/go/build/deps_test.go index 147eaf6aba3..1ee23eda69e 100644 --- a/src/go/build/deps_test.go +++ b/src/go/build/deps_test.go @@ -175,7 +175,7 @@ var pkgDeps = map[string][]string{ "regexp/syntax": {"L2"}, "runtime/debug": {"L2", "fmt", "io/ioutil", "os", "time"}, "runtime/pprof/internal/protopprof": {"L2", "fmt", "internal/pprof/profile", "os", "time"}, - "runtime/pprof": {"L2", "fmt", "internal/pprof/profile", "os", "runtime/pprof/internal/protopprof", "text/tabwriter", "time"}, + "runtime/pprof": {"L2", "context", "fmt", "internal/pprof/profile", "os", "runtime/pprof/internal/protopprof", "text/tabwriter", "time"}, "runtime/trace": {"L0"}, "text/tabwriter": {"L2"}, diff --git a/src/runtime/pprof/label.go b/src/runtime/pprof/label.go new file mode 100644 index 00000000000..6643336db45 --- /dev/null +++ b/src/runtime/pprof/label.go @@ -0,0 +1,77 @@ +// Copyright 2016 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package pprof + +import ( + "context" +) + +type label struct { + key string + value string +} + +// LabelSet is a set of labels. +type LabelSet struct { + list []label +} + +// labelContextKey is the type of contextKeys used for profiler labels. +type labelContextKey struct{} + +// labelMap is the representation of the label set held in the context type. +// This is an initial implementation, but it will be replaced with something +// that admits incremental immutable modification more efficiently. +type labelMap map[string]string + +// WithLabels returns a new context.Context with the given labels added. +// A label overwrites a prior label with the same key. +func WithLabels(ctx context.Context, labels LabelSet) context.Context { + childLabels := make(labelMap) + parentLabels, _ := ctx.Value(labelContextKey{}).(labelMap) + // TODO(matloob): replace the map implementation with something + // more efficient so creating a child context WithLabels doesn't need + // to clone the map. + for k, v := range parentLabels { + childLabels[k] = v + } + for _, label := range labels.list { + childLabels[label.key] = label.value + } + return context.WithValue(ctx, labelContextKey{}, childLabels) +} + +// Labels takes an even number of strings representing key-value pairs +// and makes a LabelList containing them. +// A label overwrites a prior label with the same key. +func Labels(args ...string) LabelSet { + if len(args)%2 != 0 { + panic("uneven number of arguments to pprof.Labels") + } + labels := LabelSet{} + for i := 0; i+1 < len(args); i += 2 { + labels.list = append(labels.list, label{key: args[i], value: args[i+1]}) + } + return labels +} + +// Label returns the value of the label with the given key on ctx, and a boolean indicating +// whether that label exists. +func Label(ctx context.Context, key string) (string, bool) { + ctxLabels, _ := ctx.Value(labelContextKey{}).(labelMap) + v, ok := ctxLabels[key] + return v, ok +} + +// ForLabels invokes f with each label set on the context. +// The function f should return true to continue iteration or false to stop iteration early. +func ForLabels(ctx context.Context, f func(key, value string) bool) { + ctxLabels, _ := ctx.Value(labelContextKey{}).(labelMap) + for k, v := range ctxLabels { + if !f(k, v) { + break + } + } +} diff --git a/src/runtime/pprof/label_test.go b/src/runtime/pprof/label_test.go new file mode 100644 index 00000000000..d4a3fe92a91 --- /dev/null +++ b/src/runtime/pprof/label_test.go @@ -0,0 +1,82 @@ +package pprof + +import ( + "context" + "reflect" + "sort" + "testing" +) + +func labelsSorted(ctx context.Context) []label { + ls := []label{} + ForLabels(ctx, func(key, value string) bool { + ls = append(ls, label{key, value}) + return true + }) + sort.Sort(labelSorter(ls)) + return ls +} + +type labelSorter []label + +func (s labelSorter) Len() int { return len(s) } +func (s labelSorter) Swap(i, j int) { s[i], s[j] = s[j], s[i] } +func (s labelSorter) Less(i, j int) bool { return s[i].key < s[j].key } + +func TestContextLabels(t *testing.T) { + // Background context starts with no lablels. + ctx := context.Background() + labels := labelsSorted(ctx) + if len(labels) != 0 { + t.Errorf("labels on background context: want [], got %v ", labels) + } + + // Add a single label. + ctx = WithLabels(ctx, Labels("key", "value")) + // Retreive it with Label. + v, ok := Label(ctx, "key") + if !ok || v != "value" { + t.Errorf(`Label(ctx, "key"): got %v, %v; want "value", ok`, v, ok) + } + gotLabels := labelsSorted(ctx) + wantLabels := []label{{"key", "value"}} + if !reflect.DeepEqual(gotLabels, wantLabels) { + t.Errorf("(sorted) labels on context: got %v, want %v", gotLabels, wantLabels) + } + + // Add a label with a different key. + ctx = WithLabels(ctx, Labels("key2", "value2")) + v, ok = Label(ctx, "key2") + if !ok || v != "value2" { + t.Errorf(`Label(ctx, "key2"): got %v, %v; want "value2", ok`, v, ok) + } + gotLabels = labelsSorted(ctx) + wantLabels = []label{{"key", "value"}, {"key2", "value2"}} + if !reflect.DeepEqual(gotLabels, wantLabels) { + t.Errorf("(sorted) labels on context: got %v, want %v", gotLabels, wantLabels) + } + + // Add label with first key to test label replacement. + ctx = WithLabels(ctx, Labels("key", "value3")) + v, ok = Label(ctx, "key") + if !ok || v != "value3" { + t.Errorf(`Label(ctx, "key3"): got %v, %v; want "value3", ok`, v, ok) + } + gotLabels = labelsSorted(ctx) + wantLabels = []label{{"key", "value3"}, {"key2", "value2"}} + if !reflect.DeepEqual(gotLabels, wantLabels) { + t.Errorf("(sorted) labels on context: got %v, want %v", gotLabels, wantLabels) + } + + // Labels called with two labels with the same key should pick the second. + ctx = WithLabels(ctx, Labels("key4", "value4a", "key4", "value4b")) + v, ok = Label(ctx, "key4") + if !ok || v != "value4b" { + t.Errorf(`Label(ctx, "key4"): got %v, %v; want "value4b", ok`, v, ok) + } + gotLabels = labelsSorted(ctx) + wantLabels = []label{{"key", "value3"}, {"key2", "value2"}, {"key4", "value4b"}} + if !reflect.DeepEqual(gotLabels, wantLabels) { + t.Errorf("(sorted) labels on context: got %v, want %v", gotLabels, wantLabels) + } +}