Bean subscription

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

@@ -520,6 +520,9 @@ void start() {
     }
 
     void stop() {
+        for (Instance in : this.getInstances()) {
+            in.cleanUp();
+        }
         this.teardownTelemetry();
         this.collectionProcessor.stop();
         this.recoveryProcessor.stop();

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

@@ -230,6 +230,15 @@ public class AppConfig {
             required = false)
     private String statusLocation;
 
+    @Parameter(
+            names = {"--enable_bean_subscription", "-B"},
+            description =
+                    "EXPERIMENTAL: If true, JMX beans will be discovered via subscription rather"
+                    + " than poll-based. Obsoletes 'initialBeanRefreshPeriod' and"
+                    + " 'beanRefreshPeriod'.",
+            required = false)
+    private boolean enableBeanSubscription;
+
     @Parameter(
             names = {"--exit_file_location", "-e"},
             description =
@@ -518,4 +527,12 @@ public int getSocketTimeout() {
     public String getVersion() {
         return MetadataHelper.getVersion();
     }
+
+    public boolean getEnableBeanSubscription() {
+        // As noted in `pkg/jmxfetch/jmxfetch.go` in the agent, using an env var
+        // for enablement is a temporary measure until the stable JMXFetch is upgraded
+        // to a version supporting this CLI arg.
+        boolean isEnvEnabled = System.getenv("DD_JMX_BEAN_SUBSCRIPTION_ENABLED") != null;
+        return isEnvEnabled || enableBeanSubscription;
+    }
 }

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

@@ -0,0 +1,58 @@
+package org.datadog.jmxfetch;
+
+import lombok.extern.slf4j.Slf4j;
+
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.LinkedBlockingQueue;
+
+import javax.management.MBeanServerNotification;
+import javax.management.Notification;
+import javax.management.NotificationListener;
+import javax.management.ObjectName;
+
+
+@Slf4j
+class BeanNotificationListener implements NotificationListener {
+    private final BlockingQueue<MBeanServerNotification> queue;
+    private final BeanTracker beanTracker;
+
+    public BeanNotificationListener(final BeanTracker bt) {
+        this.beanTracker = bt;
+        this.queue = new LinkedBlockingQueue<>();
+        new Thread(new Runnable() {
+            @Override
+            public void run() {
+                while (true) {
+                    try {
+                        MBeanServerNotification mbs = queue.take();
+                        processMBeanServerNotification(mbs);
+                    } catch (InterruptedException e) {
+                        // ignore
+                    }
+                }
+            }
+        }).start();
+    }
+
+    @Override
+    public void handleNotification(Notification notification, Object handback) {
+        if (!(notification instanceof MBeanServerNotification)) {
+            return;
+        }
+        MBeanServerNotification mbs = (MBeanServerNotification) notification;
+        queue.offer(mbs);
+    }
+
+    private void processMBeanServerNotification(MBeanServerNotification notif) {
+        log.debug("MBeanNotification: ts {} seqNum: {} msg: '{}'",
+            notif.getTimeStamp(),
+            notif.getSequenceNumber(),
+            notif.getMessage());
+        ObjectName beanName = notif.getMBeanName();
+        if (notif.getType().equals(MBeanServerNotification.REGISTRATION_NOTIFICATION)) {
+            beanTracker.trackBean(beanName);
+        } else if (notif.getType().equals(MBeanServerNotification.UNREGISTRATION_NOTIFICATION)) {
+            beanTracker.untrackBean(beanName);
+        }
+    }
+}

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

@@ -0,0 +1,9 @@
+package org.datadog.jmxfetch;
+
+import javax.management.ObjectName;
+
+public interface BeanTracker {
+    public void trackBean(ObjectName beanName);
+
+    public void untrackBean(ObjectName beanName);
+}

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

@@ -3,27 +3,25 @@
 import lombok.extern.slf4j.Slf4j;
 
 import java.io.IOException;
-import java.io.InterruptedIOException;
-import java.net.SocketTimeoutException;
-import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.ArrayBlockingQueue;
-import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ThreadFactory;
-import java.util.concurrent.TimeUnit;
 import javax.management.Attribute;
 import javax.management.AttributeNotFoundException;
 import javax.management.InstanceNotFoundException;
 import javax.management.IntrospectionException;
-import javax.management.MBeanAttributeInfo;
+import javax.management.ListenerNotFoundException;
 import javax.management.MBeanException;
 import javax.management.MBeanInfo;
 import javax.management.MBeanServerConnection;
+import javax.management.MBeanServerDelegate;
+import javax.management.MalformedObjectNameException;
+import javax.management.Notification;
+import javax.management.NotificationListener;
 import javax.management.ObjectName;
 import javax.management.ReflectionException;
+import javax.management.relation.MBeanServerNotificationFilter;
+import javax.management.remote.JMXConnectionNotification;
 import javax.management.remote.JMXConnector;
 import javax.management.remote.JMXConnectorFactory;
 import javax.management.remote.JMXServiceURL;
@@ -35,6 +33,44 @@ public class Connection {
     protected MBeanServerConnection mbs;
     protected Map<String, Object> env;
     protected JMXServiceURL address;
+    private NotificationListener connectionNotificationListener;
+    private boolean seenConnectionIssues;
+
+    private static class ConnectionNotificationListener implements NotificationListener {
+        public void handleNotification(Notification notification, Object handback) {
+            if (!(notification instanceof JMXConnectionNotification)) {
+                return;
+            }
+            if (!(handback instanceof Connection)) {
+                return;
+            }
+
+            JMXConnectionNotification connNotif = (JMXConnectionNotification) notification;
+            Connection conn = (Connection) handback;
+
+            if (connNotif.getType() == JMXConnectionNotification.CLOSED
+                    || connNotif.getType() == JMXConnectionNotification.FAILED
+                    || connNotif.getType() == JMXConnectionNotification.NOTIFS_LOST) {
+                log.warn("Marking connection issues due to {} - {}",
+                     connNotif.getType(), connNotif.getMessage());
+                conn.seenConnectionIssues = true;
+            }
+            log.debug("Received connection notification: {} Message: {}",
+                 connNotif.getType(), connNotif.getMessage());
+        }
+    }
+
+    /** Subscribes for bean registration/deregistration events under the specified bean scopes. */
+    public void subscribeToBeanScopes(List<String> beanScopes, BeanTracker 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 MBeanInfo getMBeanInfo(ObjectName beanName)
@@ -56,6 +92,10 @@ protected void createConnection() throws IOException {
         log.info("Connecting to: " + this.address);
         connector = JMXConnectorFactory.connect(this.address, this.env);
         mbs = connector.getMBeanServerConnection();
+
+        this.connectionNotificationListener = new ConnectionNotificationListener();
+        connector.addConnectionNotificationListener(
+            this.connectionNotificationListener, null, this);
     }
 
     /** Gets attribute for matching bean and attribute name. */
@@ -73,13 +113,21 @@ public Object getAttribute(ObjectName objectName, String attributeName)
     public void closeConnector() {
         if (connector != null) {
             try {
+                this.connector.removeConnectionNotificationListener(
+                    this.connectionNotificationListener);
                 connector.close();
-            } catch (IOException e) {
+                connector = null;
+            } catch (IOException | ListenerNotFoundException e) {
                 // ignore
             }
         }
     }
 
+    /** True if connection has been notified of failure/lost notifications. */
+    public boolean hasSeenConnectionIssues() {
+        return this.seenConnectionIssues;
+    }
+
     /** Returns a boolean describing if the connection is still alive. */
     public boolean isAlive() {
         if (connector == null) {