feat: Ability to add additional Exemplar metadata

prometheus-metrics-core/src/main/java/io/prometheus/metrics/core/exemplars/ExemplarSampler.java

@@ -12,6 +12,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.LongSupplier;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /**
@@ -45,8 +46,10 @@ public class ExemplarSampler {
   private final SpanContext
       spanContext; // may be null, in that case SpanContextSupplier.getSpanContext() is used.
 
+  @Nullable private final Supplier<Labels> additionalLabelsSupplier;
+
   public ExemplarSampler(ExemplarSamplerConfig config) {
-    this(config, null);
+    this(config, null, null);
   }
 
   /**
@@ -58,10 +61,29 @@ public ExemplarSampler(ExemplarSamplerConfig config) {
    * SpanContextSupplier.getSpanContext()} is called to find a span context.
    */
   public ExemplarSampler(ExemplarSamplerConfig config, @Nullable SpanContext spanContext) {
+    this(config, spanContext, null);
+  }
+
+  /**
+   * Constructor that accepts a supplier of additional labels to be merged into every
+   * automatically-sampled exemplar. The supplier is called each time an exemplar is sampled from a
+   * span context, so it can return dynamic values (e.g. a request-scoped identifier). The supplier
+   * is only called when a valid, sampled span context is present.
+   */
+  public ExemplarSampler(
+      ExemplarSamplerConfig config, @Nullable Supplier<Labels> additionalLabelsSupplier) {
+    this(config, null, additionalLabelsSupplier);
+  }
+
+  public ExemplarSampler(
+      ExemplarSamplerConfig config,
+      @Nullable SpanContext spanContext,
+      @Nullable Supplier<Labels> additionalLabelsSupplier) {
     this.config = config;
     this.exemplars = new Exemplar[config.getNumberOfExemplars()];
     this.customExemplars = new Exemplar[exemplars.length];
     this.spanContext = spanContext;
+    this.additionalLabelsSupplier = additionalLabelsSupplier;
   }
 
   public Exemplars collect() {
@@ -355,7 +377,14 @@ private Labels doSampleExemplar() {
           String traceId = spanContext.getCurrentTraceId();
           if (spanId != null && traceId != null) {
             spanContext.markCurrentSpanAsExemplar();
-            return Labels.of(Exemplar.TRACE_ID, traceId, Exemplar.SPAN_ID, spanId);
+            Labels base = Labels.of(Exemplar.TRACE_ID, traceId, Exemplar.SPAN_ID, spanId);
+            if (additionalLabelsSupplier != null) {
+              Labels extra = additionalLabelsSupplier.get();
+              if (!extra.isEmpty()) {
+                return base.merge(extra);
+              }
+            }
+            return base;
           }
         }
       }

prometheus-metrics-core/src/main/java/io/prometheus/metrics/core/metrics/Counter.java

