2 files changed, 17 insertions, 30 deletions

diff --git a/test/916-obsolete-jit/src/Main.java b/test/916-obsolete-jit/src/Main.java
index 74eb003d5c..1e43f7ee9d 100644
--- a/test/916-obsolete-jit/src/Main.java
+++ b/test/916-obsolete-jit/src/Main.java
@@ -157,38 +157,13 @@ public class Main {
         doCommonClassRedefinition(Transform.class, CLASS_BYTES, DEX_BYTES);
       }
     };
-    // This does nothing.
-    Runnable noop = () -> {};
     // This just prints something out to show we are running the Runnable.
     Runnable say_nothing = () -> { w.accept("Not doing anything here"); };
-    // This checks to see if we have jitted the methods we are testing.
-    Runnable check_interpreting = () -> {
-      // TODO remove the second check when we remove the doCall function. We need to check that
-      // both of these functions aren't being interpreted because if sayHi is the test doesn't do
-      // anything and if doCall is then there will be a runtime call right above the sayHi
-      // function preventing sayHi from being deoptimized.
-      interpreting = has_jit && (Main.isInterpretedFunction(say_hi_method, true) ||
-                                 Main.isInterpretedFunction(do_call_method, false));
-    };
     do {
-      w.clear();
-      // Wait for the methods to be jitted
-      long j = 0;
-      do {
-        for (int i = 0; i < 10000; i++) {
-          t.sayHi(noop, w);
-          j++;
-          // Clear so that we won't OOM if we go around a few times.
-          w.clear();
-        }
-        t.sayHi(check_interpreting, w);
-        if (j >= 1000000) {
-          System.out.println("FAIL: Could not make sayHi be Jitted!");
-          return;
-        }
-        j++;
-      } while(interpreting);
-      // Clear output. Now we try for real.
+      // Run ensureJitCompiled here since it might get GCd
+      ensureJitCompiled(Transform.class, "sayHi");
+      ensureJitCompiled(Main.class, "doCall");
+      // Clear output.
       w.clear();
       // Try and redefine.
       t.sayHi(say_nothing, w);
@@ -203,6 +178,8 @@ public class Main {
 
   private static native boolean isInterpretedFunction(Method m, boolean require_deoptimizable);
 
+  private static native void ensureJitCompiled(Class c, String name);
+
   // Transforms the class
   private static native void doCommonClassRedefinition(Class<?> target,
                                                        byte[] classfile,
diff --git a/test/common/runtime_state.cc b/test/common/runtime_state.cc
index 7451cf97de..271657fec3 100644
--- a/test/common/runtime_state.cc
+++ b/test/common/runtime_state.cc
@@ -152,7 +152,12 @@ extern "C" JNIEXPORT void JNICALL Java_Main_ensureJitCompiled(JNIEnv* env,
                                                              jclass,
                                                              jclass cls,
                                                              jstring method_name) {
-  jit::Jit* jit = Runtime::Current()->GetJit();
+  Runtime* runtime = Runtime::Current();
+  if (runtime == nullptr) {
+    // We must be on the RI. We will just roll with it.
+    return;
+  }
+  jit::Jit* jit = runtime->GetJit();
   if (jit == nullptr) {
     return;
   }
@@ -166,6 +171,11 @@ extern "C" JNIEXPORT void JNICALL Java_Main_ensureJitCompiled(JNIEnv* env,
     CHECK(chars.c_str() != nullptr);
     method = soa.Decode<mirror::Class>(cls)->FindDeclaredDirectMethodByName(
         chars.c_str(), kRuntimePointerSize);
+    if (method == nullptr) {
+      method = soa.Decode<mirror::Class>(cls)->FindDeclaredVirtualMethodByName(
+          chars.c_str(), kRuntimePointerSize);
+    }
+    DCHECK(method != nullptr) << "Unable to find method called " << chars.c_str();
   }
 
   jit::JitCodeCache* code_cache = jit->GetCodeCache();