fix: keep custom exemplar labels caller-controlled (#2194)

Author: zeitlinger

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

@@ -339,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 = mergeLabels(labels, doSampleExemplar());
+      labels = mergeLabels(labels, sampleTraceContextLabels());
     }
     customExemplars[index] =
         Exemplar.builder().value(value).labels(labels).timestampMillis(now).build();
@@ -358,6 +358,19 @@ private long updateExemplar(int index, double value, long now) {
   }
 
   private Labels doSampleExemplar() {
+    Labels labels = sampleTraceContextLabels();
+    if (labels.isEmpty()) {
+      return labels;
+    }
+    // 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;
+  }
+
+  private Labels sampleTraceContextLabels() {
     // Using the qualified name so that Micrometer can exclude the dependency on
     // prometheus-metrics-tracer-initializer
     // as they provide their own implementation of SpanContextSupplier.
@@ -374,14 +387,7 @@ private Labels doSampleExemplar() {
           String traceId = spanContext.getCurrentTraceId();
           if (spanId != null && traceId != null) {
             spanContext.markCurrentSpanAsExemplar();
-            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;
+            return Labels.of(Exemplar.TRACE_ID, traceId, Exemplar.SPAN_ID, spanId);
           }
         }
       }

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

@@ -502,22 +502,22 @@ void globalSupplierWithoutSpanContextProducesNoExemplar() {
   }
 
   @Test
-  void globalSupplierMergedIntoCustomExemplar() throws Exception {
+  void globalSupplierDoesNotApplyToCustomExemplar() throws Exception {
     SpanContextSupplier.setSpanContext(sampledSpanContext("trace-abc", "span-def"));
     ExemplarLabelsSupplier.setExemplarLabelsSupplier(() -> Labels.of("management_id", "mgmt-42"));
 
     Counter counter = Counter.builder().name("requests_total").build();
     counter.incWithExemplar(Labels.of("k", "v"));
 
     assertThat(openMetrics(counter))
-        .contains("management_id=\"mgmt-42\"")
         .contains("k=\"v\"")
         .contains("trace_id=\"trace-abc\"")
-        .contains("span_id=\"span-def\"");
+        .contains("span_id=\"span-def\"")
+        .doesNotContain("management_id=\"mgmt-42\"");
   }
 
   @Test
-  void callerLabelsWinOverGlobalSupplierInCustomExemplar() throws Exception {
+  void callerControlsCustomExemplarLabels() throws Exception {
     SpanContextSupplier.setSpanContext(sampledSpanContext("trace-abc", "span-def"));
     ExemplarLabelsSupplier.setExemplarLabelsSupplier(
         () -> Labels.of("k", "global", "management_id", "mgmt-42"));
@@ -527,10 +527,10 @@ void callerLabelsWinOverGlobalSupplierInCustomExemplar() throws Exception {
 
     assertThat(openMetrics(counter))
         .contains("k=\"caller\"")
-        .contains("management_id=\"mgmt-42\"")
         .contains("trace_id=\"trace-abc\"")
         .contains("span_id=\"span-def\"")
-        .doesNotContain("k=\"global\"");
+        .doesNotContain("k=\"global\"")
+        .doesNotContain("management_id=\"mgmt-42\"");
   }
 
   @Test