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  	"slices"
    11  	"strings"
    12  )
    13  
    14  type label struct {
    15  	key   string
    16  	value string
    17  }
    18  
    19  // LabelSet is a set of labels.
    20  type LabelSet struct {
    21  	list []label
    22  }
    23  
    24  // labelContextKey is the type of contextKeys used for profiler labels.
    25  type labelContextKey struct{}
    26  
    27  func 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.
    38  type labelMap map[string]string
    39  
    40  // String satisfies Stringer and returns key, value pairs in a consistent
    41  // order.
    42  func (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.
    59  func 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.
    80  func 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.
    93  func 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.
   101  func 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
View as plain text