CASSANDRASC-116: Allow for JmxClient to be extensible

Patch by Francisco Guerrero; Reviewed by Yifan Cai for CASSANDRASC-116
diff --git a/CHANGES.txt b/CHANGES.txt
index 3e79c80..7a88255 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -1,5 +1,6 @@
 1.0.0
 -----
+ * Allow for JmxClient to be extensible (CASSANDRASC-116)
  * Improve observability in Sidecar (CASSANDRASC-111)
  * Improve logging for slice restore task (CASSANDRASC-107)
  * Add restore task watcher to report long running tasks (CASSANDRASC-106)
diff --git a/common/src/main/java/org/apache/cassandra/sidecar/common/JmxClient.java b/common/src/main/java/org/apache/cassandra/sidecar/common/JmxClient.java
index 44ab18c..d9b0325 100644
--- a/common/src/main/java/org/apache/cassandra/sidecar/common/JmxClient.java
+++ b/common/src/main/java/org/apache/cassandra/sidecar/common/JmxClient.java
@@ -77,7 +77,7 @@
      *
      * @param builder the builder options
      */
-    private JmxClient(Builder builder)
+    protected JmxClient(Builder builder)
     {
         if (builder.jmxServiceURL != null)
         {
@@ -151,7 +151,7 @@
                : RMISocketFactory.getDefaultSocketFactory();
     }
 
-    private synchronized void checkConnection()
+    protected synchronized void checkConnection()
     {
         if (!this.connected)
         {
@@ -159,7 +159,7 @@
         }
     }
 
-    private void connect()
+    protected void connect()
     {
         int attempts = 1;
         int maxAttempts = connectionMaxRetries;
@@ -204,7 +204,7 @@
         throw new RuntimeException(error, lastThrown);
     }
 
-    private void connectInternal(int currentAttempt) throws IOException
+    protected void connectInternal(int currentAttempt) throws IOException
     {
         jmxConnector = JMXConnectorFactory.connect(jmxServiceURL, buildJmxEnv());
         jmxConnector.addConnectionNotificationListener(this, null, null);
diff --git a/common/src/test/java/org/apache/cassandra/sidecar/common/JmxClientTest.java b/common/src/test/java/org/apache/cassandra/sidecar/common/JmxClientTest.java
index a3cfe86..6c53190 100644
--- a/common/src/test/java/org/apache/cassandra/sidecar/common/JmxClientTest.java
+++ b/common/src/test/java/org/apache/cassandra/sidecar/common/JmxClientTest.java
@@ -336,6 +336,51 @@
         }
     }
 
+    @Test
+    void testExtendingJmxClient() throws IOException
+    {
+        try (JmxClientExtended extended = new JmxClientExtended(JmxClient.builder()
+                                                                         .host("127.0.0.1")
+                                                                         .port(port)
+                                                                         .roleSupplier(() -> "controlRole")
+                                                                         .passwordSupplier(() -> "password")))
+        {
+            assertThat(extended.proxyInterceptionCount.get()).isEqualTo(0);
+            List<String> result = extended.proxy(Import.class, objectName)
+                                          .importNewSSTables(Sets.newHashSet("foo", "bar"), true,
+                                                             true, true, true, true,
+                                                             true);
+            assertThat(result.size()).isEqualTo(0);
+            assertThat(extended.proxyInterceptionCount.get()).isEqualTo(1);
+
+            extended.proxy(Import.class, objectName)
+                    .importNewSSTables(Sets.newHashSet("foo", "bar"), true,
+                                       true, true, true, true,
+                                       true);
+            assertThat(extended.proxyInterceptionCount.get()).isEqualTo(2);
+        }
+    }
+
+    /**
+     * Extends from JmxClient and keeps tracks of the number of proxy calls
+     */
+    static class JmxClientExtended extends JmxClient
+    {
+        AtomicInteger proxyInterceptionCount = new AtomicInteger(0);
+
+        protected JmxClientExtended(Builder builder)
+        {
+            super(builder);
+        }
+
+        @Override
+        public <C> C proxy(Class<C> clientClass, String remoteName)
+        {
+            proxyInterceptionCount.incrementAndGet();
+            return super.proxy(clientClass, remoteName);
+        }
+    }
+
     /**
      * Simulates to C*'s `nodetool import` call
      */