label.go

     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  
     5  package pprof
     6  
     7  import (
     8  	"context"
     9  	"fmt"
    10  	"internal/runtime/pprof/label"
    11  	"slices"
    12  	"strings"
    13  )
    14  
    15  // LabelSet is a set of labels.
    16  type LabelSet struct {
    17  	list []label.Label
    18  }
    19  
    20  // labelContextKey is the type of contextKeys used for profiler labels.
    21  type labelContextKey struct{}
    22  
    23  func labelValue(ctx context.Context) labelMap {
    24  	labels, _ := ctx.Value(labelContextKey{}).(*labelMap)
    25  	if labels == nil {
    26  		return labelMap{}
    27  	}
    28  	return *labels
    29  }
    30  
    31  // labelMap is the representation of the label set held in the context type.
    32  // This is an initial implementation, but it will be replaced with something
    33  // that admits incremental immutable modification more efficiently.
    34  type labelMap struct {
    35  	label.Set
    36  }
    37  
    38  // String satisfies Stringer and returns key, value pairs in a consistent
    39  // order.
    40  func (l *labelMap) String() string {
    41  	if l == nil {
    42  		return ""
    43  	}
    44  	keyVals := make([]string, 0, len(l.Set.List))
    45  
    46  	for _, lbl := range l.Set.List {
    47  		keyVals = append(keyVals, fmt.Sprintf("%q:%q", lbl.Key, lbl.Value))
    48  	}
    49  
    50  	slices.Sort(keyVals)
    51  	return "{" + strings.Join(keyVals, ", ") + "}"
    52  }
    53  
    54  // WithLabels returns a new [context.Context] with the given labels added.
    55  // A label overwrites a prior label with the same key.
    56  func WithLabels(ctx context.Context, labels LabelSet) context.Context {
    57  	parentLabels := labelValue(ctx)
    58  	return context.WithValue(ctx, labelContextKey{}, &labelMap{mergeLabelSets(parentLabels.Set, labels)})
    59  }
    60  
    61  func mergeLabelSets(left label.Set, right LabelSet) label.Set {
    62  	if len(left.List) == 0 {
    63  		return label.NewSet(right.list)
    64  	} else if len(right.list) == 0 {
    65  		return left
    66  	}
    67  
    68  	lList, rList := left.List, right.list
    69  	l, r := 0, 0
    70  	result := make([]label.Label, 0, len(rList))
    71  	for l < len(lList) && r < len(rList) {
    72  		switch strings.Compare(lList[l].Key, rList[r].Key) {
    73  		case -1: // left key < right key
    74  			result = append(result, lList[l])
    75  			l++
    76  		case 1: // right key < left key
    77  			result = append(result, rList[r])
    78  			r++
    79  		case 0: // keys are equal, right value overwrites left value
    80  			result = append(result, rList[r])
    81  			l++
    82  			r++
    83  		}
    84  	}
    85  
    86  	// Append the remaining elements
    87  	result = append(result, lList[l:]...)
    88  	result = append(result, rList[r:]...)
    89  
    90  	return label.NewSet(result)
    91  }
    92  
    93  // Labels takes an even number of strings representing key-value pairs
    94  // and makes a [LabelSet] containing them.
    95  // A label overwrites a prior label with the same key.
    96  // Currently only the CPU and goroutine profiles utilize any labels
    97  // information.
    98  // See https://golang.org/issue/23458 for details.
    99  func Labels(args ...string) LabelSet {
   100  	if len(args)%2 != 0 {
   101  		panic("uneven number of arguments to pprof.Labels")
   102  	}
   103  	list := make([]label.Label, 0, len(args)/2)
   104  	sortedNoDupes := true
   105  	for i := 0; i+1 < len(args); i += 2 {
   106  		list = append(list, label.Label{Key: args[i], Value: args[i+1]})
   107  		sortedNoDupes = sortedNoDupes && (i < 2 || args[i] > args[i-2])
   108  	}
   109  	if !sortedNoDupes {
   110  		// slow path: keys are unsorted, contain duplicates, or both
   111  		slices.SortStableFunc(list, func(a, b label.Label) int {
   112  			return strings.Compare(a.Key, b.Key)
   113  		})
   114  		deduped := make([]label.Label, 0, len(list))
   115  		for i, lbl := range list {
   116  			if i == 0 || lbl.Key != list[i-1].Key {
   117  				deduped = append(deduped, lbl)
   118  			} else {
   119  				deduped[len(deduped)-1] = lbl
   120  			}
   121  		}
   122  		list = deduped
   123  	}
   124  	return LabelSet{list: list}
   125  }
   126  
   127  // Label returns the value of the label with the given key on ctx, and a boolean indicating
   128  // whether that label exists.
   129  func Label(ctx context.Context, key string) (string, bool) {
   130  	ctxLabels := labelValue(ctx)
   131  	for _, lbl := range ctxLabels.Set.List {
   132  		if lbl.Key == key {
   133  			return lbl.Value, true
   134  		}
   135  	}
   136  	return "", false
   137  }
   138  
   139  // ForLabels invokes f with each label set on the context.
   140  // The function f should return true to continue iteration or false to stop iteration early.
   141  func ForLabels(ctx context.Context, f func(key, value string) bool) {
   142  	ctxLabels := labelValue(ctx)
   143  	for _, lbl := range ctxLabels.Set.List {
   144  		if !f(lbl.Key, lbl.Value) {
   145  			break
   146  		}
   147  	}
   148  }
   149
View as plain text