Make instrumentation trampoline able to jump to jit code

In order to get the benefit of the instrumentation trampoline it must
be able to jump to jit code. This patch adds support for doing that
and adds code to ensure that we will not be jumping back and forth
between the trampoline and the interpreter all the time if the jit has
not yet compiled the current method.

Note we also disable the jit-gc when turning on these trampolines
since otherwise we could end up either sending instrumentation events
multiple times or running uninitialized memory.

Bug: 110263880
Test: ./test/testrunner/testrunner.py --host --runtime-option=-Xplugin:libtracefast-trampolined.so
Test: ./test/testrunner/testrunner.py --host --runtime-option=-Xplugin:libtracefast-trampolined.so --runtime-option=-Xjitthreshhold:0
Test: ./test.py --host

Change-Id: Ie6e92ec6367452fe4fde24d520d808b7af91d1b5
diff --git a/runtime/class_linker.cc b/runtime/class_linker.cc
index 4a5da1f..966d636 100644
--- a/runtime/class_linker.cc
+++ b/runtime/class_linker.cc
@@ -2811,6 +2811,7 @@
 }
 
 bool ClassLinker::ShouldUseInterpreterEntrypoint(ArtMethod* method, const void* quick_code) {
+  ScopedAssertNoThreadSuspension sants(__FUNCTION__);
   if (UNLIKELY(method->IsNative() || method->IsProxyMethod())) {
     return false;
   }
@@ -2840,6 +2841,12 @@
     return true;
   }
 
