jitzygote: Handle case of methods having the resolution stub.

We need to cache the compiled code until the class is initialized.

Test: test.py
Test: jitzygote config boots
Bug: 119800099
Change-Id: Ib7a1efc3e101d9dabfad963621a1fe4a142e9b29
diff --git a/libdexfile/dex/modifiers.h b/libdexfile/dex/modifiers.h
index bc232d5..5caa402 100644
--- a/libdexfile/dex/modifiers.h
+++ b/libdexfile/dex/modifiers.h
@@ -79,9 +79,9 @@
 // Set by the verifier for a method we do not want the compiler to compile.
 static constexpr uint32_t kAccCompileDontBother =     0x02000000;  // method (runtime)
 
-// Used in conjunction with kAccCompileDontBother to mark the method as zygote
-// compiled.
-static constexpr uint32_t kAccZygoteCompiled =        0x00200000;  // method (runtime)
+// Used in conjunction with kAccCompileDontBother to mark the method as pre
+// compiled by the JIT compiler.
+static constexpr uint32_t kAccPreCompiled =           0x00200000;  // method (runtime)
 
 // Set by the verifier for a method that could not be verified to follow structured locking.
 static constexpr uint32_t kAccMustCountLocks =        0x04000000;  // method (runtime)
diff --git a/runtime/art_method.h b/runtime/art_method.h
index 9c74d83..450fe41 100644
--- a/runtime/art_method.h
+++ b/runtime/art_method.h
@@ -212,19 +212,19 @@
     return !IsAbstract() && !IsDefaultConflicting();
   }
 
