Fall back to Class.forName in case boot class is already defined (#41)

* Remove unnecessary DUP
* Align non-boot and boot define-class approaches
* Fall back to Class.forName in case boot class is already defined

Author: mcculls

class-inject/src/glue/java/datadog/instrument/glue/DefineClassGlueGenerator.java

@@ -81,6 +81,9 @@ final class DefineClassGlueGenerator {
           + CLASS_CLASS
           + ";";
 
+  private static final String CLASS_FORNAME_DESCRIPTOR =
+      "(L" + STRING_CLASS + ";ZL" + CLASSLOADER_CLASS + ";)L" + CLASS_CLASS + ";";
+
   private DefineClassGlueGenerator() {}
 
   /**
@@ -166,6 +169,10 @@ private static byte[] generateBytecode(String unsafeNamespace) {
     final Label unlockClassLoaderAndThrow = new Label();
     final Label setupUnsafeDefiner = new Label();
     final Label hasNextBootClass = new Label();
+    final Label defineBootClass = new Label();
+    final Label storeBootClass = new Label();
+    final Label loadBootClass = new Label();
+    final Label rethrowOriginal = new Label();
     final Label returnDefinedClasses = new Label();
 
     // common bytecode variables
@@ -175,20 +182,25 @@ private static byte[] generateBytecode(String unsafeNamespace) {
     final int mapEntrySetIterator = 4;
     final int protectionDomain = 5;
     final int classLoader = 6;
+    final int className = 7;
 
     // bytecode variables for non-boot classes
-    final int classLoadingLock = 7;
-    final int bytecodeMapEntry = 8;
-    final int className = 9;
+    final int classLoadingLock = 8;
+    final int bytecodeMapEntry = 9;
 
     // bytecode variables for boot classes
-    final int unsafeInstance = 7;
+    final int unsafeInstance = 8;
+    final int originalError = 9;
 
     mv.visitCode();
 
     // make sure we unlock the class-loader if an exception occurs while defining a class
     mv.visitTryCatchBlock(classLoaderLocked, classLoaderUnlocked, unlockClassLoaderAndThrow, null);
 
+    // fall back and see if the boot class is already loaded if defining it fails
+    mv.visitTryCatchBlock(defineBootClass, storeBootClass, loadBootClass, null);
+    mv.visitTryCatchBlock(loadBootClass, rethrowOriginal, rethrowOriginal, null);
+
     // -------- SHARED SETUP CODE  --------
 
     // store the defined classes in a list
@@ -278,11 +290,13 @@ private static byte[] generateBytecode(String unsafeNamespace) {
         false);
     mv.visitInsn(DUP);
     mv.visitJumpInsn(IFNONNULL, storeNonBootClass);
-
-    // define the class using the given class-name and bytecode
     mv.visitInsn(POP);
+
+    // not yet defined, prepare arguments to define it
     mv.visitVarInsn(ALOAD, classLoader);
     mv.visitVarInsn(ALOAD, className);
+
+    // extract the bytecode of the next class to define
     mv.visitVarInsn(ALOAD, bytecodeMapEntry);
     mv.visitMethodInsn(
         INVOKEINTERFACE, MAP_ENTRY_CLASS, "getValue", "()L" + OBJECT_CLASS + ";", true);
@@ -291,7 +305,10 @@ private static byte[] generateBytecode(String unsafeNamespace) {
     mv.visitInsn(ARRAYLENGTH);
     mv.visitInsn(ICONST_0);
     mv.visitInsn(SWAP);
+
     mv.visitVarInsn(ALOAD, protectionDomain);
+
+    // define the class using the given class-name and bytecode
     mv.visitMethodInsn(
         INVOKEVIRTUAL, CLASSLOADER_CLASS, "defineClass", CLASSLOADER_DEFINECLASS_DESCRIPTOR, false);
 
@@ -312,7 +329,6 @@ private static byte[] generateBytecode(String unsafeNamespace) {
 
     // unlock the class-loader if something goes wrong
     mv.visitLabel(unlockClassLoaderAndThrow);
-    mv.visitInsn(DUP);
     mv.visitVarInsn(ALOAD, classLoadingLock);
     mv.visitInsn(MONITOREXIT);
     mv.visitInsn(ATHROW);
@@ -330,28 +346,38 @@ private static byte[] generateBytecode(String unsafeNamespace) {
     mv.visitMethodInsn(INVOKEINTERFACE, ITERATOR_CLASS, "hasNext", "()Z", true);
     mv.visitJumpInsn(IFEQ, returnDefinedClasses);
 
-    // define the class using the given class-name and bytecode
     mv.visitVarInsn(ALOAD, unsafeInstance);
+
+    // extract the name of the next class to define
     mv.visitVarInsn(ALOAD, mapEntrySetIterator);
     mv.visitMethodInsn(INVOKEINTERFACE, ITERATOR_CLASS, "next", "()L" + OBJECT_CLASS + ";", true);
     mv.visitInsn(DUP);
     mv.visitMethodInsn(
         INVOKEINTERFACE, MAP_ENTRY_CLASS, "getKey", "()L" + OBJECT_CLASS + ";", true);
     mv.visitTypeInsn(CHECKCAST, STRING_CLASS);
+    mv.visitInsn(DUP);
+    mv.visitVarInsn(ASTORE, className);
     mv.visitInsn(SWAP);
+
+    // extract the bytecode of the next class to define
     mv.visitMethodInsn(
         INVOKEINTERFACE, MAP_ENTRY_CLASS, "getValue", "()L" + OBJECT_CLASS + ";", true);
     mv.visitTypeInsn(CHECKCAST, BYTE_ARRAY);
     mv.visitInsn(DUP);
     mv.visitInsn(ARRAYLENGTH);
     mv.visitInsn(ICONST_0);
     mv.visitInsn(SWAP);
+
     mv.visitInsn(ACONST_NULL);
     mv.visitVarInsn(ALOAD, protectionDomain);
+
+    // define the boot class using the given class-name and bytecode
+    mv.visitLabel(defineBootClass);
     mv.visitMethodInsn(
         INVOKEVIRTUAL, unsafeClass, "defineClass", UNSAFE_DEFINECLASS_DESCRIPTOR, false);
 
-    // store the class in the list
+    // store the class in the list whether we defined it or it already existed
+    mv.visitLabel(storeBootClass);
     mv.visitVarInsn(ALOAD, definedClasses);
     mv.visitInsn(SWAP);
     mv.visitMethodInsn(INVOKEVIRTUAL, ARRAYLIST_CLASS, "add", "(L" + OBJECT_CLASS + ";)Z", false);
@@ -360,6 +386,21 @@ private static byte[] generateBytecode(String unsafeNamespace) {
     // check again if we've defined all the given bytecode
     mv.visitJumpInsn(GOTO, hasNextBootClass);
 
+    // see if the boot class has already been defined
+    mv.visitLabel(loadBootClass);
+    mv.visitVarInsn(ASTORE, originalError);
+    mv.visitVarInsn(ALOAD, className);
+    mv.visitInsn(ICONST_0); // initialize=false
+    mv.visitInsn(ACONST_NULL);
+    mv.visitMethodInsn(INVOKESTATIC, CLASS_CLASS, "forName", CLASS_FORNAME_DESCRIPTOR, false);
+    mv.visitJumpInsn(GOTO, storeBootClass);
+
+    // boot class is not defined, report original error
+    mv.visitLabel(rethrowOriginal);
+    mv.visitInsn(POP);
+    mv.visitVarInsn(ALOAD, originalError);
+    mv.visitInsn(ATHROW);
+
     // -------- SHARED RETURN CODE  --------
 
     // return the defined classes