Add support for direct local jmx connections (#201)

This allows access to the mBeanServer without going through a jmx connection.  This is important to dd-java-agent because it also avoids initializing `java.util.logging` too early in the boot process.

With tests to verify class loading behavior, which will help avoid regressions in the future.

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

@@ -34,7 +34,7 @@ public class Connection {
     private final static Logger LOGGER = Logger.getLogger(Connection.class.getName());
     private static final ThreadFactory daemonThreadFactory = new DaemonThreadFactory();
     private JMXConnector connector;
-    private MBeanServerConnection mbs;
+    protected MBeanServerConnection mbs;
     protected HashMap<String, Object> env;
     protected JMXServiceURL address;
 

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

@@ -25,6 +25,14 @@ public static Connection createConnection(LinkedHashMap<String, Object> connecti
             return new AttachApiConnection(connectionParams);
 
         }
+
+        // This is used by dd-java-agent to enable directly connecting to the mbean server.
+        // This works because jmxfetch is being run as a library inside the process.
+        if("service:jmx:local:///".equals(connectionParams.get("jmx_url"))) {
+            LOGGER.info("Connecting using JMX Local");
+            return new LocalConnection();
+        }
+
         LOGGER.info("Connecting using JMX Remote");
         return new RemoteConnection(connectionParams);
 

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

@@ -0,0 +1,25 @@
+package org.datadog.jmxfetch;
+
+import java.io.IOException;
+import java.lang.management.ManagementFactory;
+import org.apache.log4j.Logger;
+
+public class LocalConnection extends Connection {
+    private final static Logger LOGGER = Logger.getLogger(LocalConnection.class.getName());
+
+    public LocalConnection() throws IOException {
+        createConnection();
+    }
+
+    protected void createConnection() throws IOException {
+        mbs = ManagementFactory.getPlatformMBeanServer();
+    }
+
+    public void closeConnector() {
+        // ignore
+    }
+
+    public boolean isAlive() {
+        return true;
+    }
+}

src/test/java/org/datadog/jmxfetch/TestLogInitialization.java

@@ -0,0 +1,195 @@
+package org.datadog.jmxfetch;
+
+import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Sets;
+import com.sun.jmx.remote.util.ClassLogger;
+import java.io.File;
+import java.lang.management.ManagementFactory;
+import java.lang.reflect.Field;
+import java.lang.reflect.Method;
+import java.net.MalformedURLException;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.util.Collections;
+import java.util.Set;
+import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.logging.LogManager;
+import javax.management.ObjectName;
+import javax.management.remote.JMXServiceURL;
+import org.junit.After;
+import org.junit.Ignore;
+import org.junit.Test;
+
+import static com.google.common.base.StandardSystemProperty.JAVA_CLASS_PATH;
+import static com.google.common.base.StandardSystemProperty.PATH_SEPARATOR;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+public class TestLogInitialization {
+    public static final Callable<AppConfig> LOCAL_CONFIG = new Callable<AppConfig>() {
+        public AppConfig call() throws Exception {
+            return AppConfig.create(
+                    ImmutableList.of("org/datadog/jmxfetch/dd-java-agent-jmx.yaml"),
+                    Collections.<String>emptyList(),
+                    (int) TimeUnit.SECONDS.toMillis(30),
+                    (int) TimeUnit.SECONDS.toMillis(30),
+                    Collections.<String, String>emptyMap(),
+                    "console",
+                    "System.err",
+                    "DEBUG");
+        }
+    };
+
+    public static final Callable<AppConfig> REMOTE_CONFIG = new Callable<AppConfig>() {
+        public AppConfig call() throws Exception {
+            return AppConfig.create(
+                    ImmutableList.of("org/datadog/jmxfetch/remote-jmx.yaml"),
+                    Collections.<String>emptyList(),
+                    (int) TimeUnit.SECONDS.toMillis(30),
+                    (int) TimeUnit.SECONDS.toMillis(30),
+                    Collections.<String, String>emptyMap(),
+                    "console",
+                    "System.err",
+                    "DEBUG");
+        }
+    };
+
+    @After
+    public void unregisterLockingMBean() {
+        try {
+        ManagementFactory.getPlatformMBeanServer()
+                .unregisterMBean(new ObjectName("org.datadog.jmxfetch.log_init_test:type=TriggeringMBean"));
+
+        } catch (Exception e) {
+            // Ignore
+        }
+    }
+
+    @Test
+    public void testLocalUsageDoesNotInitalizeLogManager() throws Exception {
+        CountDownLatch latch = registerLockingMBean();
+
+        TrackingClassLoader classLoader = new TrackingClassLoader();
+
+        final AtomicReference<Exception> errored = runInThread(classLoader, "LOCAL_CONFIG");
+
+        latch.await(15, TimeUnit.SECONDS);
+
+        assertNull(errored.get());
+        assertTrue(classLoader.classLoaded(LocalConnection.class.getName()));
+        assertFalse(classLoader.classLoaded(JMXServiceURL.class.getName()));
+    }
+
+    @Test
+    public void testRemoteUsageDoesInitalizeLogManager() throws Exception {
+        registerLockingMBean();
+
+        TrackingClassLoader classLoader = new TrackingClassLoader();
+
+        runInThread(classLoader, "REMOTE_CONFIG");
+
+        // We don't know how long until the error triggers.  No good way to verify.
+        Thread.sleep(TimeUnit.SECONDS.toMillis(1));
+
+        assertTrue(classLoader.classLoaded(RemoteConnection.class.getName()));
+        assertTrue(classLoader.classLoaded(JMXServiceURL.class.getName()));
+    }
+
+    private AtomicReference<Exception> runInThread(final TrackingClassLoader classLoader,
+                                                   String configName) throws Exception {
+        final AtomicReference<Exception> errored = new AtomicReference<Exception>();
+
+        Class<?> appClass = classLoader.loadClass(App.class.getName());
+        Class<?> appConfigClass = classLoader.loadClass(AppConfig.class.getName());
+        assertTrue(classLoader.classLoaded(App.class.getName()));
+        assertTrue(classLoader.classLoaded(AppConfig.class.getName()));
+        final Method runMethod = appClass.getMethod("run", appConfigClass);
+
+        Class<?> thisClass = classLoader.loadClass(getClass().getName());
+        Field appConfigField = thisClass.getField(configName);
+        final Callable config = (Callable)appConfigField.get(null);
+
+        Thread task = new Thread(new Runnable() {
+            public void run() {
+                // This will run forever, so we need to run in a different thread.
+                try {
+                    // App.run(appConfig);
+                    runMethod.invoke(null, config.call());
+                } catch (Exception e) {
+                    errored.set(e);
+                    e.printStackTrace();
+                }
+            }
+        });
+        task.start();
+
+        return errored;
+    }
+
+    private CountDownLatch registerLockingMBean() throws Exception {
+        CountDownLatch latch = new CountDownLatch(1);
+        TriggeringMBean bean = new Triggering(latch);
+        ManagementFactory.getPlatformMBeanServer()
+                .registerMBean(bean, new ObjectName("org.datadog.jmxfetch.log_init_test:type=TriggeringMBean"));
+        return latch;
+    }
+
+    public interface TriggeringMBean {
+        boolean isTriggered();
+    }
+
+    class Triggering implements TriggeringMBean {
+        private final CountDownLatch latch;
+
+        Triggering(CountDownLatch latch) {
+            this.latch = latch;
+        }
+
+        public boolean isTriggered() {
+            System.out.println("Triggering!");
+            latch.countDown();
+            return true;
+        }
+    }
+
+    static class TrackingClassLoader extends URLClassLoader {
+        private final Set<String> loadedClasses = Sets.newConcurrentHashSet();
+
+        TrackingClassLoader() throws MalformedURLException {
+            // Don't delegate to the parent as that already has the classes loaded.
+            super(getClasspathUrls(), null);
+        }
+
+        public Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
+            loadedClasses.add(name);
+            if(name.startsWith("org.apache.log4j")) {
+                return getSystemClassLoader().loadClass(name);
+            }
+            return super.loadClass(name, resolve);
+        }
+
+        boolean classLoaded(String name) {
+            return loadedClasses.contains(name);
+        }
+
+        private static URL[] getClasspathUrls() throws MalformedURLException {
+
+            ImmutableList.Builder<URL> urls = ImmutableList.builder();
+            for (String entry : Splitter.on(PATH_SEPARATOR.value()).split(JAVA_CLASS_PATH.value())) {
+                try {
+                    urls.add(new File(entry).toURI().toURL());
+                } catch (SecurityException e) { // File.toURI checks to see if the file is a directory
+                    urls.add(new URL("file", null, new File(entry).getAbsolutePath()));
+                }
+
+            }
+            return urls.build().toArray(new URL[0]);
+        }
+    }
+}

src/test/resources/org/datadog/jmxfetch/dd-java-agent-jmx.yaml

@@ -0,0 +1,12 @@
+init_config:
+  is_jmx: true
+  new_gc_metrics: true
+
+instances:
+- jmx_url: service:jmx:local:///
+  conf:
+  - include:
+      domain: org.datadog.jmxfetch.log_init_test
+      attribute:
+        Triggered:
+          metric_type: gauge

src/test/resources/org/datadog/jmxfetch/remote-jmx.yaml

@@ -0,0 +1,13 @@
+init_config:
+  is_jmx: true
+  new_gc_metrics: true
+
+instances:
+#  this is expected to fail with a socket connection error
+- jmx_url: service:jmx:rmi:///jndi/rmi://localhost:9999/jmxrmi
+  conf:
+  - include:
+      domain: org.datadog.jmxfetch.log_init_test
+      attribute:
+        Triggered:
+          metric_type: gauge