add exmplar supplier

Signed-off-by: Jay DeLuca <jaydeluca4@gmail.com>

Author: jaydeluca

docs/apidiffs/current_vs_latest/prometheus-metrics-core.txt

@@ -75,3 +75,43 @@ The [examples/example-exemplar-tail-sampling/](https://github.com/prometheus/cli
 directory has a complete end-to-end example, with a distributed Java application with two services,
 an OpenTelemetry collector, Prometheus, Tempo as a trace database, and Grafana dashboards. Use
 docker-compose as described in the example's readme to run the example and explore the results.
+
+## Adding custom labels to exemplars
+
+Automatically-sampled exemplars carry the `trace_id` and `span_id` labels. You can attach
+additional, custom labels (for example an internal identifier) to every automatically-sampled
+exemplar. There are two options.
+
+### Global (all metrics)
+
+Register a global supplier to add custom labels to the exemplars of _all_ metrics, including
+metrics registered by third-party libraries that you do not control. This is the right option when
+you cannot modify the code that creates the metric:
+
+```java
+ExemplarLabelsSupplier.setExemplarLabelsSupplier(
+    () -> Labels.of("management_id", currentManagementId()));
+```
+
+### Per metric
+
+If you only want the extra labels on a specific metric you define yourself, use the builder:
+
+```java
+Counter counter =
+    Counter.builder()
+        .name("requests_total")
+        .exemplarLabelsSupplier(() -> Labels.of("management_id", currentManagementId()))
+        .build();
+```
+
+### Notes
+
+- The supplier is invoked on the (rate-limited) hot path each time an exemplar is sampled, so it
+  should be cheap. It may return dynamic, request-scoped values (e.g. read from a thread-local).
+- Custom labels are only added when a valid, sampled span context is present; the supplier never
+  causes an exemplar to be created on its own.
+- Precedence on a label-name collision: the reserved `trace_id`/`span_id` labels always win, then
+  the per-metric supplier, then the global supplier. Colliding labels are silently dropped.
+- If the supplier throws, the exception is swallowed and the exemplar is created without the
+  additional labels, so a misbehaving supplier never breaks metric collection.

docs/content/otel/tracing.md

@@ -0,0 +1,55 @@
+package io.prometheus.metrics.core.exemplars;
+
+import io.prometheus.metrics.annotations.StableApi;
+import io.prometheus.metrics.model.snapshots.Labels;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+
+/**
+ * Global holder for a {@link Supplier} of additional {@link Labels} that are merged into every
+ * automatically-sampled Exemplar across the entire application.
+ *
+ * <p>This is the global counterpart to the per-metric {@code exemplarLabelsSupplier(...)} builder
+ * method. Registering a supplier here affects <em>all</em> metrics, including metrics registered by
+ * third-party libraries that the application does not control. This makes it the right tool when
+ * you cannot modify the code that creates the metrics.
+ *
+ * <p>The supplier is invoked on the metric hot path (rate-limited by the exemplar sampler), each
+ * time an Exemplar is sampled from a valid, sampled span context. It should therefore be cheap and
+ * non-blocking. It may return dynamic, request-scoped values, for example an identifier read from a
+ * thread-local:
+ *
+ * <pre>{@code
+ * ExemplarLabelsSupplier.setExemplarLabelsSupplier(
+ *     () -> Labels.of("management_id", currentManagementId()));
+ * }</pre>
+ *
+ * <p>Labels returned by the supplier that collide with {@code trace_id}/{@code span_id} (or, when a
+ * per-metric supplier is also configured, with that supplier's labels) are silently dropped rather
+ * than causing an error: the per-metric supplier takes precedence over the global one, and the
+ * reserved {@code trace_id}/{@code span_id} labels always win. If the supplier throws, the
+ * exception is swallowed and the Exemplar is created without the additional labels, so a
+ * misbehaving supplier never breaks metric collection.
+ */
+@StableApi
+public class ExemplarLabelsSupplier {
+
+  private static final AtomicReference<Supplier<Labels>> supplierRef = new AtomicReference<>();
+
+  private ExemplarLabelsSupplier() {}
+
+  /**
+   * Register a global supplier of additional exemplar labels. Pass {@code null} to remove a
+   * previously registered supplier. The most recently registered supplier wins.
+   */
+  public static void setExemplarLabelsSupplier(@Nullable Supplier<Labels> supplier) {
+    supplierRef.set(supplier);
+  }
+
+  /** Returns the registered global supplier, or {@code null} if none has been set. */
+  @Nullable
+  public static Supplier<Labels> getExemplarLabelsSupplier() {
+    return supplierRef.get();
+  }
+}

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

@@ -13,6 +13,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;
 
 /**
@@ -47,8 +48,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);
   }
 
   /**
@@ -60,10 +63,24 @@ 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 additionally accepts a supplier of 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 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() {
@@ -322,7 +339,7 @@ private long durationUntilNextExemplarExpires(long now) {
 
   private long updateCustomExemplar(int index, double value, Labels labels, long now) {
     if (!labels.contains(Exemplar.TRACE_ID) && !labels.contains(Exemplar.SPAN_ID)) {
-      labels = labels.merge(doSampleExemplar());
+      labels = mergeLabels(labels, doSampleExemplar());
     }
     customExemplars[index] =
         Exemplar.builder().value(value).labels(labels).timestampMillis(now).build();
@@ -357,7 +374,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 labels = Labels.of(Exemplar.TRACE_ID, traceId, Exemplar.SPAN_ID, spanId);
+            // Per-metric supplier first (more specific), then the global supplier. On a name
+            // collision the earlier (more specific) value is kept; the reserved trace_id/span_id
+            // labels always win over both.
+            labels = mergeAdditionalLabels(labels, additionalLabelsSupplier);
+            labels =
+                mergeAdditionalLabels(labels, ExemplarLabelsSupplier.getExemplarLabelsSupplier());
+            return labels;
           }
         }
       }
@@ -366,4 +390,68 @@ private Labels doSampleExemplar() {
     }
     return Labels.EMPTY;
   }
+
+  /**
+   * Merge labels from {@code supplier} into {@code base}, dropping any label whose name already
+   * exists in {@code base}. Never throws: a {@code null} supplier, a {@code null}/empty result, a
+   * colliding label name, or an exception thrown by the supplier all result in {@code base} being
+   * returned unchanged (minus the offending labels). A misbehaving supplier must never break metric
+   * collection.
+   */
+  private static Labels mergeAdditionalLabels(Labels base, @Nullable Supplier<Labels> supplier) {
+    if (supplier == null) {
+      return base;
+    }
+    Labels extra;
+    try {
+      extra = supplier.get();
+    } catch (Throwable ignored) {
+      // A misbehaving supplier (any RuntimeException or Error) must never break metric collection.
+      return base;
+    }
+    if (extra == null || extra.isEmpty()) {
+      return base;
+    }
+    return mergeLabels(base, extra);
+  }
+
+  /**
+   * Merge {@code extra} into {@code base}, dropping any label whose name already exists in {@code
+   * base}.
+   */
+  private static Labels mergeLabels(Labels base, Labels extra) {
+    if (extra.isEmpty()) {
+      return base;
+    }
+    // Count name collisions with base in a single pass so we can merge exactly once below: base
+    // (trace_id/span_id and any more-specific supplier) always wins, so colliding labels are
+    // dropped. extra is itself a valid Labels (no internal duplicates), so the surviving labels
+    // never collide with each other and merge() cannot throw on a duplicate name.
+    int size = extra.size();
+    int collisions = 0;
+    for (int i = 0; i < size; i++) {
+      if (base.contains(extra.getName(i))) {
+        collisions++;
+      }
+    }
+    if (collisions == 0) {
+      return base.merge(extra);
+    }
+    if (collisions == size) {
+      return base;
+    }
+    int kept = size - collisions;
+    String[] names = new String[kept];
+    String[] values = new String[kept];
+    int j = 0;
+    for (int i = 0; i < size; i++) {
+      String name = extra.getName(i);
+      if (!base.contains(name)) {
+        names[j] = name;
+        values[j] = extra.getValue(i);
+        j++;
+      }
+    }
+    return base.merge(names, values);
+  }
 }

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

