metrics: add exemplar support for counters

adds exemplar support for counters
* utilizes Prometheus' underlying exemplar machinery
* introduces contextual counters (which were a no-op till now)
* adds testcases

addresses (a part of): kubernetes#119697

Author: rexagod

staging/src/k8s.io/component-base/metrics/counter.go

@@ -27,6 +27,7 @@ import (
 // Counter is our internal representation for our wrapping struct around prometheus
 // counters. Counter implements both kubeCollector and CounterMetric.
 type Counter struct {
+	ctx context.Context
 	CounterMetric
 	*CounterOpts
 	lazyMetric
@@ -36,6 +37,9 @@ type Counter struct {
 // The implementation of the Metric interface is expected by testutil.GetCounterMetricValue.
 var _ Metric = &Counter{}
 
+// All supported exemplar metric types implement the metricWithExemplar interface.
+var _ metricWithExemplar = &Counter{}
+
 // NewCounter returns an object which satisfies the kubeCollector and CounterMetric interfaces.
 // However, the object returned will not measure anything unless the collector is first
 // registered, since the metric is lazily instantiated.
@@ -93,11 +97,25 @@ func (c *Counter) initializeDeprecatedMetric() {
 	c.initializeMetric()
 }
 
-// WithContext allows the normal Counter metric to pass in context. The context is no-op now.
+// WithContext allows the normal Counter metric to pass in context.
 func (c *Counter) WithContext(ctx context.Context) CounterMetric {
+	c.ctx = ctx
 	return c.CounterMetric
 }
 
+// withExemplar initializes the exemplarMetric object and sets the exemplar value.
+func (c *Counter) withExemplar(v float64) {
+	(&exemplarMetric{c}).withExemplar(v)
+}
+
+func (c *Counter) Add(v float64) {
+	c.withExemplar(v)
+}
+
+func (c *Counter) Inc() {
+	c.withExemplar(1)
+}
+
 // CounterVec is the internal representation of our wrapping struct around prometheus
 // counterVecs. CounterVec implements both kubeCollector and CounterVecMetric.
 type CounterVec struct {

staging/src/k8s.io/component-base/metrics/counter_test.go

@@ -18,12 +18,15 @@ package metrics
 
 import (
 	"bytes"
+	"context"
 	"testing"
 
 	"github.com/blang/semver/v4"
+	dto "github.com/prometheus/client_model/go"
 	"github.com/prometheus/common/expfmt"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/trace"
 
 	apimachineryversion "k8s.io/apimachinery/pkg/version"
 )
@@ -286,3 +289,94 @@ func TestCounterWithLabelValueAllowList(t *testing.T) {
 		})
 	}
 }
+
+func TestCounterWithExemplar(t *testing.T) {
+	// Set exemplar.
+	fn := func(offset int) []byte {
+		arr := make([]byte, 16)
+		for i := 0; i < 16; i++ {
+			arr[i] = byte(2<<7 - i - offset)
+		}
+		return arr
+	}
+	traceID := trace.TraceID(fn(1))
+	spanID := trace.SpanID(fn(2))
+	ctxForSpanCtx := trace.ContextWithSpanContext(context.Background(), trace.NewSpanContext(trace.SpanContextConfig{
+		TraceID: traceID,
+		SpanID:  spanID,
+	}))
+	toAdd := float64(40)
+
+	// Create contextual counter.
+	counter := NewCounter(&CounterOpts{
+		Name: "metric_exemplar_test",
+		Help: "helpless",
+	})
+	_ = counter.WithContext(ctxForSpanCtx)
+
+	// Register counter.
+	registry := newKubeRegistry(apimachineryversion.Info{
+		Major:      "1",
+		Minor:      "15",
+		GitVersion: "v1.15.0-alpha-1.12345",
+	})
+	registry.MustRegister(counter)
+
+	// Call underlying exemplar methods.
+	counter.Add(toAdd)
+	counter.Inc()
+	counter.Inc()
+
+	// Gather.
+	mfs, err := registry.Gather()
+	if err != nil {
+		t.Fatalf("Gather failed %v", err)
+	}
+	if len(mfs) != 1 {
+		t.Fatalf("Got %v metric families, Want: 1 metric family", len(mfs))
+	}
+
+	// Verify metric type.
+	mf := mfs[0]
+	var m *dto.Metric
+	switch mf.GetType() {
+	case dto.MetricType_COUNTER:
+		m = mfs[0].GetMetric()[0]
+	default:
+		t.Fatalf("Got %v metric type, Want: %v metric type", mf.GetType(), dto.MetricType_COUNTER)
+	}
+
+	// Verify value.
+	want := toAdd + 2
+	got := m.GetCounter().GetValue()
+	if got != want {
+		t.Fatalf("Got %f, wanted %f as the count", got, want)
+	}
+
+	// Verify exemplars.
+	e := m.GetCounter().GetExemplar()
+	if e == nil {
+		t.Fatalf("Got nil exemplar, wanted an exemplar")
+	}
+	eLabels := e.GetLabel()
+	if eLabels == nil {
+		t.Fatalf("Got nil exemplar label, wanted an exemplar label")
+	}
+	if len(eLabels) != 2 {
+		t.Fatalf("Got %v exemplar labels, wanted 2 exemplar labels", len(eLabels))
+	}
+	for _, l := range eLabels {
+		switch *l.Name {
+		case "trace_id":
+			if *l.Value != traceID.String() {
+				t.Fatalf("Got %s as traceID, wanted %s", *l.Value, traceID.String())
+			}
+		case "span_id":
+			if *l.Value != spanID.String() {
+				t.Fatalf("Got %s as spanID, wanted %s", *l.Value, spanID.String())
+			}
+		default:
+			t.Fatalf("Got unexpected label %s", *l.Name)
+		}
+	}
+}

staging/src/k8s.io/component-base/metrics/metric.go

@@ -19,11 +19,13 @@ package metrics
 import (
 	"sync"
 
+	"go.opentelemetry.io/otel/trace"
+
 	"github.com/blang/semver/v4"
 	"github.com/prometheus/client_golang/prometheus"
 	dto "github.com/prometheus/client_model/go"
-	promext "k8s.io/component-base/metrics/prometheusextension"
 
+	promext "k8s.io/component-base/metrics/prometheusextension"
 	"k8s.io/klog/v2"
 )
 
@@ -210,6 +212,33 @@ func (c *selfCollector) Collect(ch chan<- prometheus.Metric) {
 	ch <- c.metric
 }
 
+// metricWithExemplar is an interface that knows how to attach an exemplar to certain supported metric types.
+type metricWithExemplar interface {
+	withExemplar(v float64)
+}
+
+// exemplarMetric is a holds a context to extract exemplar labels from, and a metric to attach them to. It implements the metricWithExemplar interface.
+type exemplarMetric struct {
+	*Counter
+}
+
+// withExemplar attaches an exemplar to the metric.
+func (e *exemplarMetric) withExemplar(v float64) {
+	if m, ok := e.CounterMetric.(prometheus.ExemplarAdder); ok {
+		maybeSpanCtx := trace.SpanContextFromContext(e.ctx)
+		if maybeSpanCtx.IsValid() {
+			exemplarLabels := prometheus.Labels{
+				"trace_id": maybeSpanCtx.TraceID().String(),
+				"span_id":  maybeSpanCtx.SpanID().String(),
+			}
+			m.AddWithExemplar(v, exemplarLabels)
+			return
+		}
+	}
+
+	e.CounterMetric.Add(v)
+}
+
 // no-op vecs for convenience
 var noopCounterVec = &prometheus.CounterVec{}
 var noopHistogramVec = &prometheus.HistogramVec{}