@@ -14,6 +14,7 @@
 import java.util.List;
 import java.util.concurrent.atomic.DoubleAdder;
 import java.util.concurrent.atomic.LongAdder;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /**
@@ -35,6 +36,7 @@ public class Counter extends StatefulMetric<CounterDataPoint, Counter.DataPoint>
     implements CounterDataPoint {
 
   @Nullable private final ExemplarSamplerConfig exemplarSamplerConfig;
+  @Nullable private final Supplier<Labels> exemplarLabelsSupplier;
 
   private Counter(Builder builder, PrometheusProperties prometheusProperties) {
     super(builder);
@@ -47,6 +49,7 @@ private Counter(Builder builder, PrometheusProperties prometheusProperties) {
     } else {
       exemplarSamplerConfig = null;
     }
+    exemplarLabelsSupplier = builder.exemplarLabelsSupplier;
   }
 
   @Override
@@ -101,7 +104,7 @@ public MetricType getMetricType() {
   @Override
   protected DataPoint newDataPoint() {
     if (exemplarSamplerConfig != null) {
-      return new DataPoint(new ExemplarSampler(exemplarSamplerConfig));
+      return new DataPoint(new ExemplarSampler(exemplarSamplerConfig, exemplarLabelsSupplier));
     } else {
       return new DataPoint(null);
     }

prometheus-metrics-core/src/main/java/io/prometheus/metrics/core/metrics/Gauge.java

@@ -13,6 +13,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /**
@@ -42,6 +43,7 @@ public class Gauge extends StatefulMetric<GaugeDataPoint, Gauge.DataPoint>
     implements GaugeDataPoint {
 
   @Nullable private final ExemplarSamplerConfig exemplarSamplerConfig;
+  @Nullable private final Supplier<Labels> exemplarLabelsSupplier;
 
   private Gauge(Builder builder, PrometheusProperties prometheusProperties) {
     super(builder);
@@ -54,6 +56,7 @@ private Gauge(Builder builder, PrometheusProperties prometheusProperties) {
     } else {
       exemplarSamplerConfig = null;
     }
+    exemplarLabelsSupplier = builder.exemplarLabelsSupplier;
   }
 
   @Override
@@ -103,7 +106,7 @@ public MetricType getMetricType() {
   @Override
   protected DataPoint newDataPoint() {
     if (exemplarSamplerConfig != null) {
-      return new DataPoint(new ExemplarSampler(exemplarSamplerConfig));
+      return new DataPoint(new ExemplarSampler(exemplarSamplerConfig, exemplarLabelsSupplier));
     } else {
       return new DataPoint(null);
     }

prometheus-metrics-core/src/main/java/io/prometheus/metrics/core/metrics/Histogram.java

@@ -25,6 +25,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.DoubleAdder;
 import java.util.concurrent.atomic.LongAdder;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /**
@@ -71,6 +72,7 @@ public class Histogram extends StatefulMetric<DistributionDataPoint, Histogram.D
   private static final double[][] NATIVE_BOUNDS;
 
   @Nullable private final ExemplarSamplerConfig exemplarSamplerConfig;
+  @Nullable private final Supplier<Labels> exemplarLabelsSupplier;
 
   // Upper bounds for the classic histogram buckets. Contains at least +Inf.
   // An empty array indicates that this is a native histogram only.
@@ -169,6 +171,7 @@ private Histogram(Histogram.Builder builder, PrometheusProperties prometheusProp
     } else {
       exemplarSamplerConfig = null;
     }
+    exemplarLabelsSupplier = builder.exemplarLabelsSupplier;
   }
 
   @Override
@@ -210,7 +213,7 @@ public class DataPoint implements DistributionDataPoint {
 
     private DataPoint() {
       if (exemplarSamplerConfig != null) {
-        exemplarSampler = new ExemplarSampler(exemplarSamplerConfig);
+        exemplarSampler = new ExemplarSampler(exemplarSamplerConfig, exemplarLabelsSupplier);
       } else {
         exemplarSampler = null;
       }

prometheus-metrics-core/src/main/java/io/prometheus/metrics/core/metrics/StatefulMetric.java

@@ -13,6 +13,7 @@
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /**
@@ -198,11 +199,23 @@ abstract static class Builder<B extends Builder<B, M>, M extends StatefulMetric<
       extends MetricWithFixedMetadata.Builder<B, M> {
 
     @Nullable protected Boolean exemplarsEnabled;
+    @Nullable protected Supplier<Labels> exemplarLabelsSupplier;
 
     protected Builder(List<String> illegalLabelNames, PrometheusProperties config) {
       super(illegalLabelNames, config);
     }
 
+    /**
+     * Provide additional labels to be merged into every automatically-sampled exemplar. The
+     * supplier is called each time an exemplar is sampled, so it can return dynamic values (e.g. a
+     * request-scoped identifier from a thread-local). The supplier is only invoked when a valid,
+     * sampled span context is present; it has no effect when tracing is not active.
+     */
+    public B exemplarLabelsSupplier(Supplier<Labels> supplier) {
+      this.exemplarLabelsSupplier = supplier;
+      return self();
+    }
+
     /** Allow Exemplars for this metric. */
     public B withExemplars() {
       this.exemplarsEnabled = true;

prometheus-metrics-core/src/main/java/io/prometheus/metrics/core/metrics/Summary.java

@@ -19,6 +19,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.DoubleAdder;
 import java.util.concurrent.atomic.LongAdder;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /**
@@ -50,6 +51,7 @@ public class Summary extends StatefulMetric<DistributionDataPoint, Summary.DataP
   private final long maxAgeSeconds;
   private final int ageBuckets;
   @Nullable private final ExemplarSamplerConfig exemplarSamplerConfig;
+  @Nullable private final Supplier<Labels> exemplarLabelsSupplier;
 
   private Summary(Builder builder, PrometheusProperties prometheusProperties) {
     super(builder);
@@ -65,6 +67,7 @@ private Summary(Builder builder, PrometheusProperties prometheusProperties) {
     } else {
       exemplarSamplerConfig = null;
     }
+    exemplarLabelsSupplier = builder.exemplarLabelsSupplier;
   }
 
   private List<CKMSQuantiles.Quantile> makeQuantiles(MetricsProperties[] properties) {
@@ -153,7 +156,7 @@ private DataPoint() {
                 ageBuckets);
       }
       if (exemplarSamplerConfig != null) {
-        exemplarSampler = new ExemplarSampler(exemplarSamplerConfig);
+        exemplarSampler = new ExemplarSampler(exemplarSamplerConfig, exemplarLabelsSupplier);
       } else {
         exemplarSampler = null;
       }

prometheus-metrics-core/src/test/java/io/prometheus/metrics/core/metrics/CounterTest.java

@@ -9,6 +9,7 @@
 import io.prometheus.metrics.config.MetricsProperties;
 import io.prometheus.metrics.config.PrometheusProperties;
 import io.prometheus.metrics.core.exemplars.ExemplarSamplerConfigTestUtil;
+import io.prometheus.metrics.expositionformats.OpenMetricsTextFormatWriter;
 import io.prometheus.metrics.expositionformats.generated.Metrics;
 import io.prometheus.metrics.expositionformats.internal.PrometheusProtobufWriterImpl;
 import io.prometheus.metrics.expositionformats.internal.ProtobufUtil;
@@ -17,9 +18,11 @@
 import io.prometheus.metrics.model.snapshots.Exemplar;
 import io.prometheus.metrics.model.snapshots.Label;
 import io.prometheus.metrics.model.snapshots.Labels;
+import io.prometheus.metrics.model.snapshots.MetricSnapshots;
 import io.prometheus.metrics.model.snapshots.Unit;
 import io.prometheus.metrics.tracer.common.SpanContext;
 import io.prometheus.metrics.tracer.initializer.SpanContextSupplier;
+import java.io.ByteArrayOutputStream;
 import java.util.Arrays;
 import java.util.Iterator;
 import org.junit.jupiter.api.AfterEach;
@@ -321,6 +324,63 @@ void incWithExemplar2() {
         getData(counter).getExemplar());
   }
 
+  @Test
+  void incWithExemplarCustomMetadataInExposition() throws Exception {
+    Counter counter = Counter.builder().name("requests_total").build();
+    counter.incWithExemplar(
+        Labels.of(
+            Exemplar.TRACE_ID, "abc123", Exemplar.SPAN_ID, "def456", "management_id", "mgmt-42"));
+
+    ByteArrayOutputStream out = new ByteArrayOutputStream();
+    new OpenMetricsTextFormatWriter(false, true)
+        .write(out, MetricSnapshots.of(counter.collect()), EscapingScheme.ALLOW_UTF8);
+
+    assertThat(out.toString())
+        .contains("management_id=\"mgmt-42\"")
+        .contains("trace_id=\"abc123\"")
+        .contains("span_id=\"def456\"");
+  }
+
+  @Test
+  void exemplarLabelsSupplierAppearsInAutomaticallySampledExemplar() throws Exception {
+    SpanContextSupplier.setSpanContext(
+        new SpanContext() {
+          @Override
+          public String getCurrentTraceId() {
+            return "trace-abc";
+          }
+
+          @Override
+          public String getCurrentSpanId() {
+            return "span-def";
+          }
+
+          @Override
+          public boolean isCurrentSpanSampled() {
+            return true;
+          }
+
+          @Override
+          public void markCurrentSpanAsExemplar() {}
+        });
+
+    Counter counter =
+        Counter.builder()
+            .name("requests_total")
+            .exemplarLabelsSupplier(() -> Labels.of("management_id", "mgmt-42"))
+            .build();
+    counter.inc(); // automatic sampling path
+
+    ByteArrayOutputStream out = new ByteArrayOutputStream();
+    new OpenMetricsTextFormatWriter(false, true)
+        .write(out, MetricSnapshots.of(counter.collect()), EscapingScheme.ALLOW_UTF8);
+
+    assertThat(out.toString())
+        .contains("management_id=\"mgmt-42\"")
+        .contains("trace_id=\"trace-abc\"")
+        .contains("span_id=\"span-def\"");
+  }
+
   @Test
   void testExemplarSamplerDisabled() {
     Counter counter = Counter.builder().name("count_total").withoutExemplars().build();