@@ -15,6 +15,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;
 
 /**
@@ -37,6 +38,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);
@@ -49,6 +51,7 @@ private Counter(Builder builder, PrometheusProperties prometheusProperties) {
     } else {
       exemplarSamplerConfig = null;
     }
+    exemplarLabelsSupplier = builder.exemplarLabelsSupplier;
   }
 
   @Override
@@ -108,7 +111,8 @@ public MetricType getMetricType() {
   @Override
   protected DataPoint newDataPoint() {
     if (exemplarSamplerConfig != null) {
-      return new DataPoint(new ExemplarSampler(exemplarSamplerConfig));
+      return new DataPoint(
+          new ExemplarSampler(exemplarSamplerConfig, null, exemplarLabelsSupplier));
     } else {
       return new DataPoint(null);
     }

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

@@ -14,6 +14,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /**
@@ -44,6 +45,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);
@@ -56,6 +58,7 @@ private Gauge(Builder builder, PrometheusProperties prometheusProperties) {
     } else {
       exemplarSamplerConfig = null;
     }
+    exemplarLabelsSupplier = builder.exemplarLabelsSupplier;
   }
 
   @Override
@@ -110,7 +113,8 @@ public MetricType getMetricType() {
   @Override
   protected DataPoint newDataPoint() {
     if (exemplarSamplerConfig != null) {
-      return new DataPoint(new ExemplarSampler(exemplarSamplerConfig));
+      return new DataPoint(
+          new ExemplarSampler(exemplarSamplerConfig, null, exemplarLabelsSupplier));
     } else {
       return new DataPoint(null);
     }

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

@@ -26,6 +26,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;
 
 /**
@@ -73,6 +74,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.
@@ -171,6 +173,7 @@ private Histogram(Histogram.Builder builder, PrometheusProperties prometheusProp
     } else {
       exemplarSamplerConfig = null;
     }
+    exemplarLabelsSupplier = builder.exemplarLabelsSupplier;
   }
 
   @Override
@@ -212,7 +215,7 @@ public class DataPoint implements DistributionDataPoint {
 
     private DataPoint() {
       if (exemplarSamplerConfig != null) {
-        exemplarSampler = new ExemplarSampler(exemplarSamplerConfig);
+        exemplarSampler = new ExemplarSampler(exemplarSamplerConfig, null, exemplarLabelsSupplier);
       } else {
         exemplarSampler = null;
       }

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

@@ -14,6 +14,7 @@
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /**
@@ -200,11 +201,31 @@ 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 of <em>this
+     * metric</em>. 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.
+     *
+     * <p>For a global supplier that applies to all metrics (including metrics registered by
+     * third-party libraries you do not control), see {@link
+     * io.prometheus.metrics.core.exemplars.ExemplarLabelsSupplier}. When both are configured, this
+     * per-metric supplier takes precedence over the global one on a label-name collision, and the
+     * reserved {@code trace_id}/{@code span_id} labels always win over both. Labels that collide
+     * are silently dropped.
+     */
+    public B exemplarLabelsSupplier(Supplier<Labels> supplier) {
+      this.exemplarLabelsSupplier = supplier;
+      return self();
+    }
+
     /** Allow Exemplars for this metric. */
     public B withExemplars() {
       this.exemplarsEnabled = true;