Support calling entry / exit hooks from JIT code for non-debuggable

This CL extends the support to call entry / exit hooks directly
from JITed code for non-debuggable mode. This is done as a best effort
and we won't call method exit hooks for the methods that are already
on the stack when method tracing is enabled. This means it might
be less precise than the existing approach.

The idea is to basically invalidate all the JITed code when we
need method entry / exit hooks and re-JIT them with the support
to call entry / exit hooks as and when necessary.

Test: art/testrunner.py
Change-Id: I79e4547368691a21b26226508228d890369823e9
diff --git a/compiler/optimizing/optimizing_compiler.cc b/compiler/optimizing/optimizing_compiler.cc
index 16abf9d..c3ccf95 100644
--- a/compiler/optimizing/optimizing_compiler.cc
+++ b/compiler/optimizing/optimizing_compiler.cc
@@ -800,6 +800,8 @@
     dead_reference_safe = false;
   }
 
+  bool is_instrumentation_enabled =
+      Runtime::Current()->GetInstrumentation()->NeedInstrumentationSupportForJIT();
   HGraph* graph = new (allocator) HGraph(
       allocator,
       arena_stack,
@@ -809,7 +811,7 @@
       compiler_options.GetInstructionSet(),
       kInvalidInvokeType,
       dead_reference_safe,
-      compiler_options.GetDebuggable(),
+      compiler_options.GetDebuggable() || is_instrumentation_enabled,
       compilation_kind);
 
   if (method != nullptr) {
@@ -1296,7 +1298,8 @@
                             /* is_full_debug_info= */ compiler_options.GetGenerateDebugInfo(),
                             compilation_kind,
                             /* has_should_deoptimize_flag= */ false,
-                            cha_single_implementation_list)) {
+                            cha_single_implementation_list,
+                            /* has_instrumentation_support= */ false)) {
       code_cache->Free(self, region, reserved_code.data(), reserved_data.data());
       return false;
     }
@@ -1404,7 +1407,8 @@
                           /* is_full_debug_info= */ compiler_options.GetGenerateDebugInfo(),
                           compilation_kind,
                           codegen->GetGraph()->HasShouldDeoptimizeFlag(),
-                          codegen->GetGraph()->GetCHASingleImplementationList())) {
+                          codegen->GetGraph()->GetCHASingleImplementationList(),
+                          codegen->GetGraph()->IsDebuggable())) {
     code_cache->Free(self, region, reserved_code.data(), reserved_data.data());
     return false;
   }
diff --git a/runtime/instrumentation.cc b/runtime/instrumentation.cc
index 8ef5ef4..c985c6b 100644
--- a/runtime/instrumentation.cc
+++ b/runtime/instrumentation.cc
@@ -234,13 +234,6 @@
     return true;
   }
 
-  // When jiting code for debuggable apps we generate the code to call method
-  // entry / exit hooks when required. Hence it is not required to update
-  // to instrumentation entry point for JITed code in debuggable mode.
-  if (!Runtime::Current()->IsJavaDebuggable()) {
-    return true;
-  }
-
   // Native functions can have JITed entry points but we don't include support
   // for calling entry / exit hooks directly from the JITed code for native
   // functions. So we still have to install entry exit stubs for such cases.
@@ -790,6 +783,28 @@
 }
 
 void Instrumentation::UpdateInstrumentationLevel(InstrumentationLevel requested_level) {
+  if (requested_level == instrumentation_level_) {
+    return;
+  }
+
+  if (!Runtime::Current()->IsJavaDebuggable() &&
+      (requested_level == InstrumentationLevel::kInstrumentWithInstrumentationStubs ||
+       instrumentation_level_ == InstrumentationLevel::kInstrumentWithInstrumentationStubs)) {
+    // In non-debuggable apps JITed code doesn't include code to call method entry / exit hooks
+    // by default. So discard all the JITed code here so we can re-jit with the support when
+    // the methods become hot again. For debuggable apps JITed code calls method entry / exit
+    // hooks when required so nothing needs to be done there. Also discard JITed code when
+    // moving away from instrumentation with stub, since the generated code disables some
+    // optimizations when instrumentation is enabled.
+    bool needs_instrumentation_support =
+        (requested_level == InstrumentationLevel::kInstrumentWithInstrumentationStubs);
+    jit::Jit* jit = Runtime::Current()->GetJit();
+    if (jit != nullptr) {
+      jit->GetCodeCache()->InvalidateAllCompiledCodeForInstrumentation(
+          needs_instrumentation_support);
+    }
+  }
+
   instrumentation_level_ = requested_level;
 }
 
diff --git a/runtime/instrumentation.h b/runtime/instrumentation.h
index a505ce3..38441e6 100644
--- a/runtime/instrumentation.h
+++ b/runtime/instrumentation.h
@@ -340,6 +340,10 @@
            instrumentation_level_ == InstrumentationLevel::kInstrumentWithInterpreter;
   }
 
