Will this break?

DataDog · May 31, 2024 · 1d9ce49 · 1d9ce49
1 parent fc57775
commit 1d9ce49
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 21 deletions.
diff --git a/src/main/java/org/datadog/jmxfetch/App.java b/src/main/java/org/datadog/jmxfetch/App.java
@@ -133,7 +133,10 @@ public App(final AppConfig appConfig) {
         }
         this.configs = getConfigs(this.appConfig);
 
-        this.initTelemetryBean();
+        if (this.appConfig.getJmxfetchTelemetry()) {
+            log.info("Enabling JMX Fetch Telemetry");
+            this.initTelemetryBean();
+        }
     }
 
     private ObjectName getAppTelemetryBeanName() {
@@ -168,14 +171,19 @@ private void initTelemetryBean() {
          | MBeanRegistrationException
          | NotCompliantMBeanException e) {
             log.warn("Could not register bean named '{}' for instance: ",
-                appTelemetryBeanName.toString(), e);
+                appTelemetryBeanName, e);
         }
 
         this.appTelemetry = bean;
         return;
     }
 
     private void teardownTelemetry() {
+        if (!this.appConfig.getJmxfetchTelemetry()) {
+            log.debug("Skipping teardown telemetry as not enabled");
+            return;
+        }
+
         MBeanServer mbs =  ManagementFactory.getPlatformMBeanServer();
         ObjectName appTelemetryBeanName = getAppTelemetryBeanName();
         if (appTelemetryBeanName == null) {
@@ -187,7 +195,7 @@ private void teardownTelemetry() {
             log.debug("Succesfully unregistered app telemetry bean");
         } catch (MBeanRegistrationException | InstanceNotFoundException e) {
             log.warn("Could not unregister bean named '{}' for instance: ",
-                appTelemetryBeanName.toString(), e);
+                appTelemetryBeanName, e);
         }
     }
 
