Bean subscription

src/main/java/org/datadog/jmxfetch/BeanListener.java

@@ -0,0 +1,8 @@
+package org.datadog.jmxfetch;
+
+import javax.management.ObjectName;
+
+public interface BeanListener {
+    public void beanRegistered(ObjectName mBeanName);
+    public void beanUnregistered(ObjectName mBeanName);
+}

src/main/java/org/datadog/jmxfetch/BeanSubscriber.java

@@ -0,0 +1,29 @@
+package org.datadog.jmxfetch;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.List;
+
+public class BeanSubscriber implements Runnable {
+    private List<String> beanScopes;
+    private Connection connection;
+    private BeanListener listener;
+    public CompletableFuture<Boolean> subscriptionSuccessful;
+
+    BeanSubscriber(List<String> beanScopes, Connection connection, BeanListener listener) {
+        this.beanScopes = beanScopes;
+        this.connection = connection;
+        this.listener = listener;
+        this.subscriptionSuccessful = new CompletableFuture<Boolean>();
+    }
+
+    public void run() {
+        try {
+            connection.subscribeToBeanScopes(beanScopes, this.listener);
+            this.subscriptionSuccessful.complete(true);
+
+            Thread.currentThread().join();
+        } catch (Exception e) {
+            this.subscriptionSuccessful.complete(false);
+        }
+    }
+}

src/main/java/org/datadog/jmxfetch/Connection.java

@@ -7,7 +7,9 @@
 import java.net.SocketTimeoutException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.List;
 import java.util.Set;
+import java.util.function.Consumer;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ExecutorService;
@@ -21,8 +23,16 @@
 import javax.management.MBeanAttributeInfo;
 import javax.management.MBeanException;
 import javax.management.MBeanServerConnection;
+import javax.management.MBeanServerDelegate;
+import javax.management.relation.MBeanServerNotificationFilter;
+import javax.management.MBeanServerNotification;
+import javax.management.MalformedObjectNameException;
+import javax.management.Notification;
+import javax.management.NotificationFilter;
+import javax.management.NotificationListener;
 import javax.management.ObjectName;
 import javax.management.ReflectionException;
+import javax.management.remote.JMXConnectionNotification;
 import javax.management.remote.JMXConnector;
 import javax.management.remote.JMXConnectorFactory;
 import javax.management.remote.JMXServiceURL;
@@ -36,6 +46,38 @@ public class Connection {
     protected Map<String, Object> env;
     protected JMXServiceURL address;
 
+    private static class BeanNotificationListener implements NotificationListener {
+        private BeanListener bl;
+
+        public BeanNotificationListener(BeanListener bl) {
+            this.bl = bl;
+        }
+        public void handleNotification(Notification notification, Object handback) {
+            if (!(notification instanceof MBeanServerNotification)) {
+                log.warn("Got unknown notification, expected MBeanServerNotification but got {}", notification.getClass());
+                return;
+            }
+            MBeanServerNotification mbs = (MBeanServerNotification) notification;
+            log.debug("MBeanNotification: ts {} seqNum: {} msg: '{}' userData: {} ", mbs.getTimeStamp(), mbs.getSequenceNumber(), mbs.getMessage(), mbs.getUserData());
+            ObjectName mBeanName = mbs.getMBeanName();
+            if (mbs.getType().equals(MBeanServerNotification.REGISTRATION_NOTIFICATION)) {
+                bl.beanRegistered(mBeanName);
+            } else if (mbs.getType().equals(MBeanServerNotification.UNREGISTRATION_NOTIFICATION)) {
+                this.bl.beanUnregistered(mBeanName);
+            }
+        }
+    }
+
+    public void subscribeToBeanScopes(List<String> beanScopes, BeanListener bl) throws MalformedObjectNameException, IOException, InstanceNotFoundException{
+        BeanNotificationListener listener = new BeanNotificationListener(bl);
+        for (String scope : beanScopes) {
+            ObjectName name = new ObjectName(scope);
+            MBeanServerNotificationFilter filter = new MBeanServerNotificationFilter();
+            filter.enableObjectName(name);
+        }
+        mbs.addNotificationListener(MBeanServerDelegate.DELEGATE_NAME, listener, null, null);
+    }
+
     /** Gets attributes for matching bean name. */
     public MBeanAttributeInfo[] getAttributesForBean(ObjectName beanName)
             throws InstanceNotFoundException, IntrospectionException, ReflectionException,