Register transformation hook only when needed

This prevents us from attempting to transform classes when not needed.
Every transformation is expensive, hence we want to avoid them if we
can.

If another agent triggers transformations of mocked classes the objects
of the classes will become unmocked without knowledge of mockito breaking
many assuptions.

Having another agent registered is very rare, i.e. currently the only
other usage is profiling using Android Studio. Still, we want to support
this, hence we might want to revert this change later.

Wifi-tests runtime 590s -> 417s

Test: frameworks/opt/net/wifi/tests/wifitests/runtests.sh
      atest CtsInlineMockingTestCases CtsMockingTestCase
Bug: 72052158

Change-Id: I2556e749806bed4f80697d5cba38a8d8a2f7ce6a

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]);