@@ -540,7 +548,7 @@ public void doIteration() {
             for (Instance instance : this.instances) {
                 getMetricsTasks.add(new MetricCollectionTask(instance));
             }
-            if (this.appTelemetry != null) {
+            if (this.appTelemetry != null && this.appConfig.getJmxfetchTelemetry()) {
                 this.appTelemetry.setRunningInstanceCount(this.instances.size());
             }
 
@@ -906,6 +914,7 @@ private Instance instantiate(
     public void init(final boolean forceNewConnection) {
         log.info("Cleaning up instances...");
         this.clearInstances(this.instances);
+        this.instances.clear();
         this.clearInstances(this.brokenInstanceMap.values());
         this.brokenInstanceMap.clear();
 
@@ -1205,7 +1214,7 @@ private <T> void processFixedStatus(
                 this.brokenInstanceMap.remove(instance.toString());
                 this.instances.add(instance);
 
-                if (this.appTelemetry != null) {
+                if (this.appTelemetry != null && this.appConfig.getJmxfetchTelemetry()) {
                     this.appTelemetry.setBrokenInstanceCount(this.brokenInstanceMap.size());
                     this.appTelemetry.setRunningInstanceCount(this.instances.size());
                 }

diff --git a/src/main/java/org/datadog/jmxfetch/Instance.java b/src/main/java/org/datadog/jmxfetch/Instance.java
@@ -240,7 +240,9 @@ public Instance(
             log.info("collect_default_jvm_metrics is false - not collecting default JVM metrics");
         }
 
-        instanceTelemetryBean = createInstanceTelemetryBean();
+        if (this.appConfig.getJmxfetchTelemetry()) {
+            this.instanceTelemetryBean = createInstanceTelemetryBean();
+        }
     }
 
     private ObjectName getObjName(String domain,String instance)
@@ -837,11 +839,14 @@ public boolean isLimitReached() {
     }
 
     private void cleanupTelemetryBean() {
+        if (this.instanceTelemetryBean == null) {
+            return;
+        }
         try {
             mbs.unregisterMBean(instanceTelemetryBeanName);
-            log.debug("Successfully unregistered bean for instance: " + this.getCheckName());
+            log.debug("Successfully unregistered bean for instance: {}", this.getCheckName());
         } catch (MBeanRegistrationException | InstanceNotFoundException e) {
-            log.debug("Unable to unregister bean for instance: " + this.getCheckName());
+            log.debug("Unable to unregister bean for instance: {}", this.getCheckName());
         }
     }
 

diff --git a/src/test/java/org/datadog/jmxfetch/TestApp.java b/src/test/java/org/datadog/jmxfetch/TestApp.java
@@ -2,6 +2,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.when;
 
@@ -22,12 +23,19 @@ public class TestApp extends TestCommon {
     /** Tag metrics with MBean parameters based on user supplied regex */
     @Test
     public void testBeanRegexTags() throws Exception {
+        // When we enable JMXFetch telemetry
+        when(appConfig.getJmxfetchTelemetry()).thenReturn(true);
+
         // We expose a few metrics through JMX
         registerMBean(
                 new SimpleTestJavaApp(),
                 "org.datadog.jmxfetch.test:type=SimpleTestJavaApp,scope=Co|olScope,host=localhost,component=");
         initApplication("jmx_bean_regex_tags.yaml");
 
+        AppTelemetry tlm = app.getAppTelemetryBean();
+        assertNotNull("AppTelemetryBean should not be null", tlm);
+        assertEquals("Before run", 0, tlm.getRunningInstanceCount());
+
         // Run the collection
         run();
 
@@ -45,14 +53,15 @@ public void testBeanRegexTags() throws Exception {
                         "nonRegexTag:value");
 
         assertMetric("this.is.100", tags, 10);
-
-        AppTelemetry tlm = app.getAppTelemetryBean();
-        assertEquals(1, tlm.getRunningInstanceCount());
+        assertEquals("Expecting one instance running", 1, tlm.getRunningInstanceCount());
     }
 
     /** Tag metrics with MBeans parameters. */
     @Test
     public void testBeanTags() throws Exception {
+        // When we enable JMXFetch telemetry
+        when(appConfig.getJmxfetchTelemetry()).thenReturn(true);
+
         // We expose a few metrics through JMX
         registerMBean(
                 new SimpleTestJavaApp(),
@@ -61,10 +70,11 @@ public void testBeanTags() throws Exception {
 
         // Collecting metrics
         run();
-        List<Map<String, Object>> metrics = getMetrics();
 
         // 14 = 13 metrics from java.lang + 1 metric explicitly defined in the yaml config file
-        assertEquals(14, metrics.size());
+        assertEquals(
+            "Expecting 13 metrics from java.lang + 1 metric explicitly defined in the yaml config file",
+            14, getMetrics().size());
 
         List<String> tags =
                 Arrays.asList(
@@ -79,6 +89,7 @@ public void testBeanTags() throws Exception {
         assertMetric("this.is.100", tags, 7);
 
         AppTelemetry tlm = app.getAppTelemetryBean();
+        assertNotNull("AppTelemetryBean should not be null", tlm);
         assertEquals(1, tlm.getRunningInstanceCount());
     }
 
@@ -125,10 +136,9 @@ public void testBeanTagsDontNormalizeParams() throws Exception {
 
         // Collecting metrics
         run();
-        List<Map<String, Object>> metrics = getMetrics();
 
         // 14 = 13 metrics from java.lang + 1 metric explicitly defined in the yaml config file
-        assertEquals(14, metrics.size());
+        assertEquals(14, getMetrics().size());
 
         List<String> tags =
                 Arrays.asList(

diff --git a/src/test/java/org/datadog/jmxfetch/TestCommon.java b/src/test/java/org/datadog/jmxfetch/TestCommon.java
@@ -62,7 +62,7 @@ public class TestCommon {
     AppConfig appConfig = spy(AppConfig.builder().build());
     App app;
     MBeanServer mbs;
-    List<ObjectName> objectNames = new ArrayList<ObjectName>();
+    List<ObjectName> objectNames = new ArrayList<>();
     List<Map<String, Object>> metrics;
     List<Map<String, Object>> serviceChecks;
 
@@ -87,10 +87,12 @@ public static void init() throws Exception {
     protected void registerMBean(Object application, String objectStringName)
             throws InstanceAlreadyExistsException, MBeanRegistrationException,
                     NotCompliantMBeanException, MalformedObjectNameException {
-        mbs = (mbs == null) ? ManagementFactory.getPlatformMBeanServer() : mbs;
+        if(this.mbs == null) {
+            this.mbs = ManagementFactory.getPlatformMBeanServer();
+        }
         ObjectName objectName = new ObjectName(objectStringName);
+        this.mbs.registerMBean(application, objectName);
         objectNames.add(objectName);
-        mbs.registerMBean(application, objectName);
     }
 
     /**
@@ -100,10 +102,11 @@ protected void registerMBean(Object application, String objectStringName)
      * @throws MBeanRegistrationException
      */
     public void unregisterMBeans() throws MBeanRegistrationException, InstanceNotFoundException {
-        if (mbs != null) {
-            for (ObjectName objectName : objectNames) {
-                mbs.unregisterMBean(objectName);
+        if (this.mbs != null) {
+            for (ObjectName objectName : this.objectNames) {
+                this.mbs.unregisterMBean(objectName);
             }
+            this.objectNames.clear();
         }
     }