+  bool NeedInstrumentationSupportForJIT() const {
+    return instrumentation_level_ == InstrumentationLevel::kInstrumentWithInstrumentationStubs;
+  }
+
   bool InterpreterStubsInstalled() const {
     return instrumentation_level_ == InstrumentationLevel::kInstrumentWithInterpreter;
   }
@@ -586,8 +590,10 @@
   // directly call entry / exit hooks and don't need the stub.
   bool CodeNeedsEntryExitStub(const void* code, ArtMethod* method);
 
-  // Update the current instrumentation_level_.
-  void UpdateInstrumentationLevel(InstrumentationLevel level);
+  // Update the current instrumentation_level_. This takes care to discard JITed
+  // code when kInstrumentWithInstrumentationStubs is requested and JITed code
+  // isn't compiled with support for calling method entry / exit hooks.
+  void UpdateInstrumentationLevel(InstrumentationLevel level) REQUIRES(Locks::mutator_lock_);
 
   // Does the job of installing or removing instrumentation code within methods.
   // In order to support multiple clients using instrumentation at the same time,
diff --git a/runtime/jit/jit_code_cache.cc b/runtime/jit/jit_code_cache.cc
index e05ca11..4b9736f 100644
--- a/runtime/jit/jit_code_cache.cc
+++ b/runtime/jit/jit_code_cache.cc
@@ -263,6 +263,7 @@
 JitCodeCache::JitCodeCache()
     : is_weak_access_enabled_(true),
       inline_cache_cond_("Jit inline cache condition variable", *Locks::jit_lock_),
+      needs_instrumentation_support_(Runtime::Current()->IsJavaDebuggable()),
       zygote_map_(&shared_region_),
       lock_cond_("Jit code cache condition variable", *Locks::jit_lock_),
       collection_in_progress_(false),
