From 46c05232c4f8c1b01a0cd61e088a40cff74e2875 Mon Sep 17 00:00:00 2001
From: Cristian Ciutea <cristi.ciutea@gmail.com>
Date: Mon, 14 Nov 2022 11:54:32 +0100
Subject: [PATCH] improve err handling (#138)

---
 .../org/newrelic/nrjmx/v2/JMXFetcher.java     | 27 ++++++++++++++++---
 1 file changed, 23 insertions(+), 4 deletions(-)

diff --git a/src/main/java/org/newrelic/nrjmx/v2/JMXFetcher.java b/src/main/java/org/newrelic/nrjmx/v2/JMXFetcher.java
index 9020f87..d25aaad 100644
--- a/src/main/java/org/newrelic/nrjmx/v2/JMXFetcher.java
+++ b/src/main/java/org/newrelic/nrjmx/v2/JMXFetcher.java
@@ -231,8 +231,10 @@ private Set<ObjectInstance> queryMBeans(ObjectName objectName) throws JMXConnect
         }
 
         try {
+            MBeanServerConnection conn = getConnection();
+
             result = withConnectionExceptionHandler(() ->
-                    getConnection().queryMBeans(objectName, null)
+                    conn.queryMBeans(objectName, null)
             );
 
             if (internalStat != null) {
@@ -300,8 +302,10 @@ private List<String> getMBeanAttributeNames(ObjectName objectName) throws JMXCon
         }
 
         try {
+            MBeanServerConnection conn = getConnection();
+
             info = withConnectionExceptionHandler(() ->
-                    getConnection().getMBeanInfo(objectName)
+                    conn.getMBeanInfo(objectName)
             );
 
             if (internalStat != null) {
@@ -397,11 +401,12 @@ private void getMBeanAttributes(ObjectName objectName, List<String> attributes,
         List<Attribute> attrValues = new ArrayList<>();
         AttributeList attributeList;
         try {
+            MBeanServerConnection conn = getConnection();
 
             List<String> finalAttributes = attributes;
 
             attributeList = withConnectionExceptionHandler(() ->
-                    getConnection().getAttributes(objectName, finalAttributes.toArray(new String[0]))
+                    conn.getAttributes(objectName, finalAttributes.toArray(new String[0]))
             );
 
             if (internalStat != null) {
@@ -509,13 +514,17 @@ private void getMBeanAttribute(ObjectName objectName, String attribute, List<Att
 
         Object value;
         try {
+            MBeanServerConnection conn = getConnection();
+
             value = withConnectionExceptionHandler(() ->
-                    getConnection().getAttribute(objectName, attribute)
+                    conn.getAttribute(objectName, attribute)
             );
             if (value instanceof Attribute) {
                 Attribute jmxAttr = (Attribute) value;
                 value = jmxAttr.getValue();
             }
+        } catch (JMXConnectionError je) {
+            throw je;
         } catch (ConnectException ce) {
             String message = String.format("can't connect to JMX server, error: '%s'", ce.getMessage());
             throw new JMXConnectionError(message);
@@ -652,6 +661,16 @@ private <T> T withTimeout(Future<T> future, long timeoutMs) throws JMXError, JMX
     private <T> T withConnectionExceptionHandler(Callable<T> task) throws Exception {
         try {
             return task.call();
+        } catch (
+            // Not connection errors.
+                MBeanException
+                | AttributeNotFoundException
+                | InstanceNotFoundException
+                | ReflectionException
+                | RuntimeOperationsException
+                | IllegalArgumentException
+                | IntrospectionException e) {
+            throw e;
         } catch (Exception e) {
             if (e instanceof ConnectException) {
                 disconnect();