prometheus-emitter: add extraLabels parameter (apache#14728)

* prometheus-emitter: add extraLabels parameter

* prometheus-emitter: update readme to include the extraLabels parameter

* prometheus-emitter: remove nullable and surface label name issues

* remove import to make linter happy

docs/development/extensions-contrib/prometheus.md

@@ -45,7 +45,8 @@ All the configuration parameters for the Prometheus emitter are under `druid.emi
 | `druid.emitter.prometheus.addHostAsLabel`     | Flag to include the hostname as a prometheus label.                                                                                                                                                                                    | no        | false                                |
 | `druid.emitter.prometheus.addServiceAsLabel`  | Flag to include the druid service name (e.g. `druid/broker`, `druid/coordinator`, etc.) as a prometheus label.                                                                                                                         | no        | false                                |
 | `druid.emitter.prometheus.pushGatewayAddress` | Pushgateway address. Required if using `pushgateway` strategy.                                                                                                                                                                         | no        | none                                 |
-|`druid.emitter.prometheus.flushPeriod`|Emit metrics to Pushgateway every `flushPeriod` seconds. Required if `pushgateway` strategy is used.|no|15|
+| `druid.emitter.prometheus.flushPeriod`        | Emit metrics to Pushgateway every `flushPeriod` seconds. Required if `pushgateway` strategy is used.                                                                                                                                   | no        | 15                                   |
+| `druid.emitter.prometheus.extraLabels`        | JSON key-value pairs for additional labels on all metrics. Keys (label names) must match the regex `[a-zA-Z_:][a-zA-Z0-9_:]*`. Example: `{"cluster_name": "druid_cluster1", "env": "staging"}`.                                        | no        | none                                 |
 
 ### Ports for colocated Druid processes
 

extensions-contrib/prometheus-emitter/src/main/java/org/apache/druid/emitter/prometheus/Metrics.java

@@ -62,10 +62,15 @@ public DimensionsAndCollector getByName(String name, String service)
     }
   }
 
-  public Metrics(String namespace, String path, boolean isAddHostAsLabel, boolean isAddServiceAsLabel)
+  public Metrics(String namespace, String path, boolean isAddHostAsLabel, boolean isAddServiceAsLabel, Map<String, String> extraLabels)
   {
     Map<String, DimensionsAndCollector> registeredMetrics = new HashMap<>();
     Map<String, Metric> metrics = readConfig(path);
+
+    if (extraLabels == null) {
+      extraLabels = Collections.emptyMap(); // Avoid null checks later
+    }
+
     for (Map.Entry<String, Metric> entry : metrics.entrySet()) {
       String name = entry.getKey();
       Metric metric = entry.getValue();
@@ -79,6 +84,8 @@ public Metrics(String namespace, String path, boolean isAddHostAsLabel, boolean
         metric.dimensions.add(TAG_SERVICE);
       }
 
+      metric.dimensions.addAll(extraLabels.keySet());
+
       String[] dimensions = metric.dimensions.toArray(new String[0]);
       String formattedName = PATTERN.matcher(StringUtils.toLowerCase(name)).replaceAll("_");
       SimpleCollector collector = null;

extensions-contrib/prometheus-emitter/src/main/java/org/apache/druid/emitter/prometheus/PrometheusEmitter.java

@@ -69,7 +69,7 @@ public PrometheusEmitter(PrometheusEmitterConfig config)
   {
     this.config = config;
     this.strategy = config.getStrategy();
-    metrics = new Metrics(config.getNamespace(), config.getDimensionMapPath(), config.isAddHostAsLabel(), config.isAddServiceAsLabel());
+    metrics = new Metrics(config.getNamespace(), config.getDimensionMapPath(), config.isAddHostAsLabel(), config.isAddServiceAsLabel(), config.getExtraLabels());
   }
 
 
@@ -136,6 +136,9 @@ private void emitMetric(ServiceMetricEvent metricEvent)
     if (metric != null) {
       String[] labelValues = new String[metric.getDimensions().length];
       String[] labelNames = metric.getDimensions();
+
+      Map<String, String> extraLabels = config.getExtraLabels();
+
       for (int i = 0; i < labelValues.length; i++) {
         String labelName = labelNames[i];
         //labelName is controlled by the user. Instead of potential NPE on invalid labelName we use "unknown" as the dimension value
@@ -148,6 +151,8 @@ private void emitMetric(ServiceMetricEvent metricEvent)
             labelValues[i] = host;
           } else if (config.isAddServiceAsLabel() && TAG_SERVICE.equals(labelName)) {
             labelValues[i] = service;
+          } else if (extraLabels.containsKey(labelName)) {
+            labelValues[i] = config.getExtraLabels().get(labelName);
           } else {
             labelValues[i] = "unknown";
           }

extensions-contrib/prometheus-emitter/src/main/java/org/apache/druid/emitter/prometheus/PrometheusEmitterConfig.java

@@ -22,8 +22,12 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Preconditions;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.java.util.common.StringUtils;
 
 import javax.annotation.Nullable;
+import java.util.Collections;
+import java.util.Map;
 import java.util.Objects;
 import java.util.regex.Pattern;
 
@@ -63,6 +67,9 @@ public class PrometheusEmitterConfig
   @JsonProperty
   private final boolean addServiceAsLabel;
 
+  @JsonProperty
+  private final Map<String, String> extraLabels;
+
   @JsonCreator
   public PrometheusEmitterConfig(
       @JsonProperty("strategy") @Nullable Strategy strategy,
@@ -72,7 +79,8 @@ public PrometheusEmitterConfig(
       @JsonProperty("pushGatewayAddress") @Nullable String pushGatewayAddress,
       @JsonProperty("addHostAsLabel") boolean addHostAsLabel,
       @JsonProperty("addServiceAsLabel") boolean addServiceAsLabel,
-      @JsonProperty("flushPeriod") Integer flushPeriod
+      @JsonProperty("flushPeriod") Integer flushPeriod,
+      @JsonProperty("extraLabels") @Nullable Map<String, String> extraLabels
   )
   {
     this.strategy = strategy != null ? strategy : Strategy.exporter;
@@ -94,6 +102,21 @@ public PrometheusEmitterConfig(
     this.flushPeriod = flushPeriod;
     this.addHostAsLabel = addHostAsLabel;
     this.addServiceAsLabel = addServiceAsLabel;
+    this.extraLabels = extraLabels != null ? extraLabels : Collections.emptyMap();
+    // Validate label names early to prevent Prometheus exceptions later.
+    for (String key : this.extraLabels.keySet()) {
+      if (!PATTERN.matcher(key).matches()) {
+        throw DruidException.forPersona(DruidException.Persona.OPERATOR)
+                            .ofCategory(DruidException.Category.INVALID_INPUT)
+                            .build(
+                                StringUtils.format(
+                                    "Invalid metric label name [%s]. Label names must conform to the pattern [%s].",
+                                    key,
+                                    PATTERN.pattern()
+                                )
+                            );
+      }
+    }
   }
 
   public String getNamespace()
@@ -137,6 +160,11 @@ public boolean isAddServiceAsLabel()
     return addServiceAsLabel;
   }
 
+  public Map<String, String> getExtraLabels()
+  {
+    return extraLabels;
+  }
+
   public enum Strategy
   {
     exporter, pushgateway

extensions-contrib/prometheus-emitter/src/test/java/org/apache/druid/emitter/prometheus/MetricsTest.java

@@ -23,12 +23,15 @@
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.util.HashMap;
+import java.util.Map;
+
 public class MetricsTest
 {
   @Test
   public void testMetricsConfiguration()
   {
-    Metrics metrics = new Metrics("test", null, true, true);
+    Metrics metrics = new Metrics("test", null, true, true, null);
     DimensionsAndCollector dimensionsAndCollector = metrics.getByName("query/time", "historical");
     Assert.assertNotNull(dimensionsAndCollector);
     String[] dimensions = dimensionsAndCollector.getDimensions();
@@ -46,4 +49,48 @@ public void testMetricsConfiguration()
     Assert.assertEquals("host_name", dims[1]);
     Assert.assertEquals("server", dims[2]);
   }
+
+  @Test
+  public void testMetricsConfigurationWithExtraLabels()
+  {
+    Map<String, String> extraLabels = new HashMap<>();
+    extraLabels.put("extra_label", "value");
+
+    Metrics metrics = new Metrics("test_2", null, true, true, extraLabels);
+    DimensionsAndCollector dimensionsAndCollector = metrics.getByName("query/time", "historical");
+    Assert.assertNotNull(dimensionsAndCollector);
+    String[] dimensions = dimensionsAndCollector.getDimensions();
+    Assert.assertEquals("dataSource", dimensions[0]);
+    Assert.assertEquals("druid_service", dimensions[1]);
+    Assert.assertEquals("extra_label", dimensions[2]);
+    Assert.assertEquals("host_name", dimensions[3]);
+    Assert.assertEquals("type", dimensions[4]);
+    Assert.assertEquals(1000.0, dimensionsAndCollector.getConversionFactor(), 0.0);
+    Assert.assertTrue(dimensionsAndCollector.getCollector() instanceof Histogram);
+
+    DimensionsAndCollector d = metrics.getByName("segment/loadQueue/count", "historical");
+    Assert.assertNotNull(d);
+    String[] dims = d.getDimensions();
+    Assert.assertEquals("druid_service", dims[0]);
+    Assert.assertEquals("extra_label", dims[1]);
+    Assert.assertEquals("host_name", dims[2]);
+    Assert.assertEquals("server", dims[3]);
+  }
+
+  @Test
+  public void testMetricsConfigurationWithBadExtraLabels()
+  {
+    Map<String, String> extraLabels = new HashMap<>();
+    extraLabels.put("extra label", "value");
+
+    // Expect an exception thrown by Prometheus code due to invalid metric label
+    Exception exception = Assert.assertThrows(IllegalArgumentException.class, () -> {
+      new Metrics("test_3", null, true, true, extraLabels);
+    });
+
+    String expectedMessage = "Invalid metric label name: extra label";
+    String actualMessage = exception.getMessage();
+
+    Assert.assertTrue(actualMessage.contains(expectedMessage));
+  }
 }

extensions-contrib/prometheus-emitter/src/test/java/org/apache/druid/emitter/prometheus/PrometheusEmitterConfigTest.java

@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.emitter.prometheus;
+
+import io.prometheus.client.CollectorRegistry;
+import org.apache.druid.error.DruidException;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class PrometheusEmitterConfigTest
+{
+  @Test
+  public void testEmitterConfigWithBadExtraLabels()
+  {
+    CollectorRegistry.defaultRegistry.clear();
+
+    Map<String, String> extraLabels = new HashMap<>();
+    extraLabels.put("label Name", "label Value");
+
+    // Expect an exception thrown by our own PrometheusEmitterConfig due to invalid label key
+    Exception exception = Assert.assertThrows(DruidException.class, () -> {
+      new PrometheusEmitterConfig(PrometheusEmitterConfig.Strategy.exporter, null, null, 0, null, false, true, 60, extraLabels);
+    });
+
+    String expectedMessage = "Invalid metric label name [label Name]. Label names must conform to the pattern [[a-zA-Z_:][a-zA-Z0-9_:]*]";
+    String actualMessage = exception.getMessage();
+
+    Assert.assertTrue(actualMessage.contains(expectedMessage));
+  }
+
+}