Fix MyClassLoader idempotency - issue #49

afterburner/src/main/java/com/fasterxml/jackson/module/afterburner/util/MyClassLoader.java

@@ -2,6 +2,9 @@
 
 import java.lang.reflect.Method;
 import java.nio.charset.Charset;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 
 /**
  * Class loader that is needed to load generated classes.
@@ -10,6 +13,11 @@ public class MyClassLoader extends ClassLoader
 {
     private final static Charset UTF8 = Charset.forName("UTF-8");
 
+    // Maps parent classloader instance and class name to the corresponding lock object.
+    // N.B. this must be static because multiple instances of MyClassLoader must all use the same lock
+    // when loading classes directly on the same parent.
+    private final static ConcurrentHashMap<String, Object> parentParallelLockMap = new ConcurrentHashMap<>();
+
     /**
      * Flag that determines if we should first try to load new class
      * using parent class loader or not; this may be done to try to
@@ -56,52 +64,156 @@ public static boolean canAddClassInPackageOf(Class<?> cls)
      *   implement
      */
     public Class<?> loadAndResolve(ClassName className, byte[] byteCode)
-        throws IllegalArgumentException
+            throws IllegalArgumentException
     {
-        // First things first: just to be sure; maybe we have already loaded it?
-        Class<?> old = findLoadedClass(className.getDottedName());
-        if (old != null) {
-            return old;
+        // first, try to loadAndResolve via the parent classloader, if configured to do so
+        Class<?> classFromParent = loadAndResolveUsingParentClassloader(className, byteCode);
+        if (classFromParent != null) {
+            return classFromParent;
         }
-        
-        Class<?> impl;
-
-        // Important: bytecode is generated with a template name (since bytecode itself
-        // is used for checksum calculation) -- must be replaced now, however
-        replaceName(byteCode, className.getSlashedTemplate(), className.getSlashedName());
-
-        // First: let's try calling it directly on parent, to be able to access protected/package-access stuff:
-        if (_cfgUseParentLoader) {
-            ClassLoader cl = getParent();
-            // if we have parent, that is
-            if (cl != null) {
-                try {
-                    Method method = ClassLoader.class.getDeclaredMethod("defineClass",
-                            new Class[] {String.class, byte[].class, int.class,
-                            int.class});
-                    method.setAccessible(true);
-                    return (Class<?>)method.invoke(getParent(),
-                            className.getDottedName(), byteCode, 0, byteCode.length);
-                } catch (Exception e) {
-                    // Should we handle this somehow?
+
+        // fall back to loading and resolving ourselves
+        synchronized (getClassLoadingLock(className.getDottedName())) {
+            // First: check to see if we have loaded it ourselves
+            Class<?> existingClass = findLoadedClass(className.getDottedName());
+            if (existingClass != null) {
+                return existingClass;
+            }
+
+            // Important: bytecode is generated with a template name (since bytecode itself
+            // is used for checksum calculation) -- must be replaced now, however
+            replaceName(byteCode, className.getSlashedTemplate(), className.getSlashedName());
+
+            // Second: define a new class instance using the bytecode
+            Class<?> newClass;
+            try {
+                newClass = defineClass(className.getDottedName(), byteCode, 0, byteCode.length);
+            } catch (LinkageError e) {
+                Throwable t = e;
+                while (t.getCause() != null) {
+                    t = t.getCause();
                 }
+                throw new IllegalArgumentException("Failed to load class '" + className + "': " + t.getMessage(), t);
             }
+            // important: must also resolve the newly-created class.
+            resolveClass(newClass);
+            return newClass;
         }
+    }
 
-        // but if that doesn't fly, try to do it from our own class loader
-
-        try {
-            impl = defineClass(className.getDottedName(), byteCode, 0, byteCode.length);
-        } catch (LinkageError e) {
-            Throwable t = e;
-            while (t.getCause() != null) {
-                t = t.getCause();
+    /**
+     * Attempt to load (and resolve) the class using the parent class loader (if it is configured and present).
+     * This method will return {@code null} if the parent classloader is not configured or cannot be retrieved.
+     *
+     * @param className Interface or abstract class that class to load should extend or implement
+     * @param byteCode  the generated bytecode for the class to load
+     * @return          the loaded class, or {@code null} if the class could not be loaded on the parent classloader.
+     */
+    private Class<?> loadAndResolveUsingParentClassloader(ClassName className, byte[] byteCode)
+    {
+        ClassLoader parentClassLoader;
+        if (!_cfgUseParentLoader || (parentClassLoader = getParent()) == null) {
+            return null;
+        }
+        // N.B. The parent-class-loading locks are shared between all instances of MyClassLoader.
+        // We can be confident that no attempt will be made to re-acquire *any* parent-class-loading lock instance
+        // inside the synchronized region (eliminating the risk of deadlock), even if the parent class loader is also
+        // an instance of MyClassLoader, because:
+        //      a) this method is the only place that attempts to acquire a parent class loading lock,
+        //      b) the only non-private method which calls this method and thus acquires this lock is
+        //          MyClassLoader#loadAndResolve,
+        //      c) nothing in the synchronized region can have the effect of calling #loadAndResolve on this
+        //          or any other instance of MyClassLoader.
+        synchronized (getParentClassLoadingLock(parentClassLoader, className.getDottedName())) {
+            // First: check to see if the parent classloader has loaded it already
+            Class<?> impl = findLoadedClassOnParent(parentClassLoader, className.getDottedName());
+            if (impl != null) {
+                return impl;
             }
-            throw new IllegalArgumentException("Failed to load class '"+className+"': "+t.getMessage(), t);
+
+            // Important: bytecode is generated with a template name (since bytecode itself
+            // is used for checksum calculation) -- must be replaced now, however
+            replaceName(byteCode, className.getSlashedTemplate(), className.getSlashedName());
+
+            // Second: define a new class instance on the parent classloder using the bytecode
+            impl = defineClassOnParent(parentClassLoader, className.getDottedName(), byteCode, 0, byteCode.length);
+            // important: must also resolve the newly-created class.
+            resolveClassOnParent(parentClassLoader, impl);
+            return impl;
+        }
+    }
+
+    /**
+     * Get the class loading lock for the parent class loader for loading the named class.
+     *
+     * This is effectively the same implementation as ClassLoader#getClassLoadingLock, but using
+     * our static parentParallelLockMap and keying off of the parent ClassLoader instance as well as
+     * the class name to load.
+     *
+     * @param parentClassLoader     The parent ClassLoader
+     * @param className             The name of the to-be-loaded class
+     */
+    private Object getParentClassLoadingLock(ClassLoader parentClassLoader, String className) {
+        // N.B. using the canonical name and identity hash code to represent the parent class loader in the key
+        // in case that ClassLoader instance (which could be anything) implements #hashCode or #toString poorly.
+        // In the event of a collision here (same key, different parent class loader), we will end up using the
+        // same lock only to synchronize loads of the same class on two different class loaders,
+        // which shouldn't ever deadlock (see proof in #loadAndResolveUsingParentClassloader);
+        // worst case is unnecessary contention for the lock.
+        String key = parentClassLoader.getClass().getCanonicalName()
+                + ":" + System.identityHashCode(parentClassLoader)
+                + ":" + className;
+        Object newLock = new Object();
+        Object lock = parentParallelLockMap.putIfAbsent(key, newLock);
+        if (lock == null) {
+            lock = newLock;
+        }
+        return lock;
+    }
+
+    private Class<?> findLoadedClassOnParent(ClassLoader parentClassLoader, String className) {
+        try {
+            Method method = ClassLoader.class.getDeclaredMethod("findLoadedClass", String.class);
+            method.setAccessible(true);
+            return (Class<?>) method.invoke(parentClassLoader, className);
+        } catch (Exception e) {
+            String msg = String.format("Exception trying 'findLoadedClass(%s)' on parent ClassLoader '%s'",
+                    className, parentClassLoader);
+            Logger.getLogger(MyClassLoader.class.getName()).log(Level.FINE, msg, e);
+            return null;
+        }
+    }
+
+    // visible for testing
+    Class<?> defineClassOnParent(ClassLoader parentClassLoader,
+                                 String className,
+                                 byte[] byteCode,
+                                 int offset,
+                                 int length) {
+        try {
+            Method method = ClassLoader.class.getDeclaredMethod("defineClass",
+                    new Class[]{String.class, byte[].class, int.class, int.class});
+            method.setAccessible(true);
+            return (Class<?>) method.invoke(parentClassLoader,
+                    className, byteCode, offset, length);
+        } catch (Exception e) {
+            String msg = String.format("Exception trying 'defineClass(%s, <bytecode>)' on parent ClassLoader '%s'",
+                    className, parentClassLoader);
+            Logger.getLogger(MyClassLoader.class.getName()).log(Level.FINE, msg, e);
+            return null;
+        }
+    }
+
+    private void resolveClassOnParent(ClassLoader parentClassLoader, Class<?> clazz) {
+        try {
+            Method method = ClassLoader.class.getDeclaredMethod("resolveClass", Class.class);
+            method.setAccessible(true);
+            method.invoke(parentClassLoader, clazz);
+        } catch (Exception e) {
+            String msg = String.format("Exception trying 'resolveClass(%s)' on parent ClassLoader '%s'",
+                    clazz, parentClassLoader);
+            Logger.getLogger(MyClassLoader.class.getName()).log(Level.FINE, msg, e);
         }
-        // important: must also resolve the class...
-        resolveClass(impl);
-        return impl;
     }
 
     public static int replaceName(byte[] byteCode,