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