New feature for global (constructor) mocks: mock instance registratio…

…n callback

I implemented this in Javassist first, then migrated the solution to ASM + BB.
Actually, the Javassist version only works with '-noverify' because there of
Javassist bug jboss-javassist/javassist#328.

Status quo:
  - ConstructorMockRegistry.isMockUnderConstruction() now returns an int instead
    of a boolean. Negative value means: no mock under construction. Positive
    values reflect the stack depth of constructors, i.e. 1 means that the
    top-level constructor (TLC) is currently being executed. This is utilised in
    order to only call ConstructorMockRegistry.registerMockInstance(this) after
    the super() returns to the TLC.
  - ConstructorMockRegistry.mockInstances is a private static Set<Object> in
    which all new object instances are registered, but it is not being used
    anywhere yet and does not even have an accessor method. TODO:
      * 'unregister' method
      * automatic unregistration if mock class is deactivated
      * grouping instances by class or class name -> turn the Set into a Map

sarek-constructor-mock-javassist/pom.xml

@@ -14,9 +14,18 @@
   <build>
     <plugins>
 
-      <!-- Note that here we use Surefire, not Failsafe, i.e. no agent JARs -->
       <plugin>
-        <artifactId>maven-surefire-plugin</artifactId>
+        <artifactId>maven-failsafe-plugin</artifactId>
+        <configuration>
+          <!-- TODO: remove '-noverify' after https://github.com/jboss-javassist/javassist/issues/328 is fixed -->
+          <argLine>-noverify</argLine>
+        </configuration>
+        <executions>
+          <execution>
+            <id>reuse-jvm</id>
+            <phase>none</phase>
+          </execution>
+        </executions>
       </plugin>
 
     </plugins>
@@ -44,6 +53,11 @@
       <artifactId>sarek-test-common</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>dev.sarek</groupId>
+      <artifactId>sarek-junit4-runner</artifactId>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>net.bytebuddy</groupId>
       <artifactId>byte-buddy-agent</artifactId>

sarek-constructor-mock-javassist/src/main/java/dev/sarek/agent/constructor_mock/ConstructorMockJavassistTransformer.java

