Save the JIT compiled code when the class isn't initialized yet.

When a class isn't initialized, we put resolution stubs to its
static methods. In order to fetch the compiled code at the point of
initialization, save it in the profiling info.

Bug: 119800099
Test: start an app, no JIT compilation of boot classpath methods
Change-Id: I32b947318dbcb1010c94a11b51ea39d992d247e3
diff --git a/runtime/class_linker.cc b/runtime/class_linker.cc
index 658fec3..fa5cb92 100644
--- a/runtime/class_linker.cc
+++ b/runtime/class_linker.cc
@@ -3380,6 +3380,10 @@
       OatFile::OatMethod oat_method = oat_class.GetOatMethod(method_index);
       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);
+    }
     // Check whether the method is native, in which case it's generic JNI.
     if (quick_code == nullptr && method->IsNative()) {
       quick_code = GetQuickGenericJniStub();
diff --git a/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc b/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc
index 160dff1..a2420af 100644
--- a/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc
+++ b/runtime/entrypoints/quick/quick_trampoline_entrypoints.cc
@@ -1485,9 +1485,14 @@
         // with the interpreter.
         code = GetQuickToInterpreterBridge();
       } else if (invoke_type == kStatic) {
-        // Class is still initializing, go to oat and grab code (trampoline must be left in place
-        // until class is initialized to stop races between threads).
-        code = linker->GetQuickOatCodeFor(called);
+        // 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);
+        }
       } else {
         // No trampoline for non-static methods.
         code = called->GetEntryPointFromQuickCompiledCode();
diff --git a/runtime/jit/jit.cc b/runtime/jit/jit.cc
index e775eb6..2eff8f4 100644
--- a/runtime/jit/jit.cc
+++ b/runtime/jit/jit.cc
@@ -727,7 +727,8 @@
       // we currently have no place to safely store the compiled code, we just don't
       // compile it for now.
       if (class_linker->IsQuickToInterpreterBridge(entry_point) ||
-          class_linker->IsQuickGenericJniStub(entry_point)) {
+          class_linker->IsQuickGenericJniStub(entry_point) ||
+          class_linker->IsQuickResolutionStub(entry_point)) {
         if (!method->IsNative()) {
           // The compiler requires a ProfilingInfo object for non-native methods.
           ProfilingInfo::Create(self, method, /* retry_allocation= */ true);
diff --git a/runtime/jit/jit_code_cache.cc b/runtime/jit/jit_code_cache.cc
index 2052f14..ff08a45 100644
--- a/runtime/jit/jit_code_cache.cc
+++ b/runtime/jit/jit_code_cache.cc
@@ -545,6 +545,26 @@
   return info->GetSavedEntryPoint();
 }
 
+const void* JitCodeCache::GetZygoteSavedEntryPoint(ArtMethod* method) {
+  if (!Runtime::Current()->IsUsingDefaultBootImageLocation() &&
+      // 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();
+      }
+    } else if (method->GetProfilingInfo(kRuntimePointerSize) != nullptr) {
+      entry_point = method->GetProfilingInfo(kRuntimePointerSize)->GetSavedEntryPoint();
+    }
+    if (Runtime::Current()->IsZygote() || IsInZygoteExecSpace(entry_point)) {
+      return entry_point;
+    }
+  }
+  return nullptr;
+}
+
 class ScopedCodeCacheWrite : ScopedTrace {
  public:
   explicit ScopedCodeCacheWrite(const JitCodeCache* const code_cache)
@@ -1066,9 +1086,9 @@
     DCHECK(cha_single_implementation_list.empty() || !Runtime::Current()->IsJavaDebuggable())
         << "Should not be using cha on debuggable apps/runs!";
 
+    ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
     for (ArtMethod* single_impl : cha_single_implementation_list) {
-      Runtime::Current()->GetClassLinker()->GetClassHierarchyAnalysis()->AddDependency(
-          single_impl, method, method_header);
+      class_linker->GetClassHierarchyAnalysis()->AddDependency(single_impl, method, method_header);
     }
 
     if (UNLIKELY(method->IsNative())) {
@@ -1081,7 +1101,9 @@
       data->SetCode(code_ptr);
       instrumentation::Instrumentation* instrum = Runtime::Current()->GetInstrumentation();
       for (ArtMethod* m : data->GetMethods()) {
-        instrum->UpdateMethodsCode(m, method_header->GetEntryPoint());
+        if (!class_linker->IsQuickResolutionStub(m->GetEntryPointFromQuickCompiledCode())) {
+          instrum->UpdateMethodsCode(m, method_header->GetEntryPoint());
+        }
       }
     } else {
       // Fill the root table before updating the entry point.
@@ -1096,6 +1118,17 @@
       if (osr) {
         number_of_osr_compilations_++;
         osr_code_map_.Put(method, code_ptr);
+      } else if (class_linker->IsQuickResolutionStub(
+          method->GetEntryPointFromQuickCompiledCode())) {
+        // This situation currently only occurs in the jit-zygote mode.
+        DCHECK(Runtime::Current()->IsZygote());
+        DCHECK(!Runtime::Current()->IsUsingDefaultBootImageLocation());
+        DCHECK(method->GetProfilingInfo(kRuntimePointerSize) != nullptr);
+        DCHECK(method->GetDeclaringClass()->GetClassLoader() == nullptr);
+        // Save the entrypoint, so it can be fethed later once the class is
+        // initialized.
+        method->GetProfilingInfo(kRuntimePointerSize)->SetSavedEntryPoint(
+            method_header->GetEntryPoint());
       } else {
         Runtime::Current()->GetInstrumentation()->UpdateMethodsCode(
             method, method_header->GetEntryPoint());
@@ -1979,14 +2012,16 @@
 
   ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
   if (class_linker->IsQuickResolutionStub(method->GetEntryPointFromQuickCompiledCode())) {
-    // We currently don't save the JIT compiled code if we cannot update the entrypoint due
-    // to having the resolution stub.
-    VLOG(jit) << "Not compiling "
-              << method->PrettyMethod()
-              << " because it has the resolution stub";
-    // Give it a new chance to be hot.
-    ClearMethodCounter(method, /*was_warm=*/ false);
-    return false;
+    if (Runtime::Current()->IsUsingDefaultBootImageLocation() || !Runtime::Current()->IsZygote()) {
+      // Unless we're running as zygote in the jitzygote experiment, we currently don't save
+      // the JIT compiled code if we cannot update the entrypoint due to having the resolution stub.
+      VLOG(jit) << "Not compiling "
+                << method->PrettyMethod()
+                << " because it has the resolution stub";
+      // Give it a new chance to be hot.
+      ClearMethodCounter(method, /*was_warm=*/ false);
+      return false;
+    }
   }
 
   MutexLock mu(self, lock_);
@@ -2017,7 +2052,9 @@
       for (ArtMethod* m : data->GetMethods()) {
         // Call the dedicated method instead of the more generic UpdateMethodsCode, because
         // `m` might be in the process of being deleted.
-        instrumentation->UpdateNativeMethodsCodeToJitCode(m, entrypoint);
+        if (!class_linker->IsQuickResolutionStub(m->GetEntryPointFromQuickCompiledCode())) {
+          instrumentation->UpdateNativeMethodsCodeToJitCode(m, entrypoint);
+        }
       }
       if (collection_in_progress_) {
         CHECK(!IsInZygoteExecSpace(data->GetCode()));
diff --git a/runtime/jit/jit_code_cache.h b/runtime/jit/jit_code_cache.h
index 8a6cebe..fabb978 100644
--- a/runtime/jit/jit_code_cache.h
+++ b/runtime/jit/jit_code_cache.h
@@ -262,6 +262,12 @@
       REQUIRES(!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)
+      REQUIRES(!lock_)
+      REQUIRES_SHARED(Locks::mutator_lock_);
+
   void PostForkChildAction(bool is_system_server, bool is_zygote);
 
   // Clear the entrypoints of JIT compiled methods that belong in the zygote space.