Unregister custom MBeans

pom.xml

@@ -3,7 +3,7 @@
   <modelVersion>4.0.0</modelVersion>
   <artifactId>classloader-leak-prevention</artifactId>
   <groupId>se.jiderhamn</groupId>
-  <version>1.0.3-SNAPSHOT</version>
+  <version>1.1.0</version>
   <packaging>jar</packaging>
   <name>ClassLoader leak prevention</name>
   <url>https://github.com/mjiderhamn/classloader-leak-prevention</url>

src/main/java/se/jiderhamn/classloader/leak/prevention/ClassLoaderLeakPreventor.java

@@ -15,6 +15,7 @@
  */
 package se.jiderhamn.classloader.leak.prevention;
 
+import java.lang.management.ManagementFactory;
 import java.lang.ref.Reference;
 import java.lang.ref.WeakReference;
 import java.lang.reflect.Field;
@@ -27,6 +28,10 @@
 import java.sql.SQLException;
 import java.util.*;
 import java.util.concurrent.ThreadPoolExecutor;
+import javax.management.InstanceNotFoundException;
+import javax.management.MBeanRegistrationException;
+import javax.management.MBeanServer;
+import javax.management.ObjectName;
 import javax.servlet.*;
 import javax.xml.parsers.ParserConfigurationException;
 
@@ -338,6 +343,9 @@ public void contextDestroyed(ServletContextEvent servletContextEvent) {
     // Deregister JDBC drivers contained in web application
     deregisterJdbcDrivers();
 
+    // Unregister MBeans loaded by the web application class loader
+    unregisterMBeans();
+
     // Deregister shutdown hooks - execute them immediately
     deregisterShutdownHooks();
 
@@ -428,6 +436,29 @@ public void deregisterJdbcDrivers() {
     }
   }
 
+  /** Unregister MBeans loaded by the web application class loader */
+  protected void unregisterMBeans() {
+    try {
+      MBeanServer mBeanServer = ManagementFactory.getPlatformMBeanServer();
+      final Set<ObjectName> allMBeanNames = mBeanServer.queryNames(new ObjectName("*:*"), null);
+      for(ObjectName objectName : allMBeanNames) {
+        try {
+          final ClassLoader mBeanClassLoader = mBeanServer.getClassLoaderFor(objectName);
+          if(isWebAppClassLoaderOrChild(mBeanClassLoader)) { // MBean loaded in web application
+            warn("MBean '" + objectName + "' was loaded in web application; unregistering");
+            mBeanServer.unregisterMBean(objectName);
+          }
+        }
+        catch(Exception e) { // MBeanRegistrationException / InstanceNotFoundException
+          error(e);
+        }
+      }
+    }
+    catch (Exception e) { // MalformedObjectNameException
+      error(e);
+    }
+  }
+
   /** Find and deregister shutdown hooks. Will by default execute the hooks after removing them. */
   protected void deregisterShutdownHooks() {
     // We will not remove known shutdown hooks, since loading the owning class of the hook,

src/test/java/se/jiderhamn/JUnitClassloaderRunner.java

@@ -113,7 +113,7 @@ public void evaluate() throws Throwable {
 
         if(expectedLeak) { // We expect this test to leak classloaders
           RedefiningClassLoader redefiningClassLoader = weak.get();
-          assertNotNull("ClassLoader has been garbage collected, while test is expected to leak " + weak.get(), redefiningClassLoader);
+          assertNotNull("ClassLoader has been garbage collected, while test is expected to leak", redefiningClassLoader);
 
           if(redefiningClassLoader != null && // Always true, otherwise assertion failure above
              preventorClass != null) {

src/test/java/se/jiderhamn/classloader/leak/prevention/MBeanTest.java

@@ -0,0 +1,43 @@
+package se.jiderhamn.classloader.leak.prevention;
+
+import java.lang.management.ManagementFactory;
+
+import javax.management.*;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import se.jiderhamn.JUnitClassloaderRunner;
+import se.jiderhamn.LeakPreventor;
+
+/**
+ * Verify that custom MBean causes leak that can be prevented
+ * @author Mattias Jiderhamn
+ */
+@RunWith(JUnitClassloaderRunner.class)
+@LeakPreventor(MBeanTest.Prevent.class)
+public class MBeanTest {
+
+  @Test
+  public void triggerMBeanLeak() throws Exception {
+    MBeanServer mBeanServer = ManagementFactory.getPlatformMBeanServer();
+    mBeanServer.registerMBean(new Custom(), new ObjectName("se.jiderhamn:foo=bar"));
+  }
+
+  public interface CustomMBean {
+  }
+
+  public static class Custom implements CustomMBean {
+  }
+
+  public static class Prevent implements Runnable {
+    public void run() {
+      new ClassLoaderLeakPreventor() {
+        { // Initializer / "Constructor"
+          unregisterMBeans();
+        }
+      };
+    }
+
+  }
+
+}