-  bool IsZygoteCompiled() {
-    uint32_t expected = (kAccZygoteCompiled | kAccCompileDontBother);
+  bool IsPreCompiled() {
+    uint32_t expected = (kAccPreCompiled | kAccCompileDontBother);
     return (GetAccessFlags() & expected) == expected;
   }
 
-  void SetZygoteCompiled() {
+  void SetPreCompiled() {
     DCHECK(IsInvokable());
     DCHECK(IsCompilable());
-    AddAccessFlags(kAccZygoteCompiled | kAccCompileDontBother);
+    AddAccessFlags(kAccPreCompiled | kAccCompileDontBother);
   }
 
-  void ClearZygoteCompiled() {
-    ClearAccessFlags(kAccZygoteCompiled | kAccCompileDontBother);
+  void ClearPreCompiled() {
+    ClearAccessFlags(kAccPreCompiled | kAccCompileDontBother);
   }
 
   bool IsCompilable() {
@@ -232,7 +232,7 @@
       // kAccCompileDontBother overlaps with kAccIntrinsicBits.
       return true;
     }
-    if (IsZygoteCompiled()) {
+    if (IsPreCompiled()) {
       return true;
     }
     return (GetAccessFlags() & kAccCompileDontBother) == 0;
diff --git a/runtime/class_linker.cc b/runtime/class_linker.cc
index 2d19bfc..c06e5e8 100644
--- a/runtime/class_linker.cc
+++ b/runtime/class_linker.cc
@@ -3345,10 +3345,19 @@
   if (method->IsProxyMethod()) {
     return GetQuickProxyInvokeHandler();
   }
-  auto* code = method->GetOatMethodQuickCode(GetImagePointerSize());
+  const void* code = method->GetOatMethodQuickCode(GetImagePointerSize());
   if (code != nullptr) {
     return code;
   }
+
+  jit::Jit* jit = Runtime::Current()->GetJit();
+  if (jit != nullptr) {
+    code = jit->GetCodeCache()->GetSavedEntryPointOfPreCompiledMethod(method);
+    if (code != nullptr) {
+      return code;
+    }
+  }
+
   if (method->IsNative()) {
     // No code and native? Use generic trampoline.
     return GetQuickGenericJniStub();
@@ -3453,8 +3462,9 @@
       quick_code = oat_method.GetQuickCode();
     }
     // Check if we have JIT compiled code for it.
-    if (quick_code == nullptr && Runtime::Current()->GetJit() != nullptr) {
-      quick_code = Runtime::Current()->GetJit()->GetCodeCache()->GetZygoteSavedEntryPoint(method);
+    jit::Jit* jit = Runtime::Current()->GetJit();
+    if (quick_code == nullptr && jit != nullptr) {
+      quick_code = jit->GetCodeCache()->GetSavedEntryPointOfPreCompiledMethod(method);
     }
     // Check whether the method is native, in which case it's generic JNI.
     if (quick_code == nullptr && method->IsNative()) {
diff --git a/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc b/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc
index e589fc9..2078186 100644
--- a/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc
+++ b/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc
@@ -1491,12 +1491,7 @@
       } else if (invoke_type == kStatic) {
         // Class is still initializing, go to JIT or oat and grab code (trampoline must be
         // left in place until class is initialized to stop races between threads).
-        if (Runtime::Current()->GetJit() != nullptr) {
-          code = Runtime::Current()->GetJit()->GetCodeCache()->GetZygoteSavedEntryPoint(called);
-        }
-        if (code == nullptr) {
-          code = linker->GetQuickOatCodeFor(called);
-        }
+        code = linker->GetQuickOatCodeFor(called);
       } else {
         // No trampoline for non-static methods.
         code = called->GetEntryPointFromQuickCompiledCode();
diff --git a/runtime/interpreter/interpreter_common.cc b/runtime/interpreter/interpreter_common.cc
index 9a5adda..15c0ccf 100644
--- a/runtime/interpreter/interpreter_common.cc
+++ b/runtime/interpreter/interpreter_common.cc
@@ -89,7 +89,7 @@
   if ((profiling_info != nullptr) && (profiling_info->GetSavedEntryPoint() != nullptr)) {
     return false;
   }
-  if (method->IsZygoteCompiled()) {
+  if (method->IsPreCompiled()) {
     return false;
   }
   return true;
diff --git a/runtime/jit/jit-inl.h b/runtime/jit/jit-inl.h
index 210a039..ee70cdf 100644
--- a/runtime/jit/jit-inl.h
+++ b/runtime/jit/jit-inl.h
@@ -35,7 +35,7 @@
                             ArtMethod* method,
                             uint16_t samples,
                             bool with_backedges) {
-  if (method->IsZygoteCompiled()) {
+  if (method->IsPreCompiled()) {
     return;
   }
   if (Jit::ShouldUsePriorityThreadWeight(self)) {
diff --git a/runtime/jit/jit.cc b/runtime/jit/jit.cc
index b6d2687..737acda 100644
--- a/runtime/jit/jit.cc
+++ b/runtime/jit/jit.cc
@@ -826,6 +826,7 @@
         // Special case ZygoteServer class so that it gets compiled before the
         // zygote enters it. This avoids needing to do OSR during app startup.
         // TODO: have a profile instead.
+        method->SetPreCompiled();
         if (!add_to_queue || method->GetDeclaringClass()->DescriptorEquals(
                 "Lcom/android/internal/os/ZygoteServer;")) {
           CompileMethod(method, self, /* baseline= */ false, /* osr= */ false, /* prejit= */ true);
@@ -833,9 +834,6 @@
           ++added_to_queue;
           thread_pool_->AddTask(self,
               new JitCompileTask(method, JitCompileTask::TaskKind::kPreCompile));
-          if (Runtime::Current()->IsZygote()) {
-            method->SetZygoteCompiled();
-          }
         }
       }
     }
@@ -844,7 +842,7 @@
 }
 
 static bool IgnoreSamplesForMethod(ArtMethod* method) REQUIRES_SHARED(Locks::mutator_lock_) {
-  if (method->IsClassInitializer() || !method->IsCompilable() || method->IsZygoteCompiled()) {
+  if (method->IsClassInitializer() || !method->IsCompilable() || method->IsPreCompiled()) {
     // We do not want to compile such methods.
     return true;
   }
@@ -973,7 +971,7 @@
     return;
   }
 
-  if (UNLIKELY(method->IsZygoteCompiled())) {
+  if (UNLIKELY(method->IsPreCompiled())) {
     const void* code_ptr = code_cache_->GetZygoteMap()->GetCodeFor(method);
     if (code_ptr != nullptr) {
       Runtime::Current()->GetInstrumentation()->UpdateMethodsCode(method, code_ptr);
diff --git a/runtime/jit/jit_code_cache.cc b/runtime/jit/jit_code_cache.cc
index ffb0117..3669973 100644
--- a/runtime/jit/jit_code_cache.cc
+++ b/runtime/jit/jit_code_cache.cc
@@ -303,24 +303,17 @@
   return info->GetSavedEntryPoint();
 }
 
-const void* JitCodeCache::GetZygoteSavedEntryPoint(ArtMethod* method) {
-  if (Runtime::Current()->IsUsingApexBootImageLocation() &&
-      // Currently only applies to boot classpath
-      method->GetDeclaringClass()->GetClassLoader() == nullptr) {
-    const void* entry_point = nullptr;
-    if (method->IsNative()) {
-      const void* code_ptr = GetJniStubCode(method);
-      if (code_ptr != nullptr) {
-        entry_point = OatQuickMethodHeader::FromCodePointer(code_ptr)->GetEntryPoint();
-      }
+const void* JitCodeCache::GetSavedEntryPointOfPreCompiledMethod(ArtMethod* method) {
+  if (Runtime::Current()->IsUsingApexBootImageLocation() && method->IsPreCompiled()) {
+    if (method->GetDeclaringClass()->GetClassLoader() == nullptr) {
+      return zygote_map_.GetCodeFor(method);
     } else {
-      ProfilingInfo* profiling_info = method->GetProfilingInfo(kRuntimePointerSize);
-      if (profiling_info != nullptr) {
-        entry_point = profiling_info->GetSavedEntryPoint();
+      MutexLock mu(Thread::Current(), *Locks::jit_lock_);
+      auto it = saved_compiled_methods_map_.find(method);
+      if (it != saved_compiled_methods_map_.end()) {
+        return it->second;
       }
-    }
-    if (Runtime::Current()->IsZygote() || IsInZygoteExecSpace(entry_point)) {
-      return entry_point;
+      return nullptr;
     }
   }
   return nullptr;
@@ -751,7 +744,7 @@
         }
       }
     } else {
-      if (method->IsZygoteCompiled() && IsSharedRegion(*region)) {
+      if (method->IsPreCompiled() && IsSharedRegion(*region)) {
         zygote_map_.Put(code_ptr, method);
       } else {
         method_code_map_.Put(code_ptr, method);
@@ -764,13 +757,11 @@
         // This situation currently only occurs in the jit-zygote mode.
         DCHECK(Runtime::Current()->IsUsingApexBootImageLocation());
         DCHECK(!garbage_collect_code_);
-        // TODO(ngeoffray): In most cases, the zygote will not have a profiling
-        // info for a compiled method. Use a map instead.
-        if (method->GetProfilingInfo(kRuntimePointerSize) != nullptr) {
-          // Save the entrypoint, so it can be fetched later once the class is
-          // initialized.
-          method->GetProfilingInfo(kRuntimePointerSize)->SetSavedEntryPoint(
-              method_header->GetEntryPoint());
+        DCHECK(method->IsPreCompiled());
+        // The shared region can easily be queried. For the private region, we
+        // use a side map.
+        if (!IsSharedRegion(*region)) {
+          saved_compiled_methods_map_.Put(method, code_ptr);
         }
       } else {
         Runtime::Current()->GetInstrumentation()->UpdateMethodsCode(
@@ -1715,10 +1706,10 @@
     }
   }
 
-  // In case the method was compiled by the zygote, clear that information so we
+  // In case the method was pre-compiled, clear that information so we
   // can recompile it ourselves.
-  if (method->IsZygoteCompiled()) {
-    method->ClearZygoteCompiled();
+  if (method->IsPreCompiled()) {
+    method->ClearPreCompiled();
   }
 }
 
diff --git a/runtime/jit/jit_code_cache.h b/runtime/jit/jit_code_cache.h
index 20b74d6..284a36a 100644
--- a/runtime/jit/jit_code_cache.h
+++ b/runtime/jit/jit_code_cache.h
@@ -316,9 +316,9 @@
       REQUIRES(!Locks::jit_lock_)
       REQUIRES_SHARED(Locks::mutator_lock_);
 
-  // Fetch the entrypoint that zygote may have saved for a method. The zygote saves an entrypoint
-  // only for the case when the method's declaring class is not initialized.
-  const void* GetZygoteSavedEntryPoint(ArtMethod* method)
+  // Fetch the code of a method that was JITted, but the JIT could not
+  // update its entrypoint due to the resolution trampoline.
+  const void* GetSavedEntryPointOfPreCompiledMethod(ArtMethod* method)
       REQUIRES(!Locks::jit_lock_)
       REQUIRES_SHARED(Locks::mutator_lock_);
 
@@ -466,6 +466,10 @@
   // Holds compiled code associated to the ArtMethod.
   SafeMap<const void*, ArtMethod*> method_code_map_ GUARDED_BY(Locks::jit_lock_);
 
+  // Holds compiled code associated to the ArtMethod. Used when pre-jitting
+  // methods whose entrypoints have the resolution stub.
+  SafeMap<ArtMethod*, const void*> saved_compiled_methods_map_ GUARDED_BY(Locks::jit_lock_);
+
   // Holds osr compiled code associated to the ArtMethod.
   SafeMap<ArtMethod*, const void*> osr_code_map_ GUARDED_BY(Locks::jit_lock_);