1// Copyright 2016 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package pprof
6
7import (
8	"context"
9	"fmt"
10	"slices"
11	"strings"
12)
13
14type label struct {
15	key   string
16	value string
17}
18
19// LabelSet is a set of labels.
20type LabelSet struct {
21	list []label
22}
23
24// labelContextKey is the type of contextKeys used for profiler labels.
25type labelContextKey struct{}
26
27func labelValue(ctx context.Context) labelMap {
28	labels, _ := ctx.Value(labelContextKey{}).(*labelMap)
29	if labels == nil {
30		return labelMap(nil)
31	}
32	return *labels
33}
34
35// labelMap is the representation of the label set held in the context type.
36// This is an initial implementation, but it will be replaced with something
37// that admits incremental immutable modification more efficiently.
38type labelMap map[string]string
39
40// String satisfies Stringer and returns key, value pairs in a consistent
41// order.
42func (l *labelMap) String() string {
43	if l == nil {
44		return ""
45	}
46	keyVals := make([]string, 0, len(*l))
47
48	for k, v := range *l {
49		keyVals = append(keyVals, fmt.Sprintf("%q:%q", k, v))
50	}
51
52	slices.Sort(keyVals)
53
54	return "{" + strings.Join(keyVals, ", ") + "}"
55}
56
57// WithLabels returns a new [context.Context] with the given labels added.
58// A label overwrites a prior label with the same key.
59func WithLabels(ctx context.Context, labels LabelSet) context.Context {
60	parentLabels := labelValue(ctx)
61	childLabels := make(labelMap, len(parentLabels))
62	// TODO(matloob): replace the map implementation with something
63	// more efficient so creating a child context WithLabels doesn't need
64	// to clone the map.
65	for k, v := range parentLabels {
66		childLabels[k] = v
67	}
68	for _, label := range labels.list {
69		childLabels[label.key] = label.value
70	}
71	return context.WithValue(ctx, labelContextKey{}, &childLabels)
72}
73
74// Labels takes an even number of strings representing key-value pairs
75// and makes a [LabelSet] containing them.
76// A label overwrites a prior label with the same key.
77// Currently only the CPU and goroutine profiles utilize any labels
78// information.
79// See https://golang.org/issue/23458 for details.
80func Labels(args ...string) LabelSet {
81	if len(args)%2 != 0 {
82		panic("uneven number of arguments to pprof.Labels")
83	}
84	list := make([]label, 0, len(args)/2)
85	for i := 0; i+1 < len(args); i += 2 {
86		list = append(list, label{key: args[i], value: args[i+1]})
87	}
88	return LabelSet{list: list}
89}
90
91// Label returns the value of the label with the given key on ctx, and a boolean indicating
92// whether that label exists.
93func Label(ctx context.Context, key string) (string, bool) {
94	ctxLabels := labelValue(ctx)
95	v, ok := ctxLabels[key]
96	return v, ok
97}
98
99// ForLabels invokes f with each label set on the context.
100// The function f should return true to continue iteration or false to stop iteration early.
101func ForLabels(ctx context.Context, f func(key, value string) bool) {
102	ctxLabels := labelValue(ctx)
103	for k, v := range ctxLabels {
104		if !f(k, v) {
105			break
106		}
107	}
108}
109