@@ -326,7 +327,13 @@
   if (method->IsPreCompiled()) {
     const void* code_ptr = nullptr;
     if (method->GetDeclaringClass()->GetClassLoader() == nullptr) {
-      code_ptr = zygote_map_.GetCodeFor(method);
+      MutexLock mu(Thread::Current(), *Locks::jit_lock_);
+      // Don't use code from zygote_map_ if we need instrumentation support in
+      // the JITed code. When instrumentation support is needed we re-jit the
+      // code by generating code to call method entry / exit hooks.
+      if (!needs_instrumentation_support_) {
+        code_ptr = zygote_map_.GetCodeFor(method);
+      }
     } else {
       MutexLock mu(Thread::Current(), *Locks::jit_lock_);
       auto it = saved_compiled_methods_map_.find(method);
@@ -645,7 +652,8 @@
                           bool is_full_debug_info,
                           CompilationKind compilation_kind,
                           bool has_should_deoptimize_flag,
-                          const ArenaSet<ArtMethod*>& cha_single_implementation_list) {
+                          const ArenaSet<ArtMethod*>& cha_single_implementation_list,
+                          bool has_instrumentation_support) {
   DCHECK(!method->IsNative() || (compilation_kind != CompilationKind::kOsr));
 
   if (!method->IsNative()) {
@@ -659,6 +667,16 @@
   const uint8_t* stack_map_data = roots_data + root_table_size;
 
   MutexLock mu(self, *Locks::jit_lock_);
+
+  // Check if a non-native was compiled with the expected instrumentation support.
+  // TODO(mythria): Currently native methods still use entry / exit stubs and we
+  // don't compile them with instrumentation support. So this check can be
+  // omitted for native methods for now. Update native methods to stop using
+  // entry / exit stubs and update this check.
+  if (!method->IsNative() && needs_instrumentation_support_ != has_instrumentation_support) {
+    return false;
+  }
+
   // We need to make sure that there will be no jit-gcs going on and wait for any ongoing one to
   // finish.
   WaitForPotentialCollectionToCompleteRunnable(self);
@@ -1712,24 +1730,45 @@
   }
 }
 
-void JitCodeCache::InvalidateAllCompiledCode() {
-  art::MutexLock mu(Thread::Current(), *Locks::jit_lock_);
+namespace {
+
+void RemoveCompiledCodeFromMethod(ArtMethod* method, ClassLinker* linker)
+    REQUIRES_SHARED(Locks::mutator_lock_) {
+  // We were compiled, so we must be warm.
+  ClearMethodCounter(method, /*was_warm=*/true);
+  if (method->IsObsolete()) {
+    linker->SetEntryPointsForObsoleteMethod(method);
+  } else {
+    linker->SetEntryPointsToInterpreter(method);
+  }
+}
+
+}  // namespace
+
+void JitCodeCache::InvalidateAllCompiledCodeWithLock() {
   VLOG(jit) << "Invalidating all compiled code";
   ClassLinker* linker = Runtime::Current()->GetClassLinker();
   for (auto it : method_code_map_) {
-    ArtMethod* meth = it.second;
-    // We were compiled, so we must be warm.
-    ClearMethodCounter(meth, /*was_warm=*/true);
-    if (meth->IsObsolete()) {
-      linker->SetEntryPointsForObsoleteMethod(meth);
-    } else {
-      linker->SetEntryPointsToInterpreter(meth);
-    }
+    RemoveCompiledCodeFromMethod(it.second, linker);
+  }
+  for (auto it : zygote_map_) {
+    RemoveCompiledCodeFromMethod(it.method, linker);
   }
   saved_compiled_methods_map_.clear();
   osr_code_map_.clear();
 }
 
+void JitCodeCache::InvalidateAllCompiledCodeForInstrumentation(bool needs_instrumentation_support) {
+  art::MutexLock mu(Thread::Current(), *Locks::jit_lock_);
+  InvalidateAllCompiledCodeWithLock();
+  needs_instrumentation_support_ = needs_instrumentation_support;
+}
+
+void JitCodeCache::InvalidateAllCompiledCode() {
+  art::MutexLock mu(Thread::Current(), *Locks::jit_lock_);
+  InvalidateAllCompiledCodeWithLock();
+}
+
 void JitCodeCache::InvalidateCompiledCodeFor(ArtMethod* method,
                                              const OatQuickMethodHeader* header) {
   DCHECK(!method->IsNative());
diff --git a/runtime/jit/jit_code_cache.h b/runtime/jit/jit_code_cache.h
index 76b7f77..712c496 100644
--- a/runtime/jit/jit_code_cache.h
+++ b/runtime/jit/jit_code_cache.h
@@ -266,7 +266,8 @@
               bool is_full_debug_info,
               CompilationKind compilation_kind,
               bool has_should_deoptimize_flag,
-              const ArenaSet<ArtMethod*>& cha_single_implementation_list)
+              const ArenaSet<ArtMethod*>& cha_single_implementation_list,
+              bool has_instrumentation_support)
       REQUIRES_SHARED(Locks::mutator_lock_)
       REQUIRES(!Locks::jit_lock_);
 
@@ -330,9 +331,15 @@
       REQUIRES(!Locks::jit_lock_)
       REQUIRES_SHARED(Locks::mutator_lock_);
 
+  void InvalidateAllCompiledCodeWithLock()
+      REQUIRES(Locks::jit_lock_)
+      REQUIRES_SHARED(Locks::mutator_lock_);
   void InvalidateAllCompiledCode()
       REQUIRES(!Locks::jit_lock_)
       REQUIRES_SHARED(Locks::mutator_lock_);
+  void InvalidateAllCompiledCodeForInstrumentation(bool needs_instrumentation_support)
+      REQUIRES(!Locks::jit_lock_)
+      REQUIRES_SHARED(Locks::mutator_lock_);
 
   void InvalidateCompiledCodeFor(ArtMethod* method, const OatQuickMethodHeader* code)
       REQUIRES(!Locks::jit_lock_)
@@ -514,6 +521,9 @@
   // Condition to wait on for accessing inline caches.
   ConditionVariable inline_cache_cond_ GUARDED_BY(Locks::jit_lock_);
 
+  // Specifies if instrumentation support is required by JITed code.
+  bool needs_instrumentation_support_ GUARDED_BY(Locks::jit_lock_);
+
   // -------------- JIT memory regions ------------------------------------- //
 
   // Shared region, inherited from the zygote.
diff --git a/runtime/trace.cc b/runtime/trace.cc
index 4b5412f..e9c8f49 100644
--- a/runtime/trace.cc
+++ b/runtime/trace.cc
@@ -426,13 +426,12 @@
             instrumentation::Instrumentation::kMethodEntered |
                 instrumentation::Instrumentation::kMethodExited |
                 instrumentation::Instrumentation::kMethodUnwind);
-        // TODO: In full-PIC mode, we don't need to fully deopt.
-        // TODO: We can only use trampoline entrypoints if we are java-debuggable since in that case
-        // we know that inlining and other problematic optimizations are disabled. We might just
-        // want to use the trampolines anyway since it is faster. It makes the story with disabling
-        // jit-gc more complex though.
-        runtime->GetInstrumentation()->EnableMethodTracing(
-            kTracerInstrumentationKey, /*needs_interpreter=*/!runtime->IsJavaDebuggable());
+        // For non-debuggable cases method tracing is best effort and we may not capture method
+        // entry / exits for the methods that are inlined. We also don't report method exits for
+        // JITed code currently on the stack. For java-debuggable it is precise since inlining is
+        // disabled and the jit code is always compiled with instrumentation support.
+        runtime->GetInstrumentation()->EnableMethodTracing(kTracerInstrumentationKey,
+                                                           /* needs_interpreter= */ false);
       }
     }
   }