@@ -157,10 +157,22 @@ private void makeConstructorsMockable(CtClass targetClass)
     for (CtConstructor ctConstructor : targetClass.getDeclaredConstructors()) {
       if (LOG_CONSTRUCTOR_MOCK)
         log("Adding constructor mock capability to constructor " + ctConstructor.getLongName());
+
+      // Actually this is not what we want because the visibility scope of a local variable declared this way is the
+      // whole constructor A variable visible only within our inserted code's scope is actually better.
+      //ctConstructor.addLocalVariable("constructorStackDepth", CtPrimitiveType.intType);
+
       String ifCondition = String.join("\n",
-        "if (dev.sarek.agent.constructor_mock.ConstructorMockRegistry.isMockUnderConstruction()) {",
-        "  " + superCall,
-        "  return;",
+        "{",
+        "  int constructorStackDepth = " + MOCK_REGISTRY + "#isMockUnderConstruction();",
+//        "  constructorStackDepth = " + MOCK_REGISTRY + "#isMockUnderConstruction();",
+        "  if (constructorStackDepth > 0) {",
+        "    " + superCall,
+        "    if (constructorStackDepth == 1) {",
+        "      " + MOCK_REGISTRY + "#registerMockInstance($0);",
+        "    }",
+        "    return;",
+        "  }",
         "}"
       );
       if (LOG_CONSTRUCTOR_MOCK)

sarek-constructor-mock-javassist/src/test/java/dev/sarek/agent/constructor_mock/JavassistConstructorMockIT.java → sarek-constructor-mock-javassist/src/test/java/dev/sarek/agent/constructor_mock/ConstructorMockJavassistIT.java

@@ -1,20 +1,20 @@
 package dev.sarek.agent.constructor_mock;
 
+import dev.sarek.junit4.SarekRunner;
 import dev.sarek.test.util.SeparateJVM;
 import net.bytebuddy.agent.ByteBuddyAgent;
 import org.acme.*;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
 
-import java.io.IOException;
 import java.lang.instrument.Instrumentation;
 import java.lang.instrument.UnmodifiableClassException;
 import java.time.Instant;
 import java.util.Date;
 import java.util.UUID;
-import java.util.jar.JarFile;
 
 import static org.junit.Assert.*;
 
@@ -31,7 +31,9 @@
  * retransformations, so they have to be done during class-loading.
  */
 @Category(SeparateJVM.class)
-public class JavassistConstructorMockIT {
+@RunWith(SarekRunner.class)
+//@Ignore("Only working with '-noverify'  because of https://github.com/jboss-javassist/javassist/issues/328")
+public class ConstructorMockJavassistIT {
   private static final Instrumentation INSTRUMENTATION = ByteBuddyAgent.install();
   private static final Class<?>[] TRANSFORMATION_TARGETS = {
     Sub.class, Base.class,
@@ -41,12 +43,12 @@ public class JavassistConstructorMockIT {
   private static ConstructorMockJavassistTransformer constructorMockTransformer;
 
   @BeforeClass
-  public static void beforeClass() throws IOException, UnmodifiableClassException {
-    // This property is usually set in Maven in order to tell us the path to the constructor mock agent.
-    // Important: The JAR needs to contain Javassist too, so it should be the '-all' or '-all-special' artifact.
-    JarFile sarekAllJar = new JarFile(System.getProperty("sarek.jar"));
-    // Inject constructor mock agent JAR into bootstrap classloader
-    INSTRUMENTATION.appendToBootstrapClassLoaderSearch(sarekAllJar);
+  public static void beforeClass() throws UnmodifiableClassException {
+
+    // Please note: There is no need to add Javassist and this module's JAR to the bootstrap class path, because in
+    // contrast to the aspect and mock APIs, the API itself or its dependencies are never used by transformed code. The
+    // only thing that must be on the bootstrap class loader is the ConstructorMockRegistry from module
+    // 'sarek-constructor-mock', which is already part of the Sarek uber JAR and will be injected via SarekRunner.
 
 //    ConstructorMockJavassistTransformer.LOG_CONSTRUCTOR_MOCK = true;
 //    ConstructorMockJavassistTransformer.DUMP_CLASS_FILES = true;

sarek-constructor-mock/src/main/java/dev/sarek/agent/constructor_mock/ConstructorMockMethodVisitor.java

@@ -7,6 +7,7 @@
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.jar.asm.Label;
 import net.bytebuddy.jar.asm.MethodVisitor;
+import net.bytebuddy.utility.visitor.LocalVariableAwareMethodVisitor;
 
 import java.util.ArrayList;
 import java.util.Comparator;
@@ -15,17 +16,25 @@
 
 import static net.bytebuddy.jar.asm.Opcodes.*;
 
-public class ConstructorMockMethodVisitor extends MethodVisitor {
+public class ConstructorMockMethodVisitor extends LocalVariableAwareMethodVisitor {
   private static final int ASM_API_VERSION = ASM8;
+  private static final String CONSTRUCTOR_MOCK_REGISTRY = "dev/sarek/agent/constructor_mock/ConstructorMockRegistry";
 
   private final String className;
   private final boolean shouldTransform;
   private final String superClassNameJVM;
   private final String superConstructorSignatureJVM;
   private final List<String> superConstructorParameterTypes;
 
-  public ConstructorMockMethodVisitor(TypeDescription instrumentedType, MethodVisitor mv) {
-    super(ASM_API_VERSION, mv);
+  public ConstructorMockMethodVisitor(
+    TypeDescription instrumentedType,
+    MethodVisitor methodVisitor,
+    MethodDescription methodDescription
+  )
+  {
+    // TODO: ASM_API_VERSION is ASM8, ByteBuddy also uses ASM8 -> check if it can be synchronised
+    // super(ASM_API_VERSION, methodVisitor);
+    super(methodVisitor, methodDescription);
 
     className = instrumentedType.getTypeName();
     shouldTransform = shouldTransform();
@@ -62,17 +71,27 @@ public void visitCode() {
     if (!shouldTransform)
       return;
 
-    // Constructor mock code starts here
+    // First free offset in local variable table
+    final int freeOffset = getFreeOffset();
+
+    // Define labels we want to insert as markers or jump targets later
     Label labelMockCode = new Label();
+    Label labelOriginalCode = new Label();
+    Label labelReturn = new Label();
+
+    // Constructor mock code starts here
     super.visitLabel(labelMockCode);
 
     // If class for instance under construction is registered for constructor mocking, call super constructor with
     // dummy values (null, 0, false), otherwise jump to original code
-    super.visitMethodInsn(INVOKESTATIC, "dev/sarek/agent/constructor_mock/ConstructorMockRegistry", "isMockUnderConstruction", "()Z", false);
-    Label labelOriginalCode = new Label();
-    super.visitJumpInsn(IFEQ, labelOriginalCode);
-
-    // Push uninitialised 'this' onto the stack
+    super.visitMethodInsn(INVOKESTATIC, CONSTRUCTOR_MOCK_REGISTRY, "isMockUnderConstruction", "()I", false);
+    final int mockUnderConstructionResult = freeOffset;
+    super.visitVarInsn(ISTORE, mockUnderConstructionResult);
+    super.visitVarInsn(ILOAD, mockUnderConstructionResult);
+    super.visitInsn(ICONST_0);
+    super.visitJumpInsn(IF_ICMPLE, labelOriginalCode);
+
+    // Push uninitialised 'this' onto the stack as first super constructor parameter
     super.visitVarInsn(ALOAD, 0);
 
     // Push dummy values for all super constructor parameters onto the stack
@@ -113,7 +132,18 @@ public void visitCode() {
 
     // Invoke super constructor
     super.visitMethodInsn(INVOKESPECIAL, superClassNameJVM, "<init>", superConstructorSignatureJVM, false);
-    // Skip original constructor code
+
+    // If we are in the top-most constructor (mockUnderConstructionResult == 1), register mock via
+    // ConstructorMockRegistry.registerMockInstance(this)
+    super.visitVarInsn(ILOAD, mockUnderConstructionResult);
+    super.visitInsn(ICONST_1);
+    super.visitJumpInsn(IF_ICMPNE, labelReturn);
+    // Now after super constructor was called, 'this' is initialised and can be used normally
+    super.visitVarInsn(ALOAD, 0);
+    super.visitMethodInsn(INVOKESTATIC, CONSTRUCTOR_MOCK_REGISTRY, "registerMockInstance", "(Ljava/lang/Object;)V", false);
+
+    // Skip original constructor code by RETURN
+    super.visitLabel(labelReturn);
     super.visitInsn(RETURN);
 
     // Original constructor code starts here

sarek-constructor-mock/src/main/java/dev/sarek/agent/constructor_mock/ConstructorMockRegistry.java

@@ -1,6 +1,8 @@
 package dev.sarek.agent.constructor_mock;
 
+import java.util.Collections;
 import java.util.HashSet;
+import java.util.IdentityHashMap;
 import java.util.Set;
 
 /**
@@ -14,6 +16,9 @@ public class ConstructorMockRegistry {
   //       Alternatively, if one day we drop Java 8 support we can use the StackWalker API.
   private static Set<String> mockClasses = new HashSet<>();
 
+  // This is effectively an IdentityHashSet, see https://stackoverflow.com/a/48096408/1082681
+  private static Set<Object> mockInstances = Collections.newSetFromMap(new IdentityHashMap<>());
+
   /**
    * Determine whether a given class has been registered for constructor mocking
    *
@@ -56,7 +61,7 @@ public static boolean deactivate(String className) {
    *
    * @return boolean value specifying if the object under construction ought to be a mock or not
    */
-  public static boolean isMockUnderConstruction() {
+  public static int isMockUnderConstruction() {
     // TODO:
     //   - Under Java 8, according to Rafael Winterhalter it would be quicker and more efficient (as in not
     //     materialising the whole stack) to use sun.misc.JavaLangAccess, e.g.
@@ -73,6 +78,17 @@ public static boolean isMockUnderConstruction() {
         break;
       constructorIndex = i;
     }
-    return constructorIndex > 0 && isMock(stackTrace[constructorIndex].getClassName());
+    return constructorIndex > 0 && isMock(stackTrace[constructorIndex].getClassName())
+      ? constructorIndex
+      : -1;
+  }
+
+  public static void registerMockInstance(Object mockInstance) {
+    // Caveat: Do not log anything here, especially not mock objects with possibly stubbed toString methods. Otherwise
+    // you might see strange exceptions in then failing tests.
+    mockInstances.add(mockInstance);
+    // TODO:
+    //   - 'unregister' method
+    //   - automatic unregistration if mock class is deactivated
   }
 }

sarek-constructor-mock/src/main/java/dev/sarek/agent/constructor_mock/ConstructorMockTransformer.java

@@ -123,7 +123,7 @@ public MethodVisitor wrap(
     {
       if (logVerbose)
         System.out.println("[Constructor Mock Transformer] Mocking constructor " + instrumentedMethod);
-      return new ConstructorMockMethodVisitor(instrumentedType, methodVisitor);
+      return new ConstructorMockMethodVisitor(instrumentedType, methodVisitor, instrumentedMethod);
     }
   }