Fix JIT for class unloading

Keep declaring class of method live to prevent unloading.
Wait for JIT to finish compiling before calling Runtime.gc(), this
prevents flaky failures due to classes not being unloaded.

Bug: 22720414

Change-Id: I9fe5e5e39d681bcd22acc2d2f34b0dbc9887708d
diff --git a/test/004-JniTest/jni_test.cc b/test/004-JniTest/jni_test.cc
index 370bd6a..be7888b 100644
--- a/test/004-JniTest/jni_test.cc
+++ b/test/004-JniTest/jni_test.cc
@@ -28,7 +28,7 @@
 
 static JavaVM* jvm = nullptr;
 
-extern "C" JNIEXPORT jint JNI_OnLoad(JavaVM *vm, void *) {
+extern "C" JNIEXPORT jint JNI_OnLoad(JavaVM *vm, void*) {
   assert(vm != nullptr);
   assert(jvm == nullptr);
   jvm = vm;
@@ -36,6 +36,13 @@
   return JNI_VERSION_1_6;
 }
 
+extern "C" JNIEXPORT void JNI_OnUnload(JavaVM*, void*) {
+  // std::cout since LOG(INFO) adds extra stuff like pid.
+  std::cout << "JNI_OnUnload called" << std::endl;
+  // Clear jvm for assert in test 004-JniTest.
+  jvm = nullptr;
+}
+
 static void* AttachHelper(void* arg) {
   assert(jvm != nullptr);
 
diff --git a/test/141-class-unload/expected.txt b/test/141-class-unload/expected.txt
index 124398f..ff65a70 100644
--- a/test/141-class-unload/expected.txt
+++ b/test/141-class-unload/expected.txt
@@ -1,9 +1,15 @@
 1
 2
+JNI_OnLoad called
+JNI_OnUnload called
 1
 2
+JNI_OnLoad called
+JNI_OnUnload called
 null
 null
+JNI_OnLoad called
+JNI_OnUnload called
 null
 loader null false
 loader null false
diff --git a/test/141-class-unload/jni_unload.cc b/test/141-class-unload/jni_unload.cc
index 894f28c..d913efe 100644
--- a/test/141-class-unload/jni_unload.cc
+++ b/test/141-class-unload/jni_unload.cc
@@ -18,7 +18,20 @@
 
 #include <iostream>
 
-extern "C" JNIEXPORT void JNI_OnUnload(JavaVM*, void *) {
-  // std::cout since LOG(INFO) adds extra stuff like pid.
-  std::cout << "JNI_OnUnload called" << std::endl;
+#include "jit/jit.h"
+#include "jit/jit_instrumentation.h"
+#include "runtime.h"
+#include "thread-inl.h"
+
+namespace art {
+namespace {
+
+extern "C" JNIEXPORT void JNICALL Java_IntHolder_waitForCompilation(JNIEnv*, jclass) {
+  jit::Jit* jit = Runtime::Current()->GetJit();
+  if (jit != nullptr) {
+    jit->GetInstrumentationCache()->WaitForCompilationToFinish(Thread::Current());
+  }
 }
+
+}  // namespace
+}  // namespace art
diff --git a/test/141-class-unload/src-ex/IntHolder.java b/test/141-class-unload/src-ex/IntHolder.java
index b4651af..e4aa6b8 100644
--- a/test/141-class-unload/src-ex/IntHolder.java
+++ b/test/141-class-unload/src-ex/IntHolder.java
@@ -34,4 +34,6 @@
     public static void loadLibrary(String name) {
         System.loadLibrary(name);
     }
+
+    public static native void waitForCompilation();
 }
diff --git a/test/141-class-unload/src/Main.java b/test/141-class-unload/src/Main.java
index da15746..105a2b9 100644
--- a/test/141-class-unload/src/Main.java
+++ b/test/141-class-unload/src/Main.java
@@ -31,7 +31,8 @@
         Constructor constructor =
             pathClassLoader.getDeclaredConstructor(String.class, ClassLoader.class);
         try {
-            testUnloadClassAndLoader(constructor);
+            testUnloadClass(constructor);
+            testUnloadLoader(constructor);
             // Test that we don't unload if we have a Method keeping the class live.
             testNoUnloadInvoke(constructor);
             // Test that we don't unload if we have an instance.
@@ -47,15 +48,14 @@
 
     private static void stressTest(Constructor constructor) throws Exception {
         for (int i = 0; i <= 100; ++i) {
-            setUpUnloadLoader(constructor);
+            setUpUnloadLoader(constructor, false);
             if (i % 10 == 0) {
                 Runtime.getRuntime().gc();
             }
         }
     }
 
-    private static void testUnloadClassAndLoader(Constructor constructor) throws Exception {
-        WeakReference<ClassLoader> loader = setUpUnloadLoader(constructor);
+    private static void testUnloadClass(Constructor constructor) throws Exception {
         WeakReference<Class> klass = setUpUnloadClass(constructor);
         // No strong refernces to class loader, should get unloaded.
         Runtime.getRuntime().gc();
@@ -64,7 +64,15 @@
         // If the weak reference is cleared, then it was unloaded.
         System.out.println(klass.get());
         System.out.println(klass2.get());
-        System.out.println(loader.get());
+    }
+
+    private static void testUnloadLoader(Constructor constructor)
+        throws Exception {
+      WeakReference<ClassLoader> loader = setUpUnloadLoader(constructor, true);
+      // No strong refernces to class loader, should get unloaded.
+      Runtime.getRuntime().gc();
+      // If the weak reference is cleared, then it was unloaded.
+      System.out.println(loader.get());
     }
 
     private static void testLoadAndUnloadLibrary(Constructor constructor) throws Exception {
@@ -96,8 +104,7 @@
         System.out.println("loader null " + isNull);
     }
 
-    private static WeakReference<Class> setUpUnloadClass(Constructor constructor)
-        throws Exception {
+    private static WeakReference<Class> setUpUnloadClass(Constructor constructor) throws Exception {
         ClassLoader loader = (ClassLoader) constructor.newInstance(
             DEX_FILE, ClassLoader.getSystemClassLoader());
         Class intHolder = loader.loadClass("IntHolder");
@@ -108,26 +115,40 @@
         System.out.println((int) getValue.invoke(intHolder));
         setValue.invoke(intHolder, 2);
         System.out.println((int) getValue.invoke(intHolder));
+        waitForCompilation(intHolder);
         return new WeakReference(intHolder);
     }
 
-    private static WeakReference<ClassLoader> setUpUnloadLoader(Constructor constructor)
+    private static WeakReference<ClassLoader> setUpUnloadLoader(Constructor constructor,
+                                                                boolean waitForCompilation)
         throws Exception {
         ClassLoader loader = (ClassLoader) constructor.newInstance(
             DEX_FILE, ClassLoader.getSystemClassLoader());
         Class intHolder = loader.loadClass("IntHolder");
         Method setValue = intHolder.getDeclaredMethod("setValue", Integer.TYPE);
         setValue.invoke(intHolder, 2);
+        if (waitForCompilation) {
+            waitForCompilation(intHolder);
+        }
         return new WeakReference(loader);
     }
 
+    private static void waitForCompilation(Class intHolder) throws Exception {
+      // Load the native library so that we can call waitForCompilation.
+      Method loadLibrary = intHolder.getDeclaredMethod("loadLibrary", String.class);
+      loadLibrary.invoke(intHolder, nativeLibraryName);
+      // Wait for JIT compilation to finish since the async threads may prevent unloading.
+      Method waitForCompilation = intHolder.getDeclaredMethod("waitForCompilation");
+      waitForCompilation.invoke(intHolder);
+    }
+
     private static WeakReference<ClassLoader> setUpLoadLibrary(Constructor constructor)
         throws Exception {
         ClassLoader loader = (ClassLoader) constructor.newInstance(
             DEX_FILE, ClassLoader.getSystemClassLoader());
         Class intHolder = loader.loadClass("IntHolder");
-        Method setValue = intHolder.getDeclaredMethod("loadLibrary", String.class);
-        setValue.invoke(intHolder, nativeLibraryName);
+        Method loadLibrary = intHolder.getDeclaredMethod("loadLibrary", String.class);
+        loadLibrary.invoke(intHolder, nativeLibraryName);
         return new WeakReference(loader);
     }
 }
diff --git a/test/457-regs/expected.txt b/test/457-regs/expected.txt
index e69de29..6a5618e 100644
--- a/test/457-regs/expected.txt
+++ b/test/457-regs/expected.txt
@@ -0,0 +1 @@
+JNI_OnLoad called