WIP

apache · Oct 16, 2024 · 1b69b17 · 1b69b17
1 parent d581be8
commit 1b69b17
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 61 deletions.
diff --git a/pinot-common/src/test/java/org/apache/pinot/common/metrics/BrokerJMXToPromMetricsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/metrics/BrokerJMXToPromMetricsTest.java
@@ -22,6 +22,7 @@
 import com.yammer.metrics.core.MetricsRegistry;
 import com.yammer.metrics.reporting.JmxReporter;
 import io.prometheus.jmx.JavaAgent;
+import io.prometheus.jmx.shaded.io.prometheus.client.exporter.HTTPServer;
 import java.net.URI;
 import java.util.List;
 import java.util.Random;
@@ -49,15 +50,13 @@ public class BrokerJMXToPromMetricsTest extends PinotJMXToPromMetricsTest {
 
   private int _exporterPort;
 
+  private HTTPServer _httpServer;
+
   @BeforeClass
   public void setup()
       throws Exception {
 
-    _exporterPort = 9000 + new Random().nextInt(1000);
-    String agentArgs = String.format("%s:%s", _exporterPort,
-        "../docker/images/pinot/etc/jmx_prometheus_javaagent/configs/broker.yml");
-
-    JavaAgent.agentmain(agentArgs, null);
+    _httpServer = startExporter(PinotComponent.BROKER);
 
     PinotConfiguration pinotConfiguration = new PinotConfiguration();
     pinotConfiguration.setProperty(CONFIG_OF_METRICS_FACTORY_CLASS_NAME,
@@ -159,7 +158,7 @@ public void brokerMeterTest() {
   @Override
   protected SimpleHttpResponse getExportedPromMetrics() {
     try {
-      return _httpClient.sendGetRequest(new URI("http://localhost:" + _exporterPort + "/metrics"));
+      return _httpClient.sendGetRequest(new URI("http://localhost:" + _httpServer.getPort() + "/metrics"));
     } catch (Exception e) {
       throw new RuntimeException(e);
     }

diff --git a/...-common/src/test/java/org/apache/pinot/common/metrics/ControllerJMXToPromMetricsTest.java b/...-common/src/test/java/org/apache/pinot/common/metrics/ControllerJMXToPromMetricsTest.java
@@ -21,22 +21,11 @@
 
 import com.yammer.metrics.core.MetricsRegistry;
 import com.yammer.metrics.reporting.JmxReporter;
-import io.prometheus.jmx.BuildInfoCollector;
-import io.prometheus.jmx.JavaAgent;
-import io.prometheus.jmx.JmxCollector;
-import io.prometheus.jmx.common.http.ConfigurationException;
-import io.prometheus.jmx.common.http.HTTPServerFactory;
-import io.prometheus.jmx.shaded.io.prometheus.client.CollectorRegistry;
 import io.prometheus.jmx.shaded.io.prometheus.client.exporter.HTTPServer;
-import io.prometheus.jmx.shaded.io.prometheus.client.hotspot.DefaultExports;
-import java.io.File;
-import java.io.IOException;
 import java.net.URI;
-import java.net.URISyntaxException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
-import java.util.Random;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Stream;
 import org.apache.commons.lang3.StringUtils;
@@ -66,25 +55,7 @@ public class ControllerJMXToPromMetricsTest extends PinotJMXToPromMetricsTest {
   public void setup()
       throws Exception {
 
-    _exporterPort = 9000 + new Random().nextInt(1000);
-    String agentArgs = String.format("%s:%s", _exporterPort,
-        "../docker/images/pinot/etc/jmx_prometheus_javaagent/configs/controller.yml");
-
-    String host = "0.0.0.0";
-
-    try {
-      JMXExporterConfig config = parseConfig(agentArgs, host);
-      CollectorRegistry registry = new CollectorRegistry();
-      (new JmxCollector(new File(config.file), JmxCollector.Mode.AGENT)).register(registry);
-      DefaultExports.register(registry);
-      _httpServer = (new HTTPServerFactory()).createHTTPServer(config.socket, registry, true, new File(config.file));
-    } catch (ConfigurationException var4) {
-      System.err.println("Configuration Exception : " + var4.getMessage());
-      System.exit(1);
-    } catch (IllegalArgumentException var5) {
-      System.err.println("Usage: -javaagent:/path/to/JavaAgent.jar=[host:]<port>:<yaml configuration file> " + var5.getMessage());
-      System.exit(1);
-    }
+    _httpServer = startExporter(PinotComponent.CONTROLLER);
 
     PinotConfiguration pinotConfiguration = new PinotConfiguration();
     pinotConfiguration.setProperty(CONFIG_OF_METRICS_FACTORY_CLASS_NAME,
@@ -259,7 +230,7 @@ private static String getStrippedMetricName(ControllerGauge controllerGauge) {
   @Override
   protected SimpleHttpResponse getExportedPromMetrics() {
     try {
-      return _httpClient.sendGetRequest(new URI("http://localhost:" + _exporterPort + "/metrics"));
+      return _httpClient.sendGetRequest(new URI("http://localhost:" + _httpServer.getPort() + "/metrics"));
     } catch (Exception e) {
       throw new RuntimeException(e);
     }

diff --git a/pinot-common/src/test/java/org/apache/pinot/common/metrics/PinotJMXToPromMetricsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/metrics/PinotJMXToPromMetricsTest.java
@@ -20,13 +20,16 @@
 package org.apache.pinot.common.metrics;
 
 import com.google.common.base.Objects;
-import io.prometheus.jmx.JavaAgent;
+import io.prometheus.jmx.JmxCollector;
+import io.prometheus.jmx.common.http.HTTPServerFactory;
+import io.prometheus.jmx.shaded.io.prometheus.client.CollectorRegistry;
+import io.prometheus.jmx.shaded.io.prometheus.client.exporter.HTTPServer;
+import io.prometheus.jmx.shaded.io.prometheus.client.hotspot.DefaultExports;
 import java.io.BufferedReader;
+import java.io.File;
 import java.io.IOException;
 import java.io.StringReader;
 import java.net.InetSocketAddress;
-import java.net.URI;
-import java.net.URISyntaxException;
 import java.util.ArrayList;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -44,6 +47,12 @@
 
 public abstract class PinotJMXToPromMetricsTest {
 
+  private static final Map<PinotComponent, String> PINOT_COMPONENT_CONFIG_FILE_MAP =
+      Map.of(PinotComponent.CONTROLLER, "controller.yml", PinotComponent.SERVER, "server.yml", PinotComponent.MINION,
+          "minion.yml", PinotComponent.BROKER, "broker.yml");
+
+  private static final String CONFIG_DIR = "../docker/images/pinot/etc/jmx_prometheus_javaagent/configs";
+
   protected HttpClient _httpClient;
 
   protected static final String LABEL_KEY_TABLE = "table";
@@ -90,6 +99,20 @@ public abstract class PinotJMXToPromMetricsTest {
   protected static final List<String> EXPORTED_LABELS_PERIODIC_TASK_TABLE_TABLETYPE =
       List.of("periodicTask", "ClusterHealthCheck", "table", "myTable", "tableType", "REALTIME");
 
+  public HTTPServer startExporter(PinotComponent pinotComponent) {
+    String args =
+        String.format("%s:%s/%s", 0, CONFIG_DIR, PINOT_COMPONENT_CONFIG_FILE_MAP.get(pinotComponent));
+      try {
+        JMXExporterConfig config = parseConfig(args, "0.0.0.0");
+        CollectorRegistry registry = new CollectorRegistry();
+        (new JmxCollector(new File(config.file), JmxCollector.Mode.AGENT)).register(registry);
+        DefaultExports.register(registry);
+        return (new HTTPServerFactory()).createHTTPServer(config.socket, registry, true, new File(config.file));
+      } catch (Exception e) {
+        throw new RuntimeException(e);
+      }
+  }
+
   protected void assertGaugeExportedCorrectly(String exportedGaugePrefix, String exportedMetricPrefix) {
     List<PromMetric> promMetrics;
     try {
@@ -315,4 +338,8 @@ public static class JMXExporterConfig {
       this.socket = socket;
     }
   }
+
+  public enum PinotComponent {
+    SERVER, BROKER, CONTROLLER, MINION
+  }
 }
diff --git a/pinot-common/src/test/java/org/apache/pinot/common/metrics/ServerJMXToPromMetricsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/metrics/ServerJMXToPromMetricsTest.java
@@ -60,26 +60,7 @@ public class ServerJMXToPromMetricsTest extends PinotJMXToPromMetricsTest {
   public void setup()
       throws Exception {
 
-    _exporterPort = 9000 + new Random().nextInt(1000);
-    String agentArgs = String.format("%s:%s", _exporterPort,
-        "../docker/images/pinot/etc/jmx_prometheus_javaagent/configs/server.yml");
-
-    String host = "0.0.0.0";
-
-    try {
-      JMXExporterConfig config = parseConfig(agentArgs, host);
-      CollectorRegistry registry = new CollectorRegistry();
-      (new JmxCollector(new File(config.file), JmxCollector.Mode.AGENT)).register(registry);
-      DefaultExports.register(registry);
-      _httpServer = (new HTTPServerFactory()).createHTTPServer(config.socket, registry, true, new File(config.file));
-    } catch (ConfigurationException var4) {
-      System.err.println("Configuration Exception : " + var4.getMessage());
-      System.exit(1);
-    } catch (IllegalArgumentException var5) {
-      System.err.println("Usage: -javaagent:/path/to/JavaAgent.jar=[host:]<port>:<yaml configuration file> " + var5.getMessage());
-      System.exit(1);
-    }
-
+    _httpServer = startExporter(PinotComponent.SERVER);
 
     PinotConfiguration pinotConfiguration = new PinotConfiguration();
     pinotConfiguration.setProperty(CONFIG_OF_METRICS_FACTORY_CLASS_NAME,
@@ -268,7 +249,7 @@ private void assertMeterExportedCorrectly(String exportedMeterName, List<String>
   @Override
   protected SimpleHttpResponse getExportedPromMetrics() {
     try {
-      return _httpClient.sendGetRequest(new URI("http://localhost:" + _exporterPort + "/metrics"));
+      return _httpClient.sendGetRequest(new URI("http://localhost:" + _httpServer.getPort() + "/metrics"));
     } catch (Exception e) {
       throw new RuntimeException(e);
     }