Merge "Register transformation hook only when needed"

3 files changed, 27 insertions, 8 deletions

diff --git a/README.version b/README.version
index 29832cf..c671035 100644
--- a/README.version
+++ b/README.version
@@ -10,3 +10,4 @@ It includes a stock code generator for class proxies. If you just want to do AOP
 
 Local Modifications:
         Allow to share classloader via dexmaker.share_classloader system property (I8c2490c3ec8e8582dc41c486f8f7a406bd635ebb)
+        Dynamically register transformation hook (I2556e749806bed4f80697d5cba38a8d8a2f7ce6a)
diff --git a/dexmaker-mockito-inline-tests/src/androidTest/java/com/android/dx/mockito/inline/tests/MultipleJvmtiAgentsInterference.java b/dexmaker-mockito-inline-tests/src/androidTest/java/com/android/dx/mockito/inline/tests/MultipleJvmtiAgentsInterference.java
index d66b128..9231202 100644
--- a/dexmaker-mockito-inline-tests/src/androidTest/java/com/android/dx/mockito/inline/tests/MultipleJvmtiAgentsInterference.java
+++ b/dexmaker-mockito-inline-tests/src/androidTest/java/com/android/dx/mockito/inline/tests/MultipleJvmtiAgentsInterference.java
@@ -19,6 +19,7 @@ package com.android.dx.mockito.inline.tests;
 import android.os.Build;
 
 import org.junit.BeforeClass;
+import org.junit.Ignore;
 import org.junit.Test;
 
 import java.io.File;
@@ -33,6 +34,8 @@ import static org.junit.Assert.assertNull;
 import static org.junit.Assume.assumeTrue;
 import static org.mockito.Mockito.mock;
 
+@Ignore("Leaving the transformation hook enabled causes a lot of unnecessary transformations. " +
+        "This is too expensive. Hence for now, we cannot support multiple agents")
 public class MultipleJvmtiAgentsInterference {
     private static final String AGENT_LIB_NAME = "multiplejvmtiagentsinterferenceagent";
 
diff --git a/dexmaker-mockito-inline/src/main/jni/dexmakerjvmtiagent/agent.cc b/dexmaker-mockito-inline/src/main/jni/dexmakerjvmtiagent/agent.cc
index e00ada8..358e375 100644
--- a/dexmaker-mockito-inline/src/main/jni/dexmakerjvmtiagent/agent.cc
+++ b/dexmaker-mockito-inline/src/main/jni/dexmakerjvmtiagent/agent.cc
@@ -779,6 +779,19 @@ Java_com_android_dx_mockito_inline_ClassTransformer_nativeRedefine(JNIEnv* env,
     return transformedArr;
 }
 
+// Register the ClassFileLoadHook hook. This causes Transform to be called for every class load and
+// for every trigger of RetransformClasses
+static jvmtiError registerClassFileLoadHook() {
+    return localJvmtiEnv->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_CLASS_FILE_LOAD_HOOK,
+                                                   NULL);
+}
+
+// Unregister the ClassFileLoadHook hook.
+static jvmtiError unregisterClassFileLoadHook() {
+    return localJvmtiEnv->SetEventNotificationMode(JVMTI_DISABLE, JVMTI_EVENT_CLASS_FILE_LOAD_HOOK,
+                                                   NULL);
+}
+
 // Initializes the agent
 extern "C" jint Agent_OnAttach(JavaVM* vm,
                                char* options,
@@ -806,12 +819,6 @@ extern "C" jint Agent_OnAttach(JavaVM* vm,
         return error;
     }
 
-    error = localJvmtiEnv->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_CLASS_FILE_LOAD_HOOK,
-                                                    NULL);
-    if (error != JVMTI_ERROR_NONE) {
-        return error;
-    }
-
     return JVMTI_ERROR_NONE;
 }
 
@@ -854,8 +861,16 @@ Java_com_android_dx_mockito_inline_JvmtiAgent_nativeRetransformClasses(JNIEnv* e
         transformedClasses[i] = (jclass) env->NewGlobalRef(env->GetObjectArrayElement(classes, i));
     }
 
-    jvmtiError error = localJvmtiEnv->RetransformClasses(numTransformedClasses,
-                                                         transformedClasses);
+    jvmtiError error = registerClassFileLoadHook();
+    if (error == JVMTI_ERROR_NONE) {
+        error = localJvmtiEnv->RetransformClasses(numTransformedClasses,
+                                                  transformedClasses);
+
+        jvmtiError unregisterError = unregisterClassFileLoadHook();
+        if (error == JVMTI_ERROR_NONE && unregisterError != JVMTI_ERROR_NONE) {
+            error = unregisterError;
+        }
+    }
 
     for (int i = 0; i < numTransformedClasses; i++) {
         env->DeleteGlobalRef(transformedClasses[i]);