+  if (quick_code == GetQuickInstrumentationEntryPoint()) {
+    const void* instr_target = instr->GetCodeForInvoke(method);
+    DCHECK_NE(instr_target, GetQuickInstrumentationEntryPoint()) << method->PrettyMethod();
+    return ShouldUseInterpreterEntrypoint(method, instr_target);
+  }
+
   if (runtime->IsJavaDebuggable()) {
     // For simplicity, we ignore precompiled code and go to the interpreter
     // assuming we don't already have jitted code.
@@ -2864,6 +2871,7 @@
 }
 
 void ClassLinker::FixupStaticTrampolines(ObjPtr<mirror::Class> klass) {
+  ScopedAssertNoThreadSuspension sants(__FUNCTION__);
   DCHECK(klass->IsInitialized()) << klass->PrettyDescriptor();
   if (klass->NumDirectMethods() == 0) {
     return;  // No direct methods => no static methods.
@@ -2923,6 +2931,7 @@
                      ArtMethod* method,
                      const OatFile::OatClass* oat_class,
                      uint32_t class_def_method_index) REQUIRES_SHARED(Locks::mutator_lock_) {
+  ScopedAssertNoThreadSuspension sants(__FUNCTION__);
   Runtime* const runtime = Runtime::Current();
   if (runtime->IsAotCompiler()) {
     // The following code only applies to a non-compiler runtime.
diff --git a/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc b/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc
index c894406..a4dea37 100644
--- a/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc
+++ b/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc
@@ -38,6 +38,7 @@
 #include "interpreter/interpreter_common.h"
 #include "interpreter/shadow_frame-inl.h"
 #include "jit/jit.h"
+#include "jit/jit_code_cache.h"
 #include "linear_alloc.h"
 #include "method_handles.h"
 #include "mirror/class-inl.h"
@@ -1107,7 +1108,20 @@
   if (instrumentation->IsDeoptimized(method)) {
     result = GetQuickToInterpreterBridge();
   } else {
-    result = instrumentation->GetQuickCodeFor(method, kRuntimePointerSize);
+    // This will get the entry point either from the oat file, the JIT or the appropriate bridge
+    // method if none of those can be found.
+    result = instrumentation->GetCodeForInvoke(method);
+    jit::Jit* jit = Runtime::Current()->GetJit();
+    DCHECK_NE(result, GetQuickInstrumentationEntryPoint()) << method->PrettyMethod();
+    DCHECK(jit == nullptr ||
+           // Native methods come through here in Interpreter entrypoints. We might not have
+           // disabled jit-gc but that is fine since we won't return jit-code for native methods.
+           method->IsNative() ||
+           !jit->GetCodeCache()->GetGarbageCollectCode());
+    DCHECK(!method->IsNative() ||
+           jit == nullptr ||
+           !jit->GetCodeCache()->ContainsPc(result))
+        << method->PrettyMethod() << " code will jump to possibly cleaned up jit code!";
   }
 
   bool interpreter_entry = (result == GetQuickToInterpreterBridge());
diff --git a/runtime/instrumentation.cc b/runtime/instrumentation.cc
index e5cdef7..1ddca87 100644
--- a/runtime/instrumentation.cc
+++ b/runtime/instrumentation.cc
@@ -175,6 +175,9 @@
     if ((forced_interpret_only_ || IsDeoptimized(method)) && !method->IsNative()) {
       new_quick_code = GetQuickToInterpreterBridge();
     } else if (is_class_initialized || !method->IsStatic() || method->IsConstructor()) {
+      // It would be great to search the JIT for its implementation here but we cannot due to the
+      // locks we hold. Instead just set to the interpreter bridge and that code will search the JIT
+      // when it gets called and replace the entrypoint then.
       if (NeedDebugVersionFor(method)) {
         new_quick_code = GetQuickToInterpreterBridge();
       } else {
@@ -192,12 +195,15 @@
       // class, all its static methods code will be set to the instrumentation entry point.
       // For more details, see ClassLinker::FixupStaticTrampolines.
       if (is_class_initialized || !method->IsStatic() || method->IsConstructor()) {
-        if (NeedDebugVersionFor(method)) {
-          // Oat code should not be used. Don't install instrumentation stub and
-          // use interpreter for instrumentation.
-          new_quick_code = GetQuickToInterpreterBridge();
-        } else if (entry_exit_stubs_installed_) {
+        if (entry_exit_stubs_installed_) {
+          // This needs to be checked first since the instrumentation entrypoint will be able to
+          // find the actual JIT compiled code that corresponds to this method.
           new_quick_code = GetQuickInstrumentationEntryPoint();
+        } else if (NeedDebugVersionFor(method)) {
+          // It would be great to search the JIT for its implementation here but we cannot due to
+          // the locks we hold. Instead just set to the interpreter bridge and that code will search
+          // the JIT when it gets called and replace the entrypoint then.
+          new_quick_code = GetQuickToInterpreterBridge();
         } else {
           new_quick_code = class_linker->GetQuickOatCodeFor(method);
         }
@@ -801,6 +807,16 @@
                  // Proxy.<init> correctly in all cases.
                  method != jni::DecodeArtMethod(WellKnownClasses::java_lang_reflect_Proxy_init)) {
         new_quick_code = GetQuickInstrumentationEntryPoint();
+        if (!method->IsNative() && Runtime::Current()->GetJit() != nullptr) {
+          // Native methods use trampoline entrypoints during interpreter tracing.
+          DCHECK(!Runtime::Current()->GetJit()->GetCodeCache()->GetGarbageCollectCode());
+          ProfilingInfo* profiling_info = method->GetProfilingInfo(kRuntimePointerSize);
+          // Tracing will look at the saved entry point in the profiling info to know the actual
+          // entrypoint, so we store it here.
+          if (profiling_info != nullptr) {
+            profiling_info->SetSavedEntryPoint(quick_code);
+          }
+        }
       } else {
         new_quick_code = quick_code;
       }
@@ -995,6 +1011,14 @@
     level = InstrumentationLevel::kInstrumentWithInterpreter;
   } else {
     level = InstrumentationLevel::kInstrumentWithInstrumentationStubs;
+    if (Runtime::Current()->GetJit() != nullptr) {
+      // TODO b/110263880 It would be better if we didn't need to do this.
+      // Since we need to hold the method entrypoint across a suspend to ensure instrumentation
+      // hooks are called correctly we have to disable jit-gc to ensure that the entrypoint doesn't
+      // go away. Furthermore we need to leave this off permanently since one could get the same
+      // effect by causing this to be toggled on and off.
+      Runtime::Current()->GetJit()->GetCodeCache()->SetGarbageCollectCode(false);
+    }
   }
   ConfigureStubs(key, level);
 }
@@ -1003,6 +1027,49 @@
   ConfigureStubs(key, InstrumentationLevel::kInstrumentNothing);
 }
 
+const void* Instrumentation::GetCodeForInvoke(ArtMethod* method) const {
+  // This is called by instrumentation entry only and that should never be getting proxy methods.
+  DCHECK(!method->IsProxyMethod()) << method->PrettyMethod();
+  ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
+  if (LIKELY(!instrumentation_stubs_installed_ && !interpreter_stubs_installed_)) {
+    // In general we just return whatever the method thinks its entrypoint is here. The only
+    // exception is if it still has the instrumentation entrypoint. That means we are racing another
+    // thread getting rid of instrumentation which is unexpected but possible. In that case we want
+    // to wait and try to get it from the oat file or jit.
+    const void* code = method->GetEntryPointFromQuickCompiledCodePtrSize(kRuntimePointerSize);
+    DCHECK(code != nullptr);
+    if (code != GetQuickInstrumentationEntryPoint()) {
+      return code;
+    } else if (method->IsNative()) {
+      return class_linker->GetQuickOatCodeFor(method);
+    }
+    // We don't know what it is. Fallthough to try to find the code from the JIT or Oat file.
+  } else if (method->IsNative()) {
+    // TODO We could have JIT compiled native entrypoints. It might be worth it to find these.
+    return class_linker->GetQuickOatCodeFor(method);
+  } else if (UNLIKELY(interpreter_stubs_installed_)) {
+    return GetQuickToInterpreterBridge();
+  }
+  // Since the method cannot be native due to ifs above we can always fall back to interpreter
+  // bridge.
+  const void* result = GetQuickToInterpreterBridge();
+  if (!NeedDebugVersionFor(method)) {
+    // If we don't need a debug version we should see what the oat file/class linker has to say.
+    result = class_linker->GetQuickOatCodeFor(method);
+  }
+  // If both those fail try the jit.
+  if (result == GetQuickToInterpreterBridge()) {
+    jit::Jit* jit = Runtime::Current()->GetJit();
+    if (jit != nullptr) {
+      const void* res = jit->GetCodeCache()->FindCompiledCodeForInstrumentation(method);
+      if (res != nullptr) {
+        result = res;
+      }
+    }
+  }
+  return result;
+}
+
 const void* Instrumentation::GetQuickCodeFor(ArtMethod* method, PointerSize pointer_size) const {
   ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
   if (LIKELY(!instrumentation_stubs_installed_)) {
diff --git a/runtime/instrumentation.h b/runtime/instrumentation.h
index 8e7a638..bc54ba2 100644
--- a/runtime/instrumentation.h
+++ b/runtime/instrumentation.h
@@ -292,6 +292,10 @@
   void UpdateMethodsCodeForJavaDebuggable(ArtMethod* method, const void* quick_code)
       REQUIRES_SHARED(Locks::mutator_lock_) REQUIRES(!deoptimized_methods_lock_);
 
+  // Return the code that we can execute for an invoke including from the JIT.
+  const void* GetCodeForInvoke(ArtMethod* method) const
+      REQUIRES_SHARED(Locks::mutator_lock_);
+
   // Get the quick code for the given method. More efficient than asking the class linker as it
   // will short-cut to GetCode if instrumentation and static method resolution stubs aren't
   // installed.
diff --git a/runtime/jit/jit.cc b/runtime/jit/jit.cc
index 5e73603..ed449b5 100644
--- a/runtime/jit/jit.cc
+++ b/runtime/jit/jit.cc
@@ -735,8 +735,11 @@
 
   ProfilingInfo* profiling_info = method->GetProfilingInfo(kRuntimePointerSize);
   // Update the entrypoint if the ProfilingInfo has one. The interpreter will call it
-  // instead of interpreting the method.
-  if ((profiling_info != nullptr) && (profiling_info->GetSavedEntryPoint() != nullptr)) {
+  // instead of interpreting the method. We don't update it for instrumentation as the entrypoint
+  // must remain the instrumentation entrypoint.
+  if ((profiling_info != nullptr) &&
+      (profiling_info->GetSavedEntryPoint() != nullptr) &&
+      (method->GetEntryPointFromQuickCompiledCode() != GetQuickInstrumentationEntryPoint())) {
     Runtime::Current()->GetInstrumentation()->UpdateMethodsCode(
         method, profiling_info->GetSavedEntryPoint());
   } else {
diff --git a/runtime/jit/jit_code_cache.cc b/runtime/jit/jit_code_cache.cc
index fe154a9..e635ad9 100644
--- a/runtime/jit/jit_code_cache.cc
+++ b/runtime/jit/jit_code_cache.cc
@@ -174,7 +174,10 @@
   bool use_ashmem = !generate_debug_info && !kIsTargetLinux && !kIsTargetFuchsia;
 
   // With 'perf', we want a 1-1 mapping between an address and a method.
-  bool garbage_collect_code = !generate_debug_info;
+  // We aren't able to keep method pointers live during the instrumentation method entry trampoline
+  // so we will just disable jit-gc if we are doing that.
+  bool garbage_collect_code = !generate_debug_info &&
+      !Runtime::Current()->GetInstrumentation()->AreExitStubsInstalled();
 
   // We need to have 32 bit offsets from method headers in code cache which point to things
   // in the data cache. If the maps are more than 4G apart, having multiple maps wouldn't work.
@@ -316,6 +319,17 @@
   return code_map_->Begin() <= ptr && ptr < code_map_->End();
 }
 
+bool JitCodeCache::WillExecuteJitCode(ArtMethod* method) {
+  ScopedObjectAccess soa(art::Thread::Current());
+  ScopedAssertNoThreadSuspension sants(__FUNCTION__);
+  if (ContainsPc(method->GetEntryPointFromQuickCompiledCode())) {
+    return true;
+  } else if (method->GetEntryPointFromQuickCompiledCode() == GetQuickInstrumentationEntryPoint()) {
+    return FindCompiledCodeForInstrumentation(method) != nullptr;
+  }
+  return false;
+}
+
 bool JitCodeCache::ContainsMethod(ArtMethod* method) {
   MutexLock mu(Thread::Current(), lock_);
   if (UNLIKELY(method->IsNative())) {
@@ -348,6 +362,20 @@
   return nullptr;
 }
 
+const void* JitCodeCache::FindCompiledCodeForInstrumentation(ArtMethod* method) {
+  if (LIKELY(!GetGarbageCollectCode())) {
+    return nullptr;
+  }
+  ProfilingInfo* info = method->GetProfilingInfo(kRuntimePointerSize);
+  if (info == nullptr) {
+    return nullptr;
+  }
+  // When GC is disabled for trampoline tracing we will use SavedEntrypoint to hold the actual
+  // jit-compiled version of the method. If jit-gc is disabled for other reasons this will just be
+  // nullptr.
+  return info->GetSavedEntryPoint();
+}
+
 class ScopedCodeCacheWrite : ScopedTrace {
  public:
   explicit ScopedCodeCacheWrite(const JitCodeCache* const code_cache)
@@ -978,6 +1006,8 @@
     // checks should always pass.
     DCHECK(!info->IsInUseByCompiler());
     new_method->SetProfilingInfo(info);
+    // Get rid of the old saved entrypoint if it is there.
+    info->SetSavedEntryPoint(nullptr);
     info->method_ = new_method;
   }
   // Update method_code_map_ to point to the new method.
@@ -1321,7 +1351,8 @@
       if (GetLiveBitmap()->Test(allocation)) {
         ++it;
       } else {
-        method_headers.insert(OatQuickMethodHeader::FromCodePointer(code_ptr));
+        OatQuickMethodHeader* header = OatQuickMethodHeader::FromCodePointer(code_ptr);
+        method_headers.insert(header);
         it = method_code_map_.erase(it);
       }
     }
@@ -1786,13 +1817,18 @@
                                              const OatQuickMethodHeader* header) {
   DCHECK(!method->IsNative());
   ProfilingInfo* profiling_info = method->GetProfilingInfo(kRuntimePointerSize);
+  const void* method_entrypoint = method->GetEntryPointFromQuickCompiledCode();
   if ((profiling_info != nullptr) &&
       (profiling_info->GetSavedEntryPoint() == header->GetEntryPoint())) {
+    // When instrumentation is set, the actual entrypoint is the one in the profiling info.
+    method_entrypoint = profiling_info->GetSavedEntryPoint();
     // Prevent future uses of the compiled code.
     profiling_info->SetSavedEntryPoint(nullptr);
   }
 
-  if (method->GetEntryPointFromQuickCompiledCode() == header->GetEntryPoint()) {
+  // Clear the method counter if we are running jitted code since we might want to jit this again in
+  // the future.
+  if (method_entrypoint == header->GetEntryPoint()) {
     // The entrypoint is the one to invalidate, so we just update it to the interpreter entry point
     // and clear the counter to get the method Jitted again.
     Runtime::Current()->GetInstrumentation()->UpdateMethodsCode(
diff --git a/runtime/jit/jit_code_cache.h b/runtime/jit/jit_code_cache.h
index b056bc3..f7829de 100644
--- a/runtime/jit/jit_code_cache.h
+++ b/runtime/jit/jit_code_cache.h
@@ -154,6 +154,10 @@
   // Return true if the code cache contains this pc.
   bool ContainsPc(const void* pc) const;
 
+  // Returns true if either the method's entrypoint is JIT compiled code or it is the
+  // instrumentation entrypoint and we can jump to jit code for this method. For testing use only.
+  bool WillExecuteJitCode(ArtMethod* method) REQUIRES(!lock_);
+
   // Return true if the code cache contains this method.
   bool ContainsMethod(ArtMethod* method) REQUIRES(!lock_);
 
@@ -269,6 +273,16 @@
     garbage_collect_code_ = value;
   }
 
+  bool GetGarbageCollectCode() const {
+    return garbage_collect_code_;
+  }
+
+  // If Jit-gc has been disabled (and instrumentation has been enabled) this will return the
+  // jit-compiled entrypoint for this method.  Otherwise it will return null.
+  const void* FindCompiledCodeForInstrumentation(ArtMethod* method)
+      REQUIRES(!lock_)
+      REQUIRES_SHARED(Locks::mutator_lock_);
+
  private:
   // Take ownership of maps.
   JitCodeCache(MemMap* code_map,