From bac2d6b62187c174291f358a59631b4b337cdc68 Mon Sep 17 00:00:00 2001
From: Sylvain Juge <763082+SylvainJuge@users.noreply.github.com>
Date: Thu, 8 Feb 2024 09:46:49 +0100
Subject: [PATCH] test jmx retry

---
 .../apm/agent/jmx/JmxMetricTracker.java       | 17 +++++---
 .../apm/agent/jmx/JmxMetricTrackerTest.java   | 40 ++++++++++++++++++-
 2 files changed, 50 insertions(+), 7 deletions(-)

diff --git a/apm-agent-plugins/apm-jmx-plugin/src/main/java/co/elastic/apm/agent/jmx/JmxMetricTracker.java b/apm-agent-plugins/apm-jmx-plugin/src/main/java/co/elastic/apm/agent/jmx/JmxMetricTracker.java
index 30283ee3c8..fb9e1405a6 100644
--- a/apm-agent-plugins/apm-jmx-plugin/src/main/java/co/elastic/apm/agent/jmx/JmxMetricTracker.java
+++ b/apm-agent-plugins/apm-jmx-plugin/src/main/java/co/elastic/apm/agent/jmx/JmxMetricTracker.java
@@ -212,12 +212,7 @@ public void onChange(ConfigurationOption<?> configurationOption, List<JmxMetric>
             retryExecutor.scheduleAtFixedRate(new Runnable() {
                 @Override
                 public void run() {
-                    List<JmxMetric> failed = JmxMetricTracker.this.failedMetrics;
-                    synchronized (failed) {
-                        List<JmxMetric> toRetry = new ArrayList<>(failed);
-                        failed.clear();
-                        register(toRetry, platformMBeanServer, failed);
-                    }
+                    retryFailedJmx(platformMBeanServer);
                 }
             }, retryMillis, retryMillis, TimeUnit.MILLISECONDS);
         }
@@ -225,6 +220,16 @@ public void run() {
         register(jmxConfiguration.getCaptureJmxMetrics().get(), platformMBeanServer, failedMetrics);
     }
 
+    // package-private for testing
+    void retryFailedJmx(MBeanServer platformMBeanServer) {
+        List<JmxMetric> failed = JmxMetricTracker.this.failedMetrics;
+        synchronized (failed) {
+            List<JmxMetric> toRetry = new ArrayList<>(failed);
+            failed.clear();
+            register(toRetry, platformMBeanServer, failed);
+        }
+    }
+
     private void registerMBeanNotificationListener(final MBeanServer server) {
         MBeanServerNotificationFilter filter = new MBeanServerNotificationFilter();
         filter.enableAllObjectNames();
diff --git a/apm-agent-plugins/apm-jmx-plugin/src/test/java/co/elastic/apm/agent/jmx/JmxMetricTrackerTest.java b/apm-agent-plugins/apm-jmx-plugin/src/test/java/co/elastic/apm/agent/jmx/JmxMetricTrackerTest.java
index c993f91b4d..cc8eeca851 100644
--- a/apm-agent-plugins/apm-jmx-plugin/src/test/java/co/elastic/apm/agent/jmx/JmxMetricTrackerTest.java
+++ b/apm-agent-plugins/apm-jmx-plugin/src/test/java/co/elastic/apm/agent/jmx/JmxMetricTrackerTest.java
@@ -181,15 +181,53 @@ void testMBeanUnregister() throws Exception {
         assertThat(metricRegistry.getGauge("jvm.jmx.Baz", labels)).isNull();
     }
 
+    @Test
+    void testMBeanExceptionWhenRegisteredThenOk() throws Exception {
+        ObjectName objectName = new ObjectName("foo:type=Foo,name=testMBeanExceptionWhenRegisteredThenOk");
+        TestMetric testMetric = new TestMetric();
+        testMetric.setValue(-1);
+        assertThatThrownBy(testMetric::getBaz).isInstanceOf(RuntimeException.class);
+
+        setConfig(JmxMetric.valueOf("object_name[foo:type=Foo,name=*] attribute[Baz]"));
+
+        registerMBean(testMetric, objectName);
+
+        Labels labels = Labels.Mutable.of("name", "testMBeanExceptionWhenRegisteredThenOk").add("type", "Foo");
+        assertThat(metricRegistry.getGaugeValue("jvm.jmx.Baz", labels))
+            .isNaN();
+        assertThat(metricRegistry.getGauge("jvm.jmx.Baz", labels))
+            .isNull();
+
+        testMetric.setValue(37);
+        // calling directly the JMX tracker to avoid testing async execution
+        jmxTracker.retryFailedJmx(mbeanServer);
+
+        assertThat(metricRegistry.getGaugeValue("jvm.jmx.Baz", labels))
+            .isEqualTo(37);
+    }
+
     public interface TestMetricMBean {
         int getBaz();
     }
 
     public static class TestMetric implements TestMetricMBean {
 
+        private int value;
+
+        public TestMetric() {
+            this.value = 42;
+        }
+
+        void setValue(int value) {
+            this.value = value;
+        }
+
         @Override
         public int getBaz() {
-            return 42;
+            if (value < 0) {
+                throw new RuntimeException("value less than zero");
+            }
+            return value;
         }
     }