jitzygote: compile system server methods at boot.

Test: m && boots
Bug: 119800099
Change-Id: Ia02a68bc3f152fe9a733577e66f327264518fe6b
diff --git a/runtime/jit/jit.cc b/runtime/jit/jit.cc
index 3087ccc..b828aaf 100644
--- a/runtime/jit/jit.cc
+++ b/runtime/jit/jit.cc
@@ -620,13 +620,79 @@
   ZygoteTask() {}
 
   void Run(Thread* self) override {
-    Runtime::Current()->GetJit()->AddNonAotBootMethodsToQueue(self);
+    Runtime* runtime = Runtime::Current();
+    std::string profile_file;
+    for (const std::string& option : runtime->GetImageCompilerOptions()) {
+      if (android::base::StartsWith(option, "--profile-file=")) {
+        profile_file = option.substr(strlen("--profile-file="));
+        break;
+      }
+    }
+
+    const std::vector<const DexFile*>& boot_class_path =
+        runtime->GetClassLinker()->GetBootClassPath();
+    ScopedNullHandle<mirror::ClassLoader> null_handle;
+    // We add to the queue for zygote so that we can fork processes in-between
+    // compilations.
+    runtime->GetJit()->CompileMethodsFromProfile(
+        self, boot_class_path, profile_file, null_handle, /* add_to_queue= */ true);
+  }
+
+  void Finalize() override {
+    delete this;
   }
 
  private:
   DISALLOW_COPY_AND_ASSIGN(ZygoteTask);
 };
 
+static std::string GetProfileFile(const std::string& dex_location) {
+  // Hardcoded assumption where the profile file is.
+  // TODO(ngeoffray): this is brittle and we would need to change change if we
+  // wanted to do more eager JITting of methods in a profile. This is
+  // currently only for system server.
+  return dex_location + ".prof";
+}
+
+class JitProfileTask final : public Task {
+ public:
+  JitProfileTask(const std::vector<std::unique_ptr<const DexFile>>& dex_files,
+                 ObjPtr<mirror::ClassLoader> class_loader) {
+    ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
+    for (const auto& dex_file : dex_files) {
+      dex_files_.push_back(dex_file.get());
+      // Register the dex file so that we can guarantee it doesn't get deleted
+      // while reading it during the task.
+      class_linker->RegisterDexFile(*dex_file.get(), class_loader);
+    }
+    ScopedObjectAccess soa(Thread::Current());
+    class_loader_ = soa.Vm()->AddGlobalRef(soa.Self(), class_loader.Ptr());
+  }
+
+  void Run(Thread* self) override {
+    ScopedObjectAccess soa(self);
+    StackHandleScope<1> hs(self);
+    Handle<mirror::ClassLoader> loader = hs.NewHandle<mirror::ClassLoader>(
+        soa.Decode<mirror::ClassLoader>(class_loader_));
+    Runtime::Current()->GetJit()->CompileMethodsFromProfile(
+        self,
+        dex_files_,
+        GetProfileFile(dex_files_[0]->GetLocation()),
+        loader,
+        /* add_to_queue= */ false);
+  }
+
+  void Finalize() override {
+    delete this;
+  }
+
+ private:
+  std::vector<const DexFile*> dex_files_;
+  jobject class_loader_;
+
+  DISALLOW_COPY_AND_ASSIGN(JitProfileTask);
+};
+
 void Jit::CreateThreadPool() {
   // There is a DCHECK in the 'AddSamples' method to ensure the tread pool
   // is not null when we instrument.
@@ -646,45 +712,49 @@
   }
 }
 
-void Jit::AddNonAotBootMethodsToQueue(Thread* self) {
-  Runtime* runtime = Runtime::Current();
-  std::string profile_location;
-  for (const std::string& option : runtime->GetImageCompilerOptions()) {
-    if (android::base::StartsWith(option, "--profile-file=")) {
-      profile_location = option.substr(strlen("--profile-file="));
-      break;
-    }
+void Jit::RegisterDexFiles(const std::vector<std::unique_ptr<const DexFile>>& dex_files,
+                           ObjPtr<mirror::ClassLoader> class_loader) {
+  if (dex_files.empty()) {
+    return;
   }
-  if (profile_location.empty()) {
-    LOG(WARNING) << "Expected a profile location in JIT zygote mode";
+  Runtime* runtime = Runtime::Current();
+  if (runtime->IsSystemServer() && runtime->IsUsingApexBootImageLocation() && UseJitCompilation()) {
+    thread_pool_->AddTask(Thread::Current(), new JitProfileTask(dex_files, class_loader));
+  }
+}
+
+void Jit::CompileMethodsFromProfile(
+    Thread* self,
+    const std::vector<const DexFile*>& dex_files,
+    const std::string& profile_file,
+    Handle<mirror::ClassLoader> class_loader,
+    bool add_to_queue) {
+
+  if (profile_file.empty()) {
+    LOG(WARNING) << "Expected a profile file in JIT zygote mode";
     return;
   }
 
   std::string error_msg;
-  ScopedFlock profile_file = LockedFile::Open(
-      profile_location.c_str(), O_RDONLY, true, &error_msg);
+  ScopedFlock profile = LockedFile::Open(
+      profile_file.c_str(), O_RDONLY, /* block= */ false, &error_msg);
 
   // Return early if we're unable to obtain a lock on the profile.
-  if (profile_file.get() == nullptr) {
+  if (profile.get() == nullptr) {
     LOG(ERROR) << "Cannot lock profile: " << error_msg;
     return;
   }
 
   ProfileCompilationInfo profile_info;
-  if (!profile_info.Load(profile_file->Fd())) {
+  if (!profile_info.Load(profile->Fd())) {
     LOG(ERROR) << "Could not load profile file";
     return;
   }
-
-  const std::vector<const DexFile*>& boot_class_path =
-      runtime->GetClassLinker()->GetBootClassPath();
   ScopedObjectAccess soa(self);
   StackHandleScope<1> hs(self);
   MutableHandle<mirror::DexCache> dex_cache = hs.NewHandle<mirror::DexCache>(nullptr);
-  ScopedNullHandle<mirror::ClassLoader> null_handle;
-  ClassLinker* class_linker = runtime->GetClassLinker();
-
-  for (const DexFile* dex_file : boot_class_path) {
+  ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
+  for (const DexFile* dex_file : dex_files) {
     if (LocationIsOnRuntimeModule(dex_file->GetLocation().c_str())) {
       // The runtime module jars are already preopted.
       continue;
@@ -714,7 +784,7 @@
 
     for (uint16_t method_idx : all_methods) {
       ArtMethod* method = class_linker->ResolveMethodWithoutInvokeType(
-          method_idx, dex_cache, null_handle);
+          method_idx, dex_cache, class_loader);
       if (method == nullptr) {
         self->ClearException();
         continue;
@@ -733,7 +803,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.
-        if (method->GetDeclaringClass()->DescriptorEquals(
+        if (!add_to_queue || method->GetDeclaringClass()->DescriptorEquals(
                 "Lcom/android/internal/os/ZygoteServer;")) {
           CompileMethod(method, self, /* baseline= */ false, /* osr= */ false);
         } else {
@@ -964,6 +1034,11 @@
       // We keep the queue for system server, as not having those methods compiled
       // impacts app startup.
       thread_pool_->RemoveAllTasks(Thread::Current());
+    } else if (Runtime::Current()->IsUsingApexBootImageLocation() && UseJitCompilation()) {
+      // Disable garbage collection: we don't want it to delete methods we're compiling
+      // through boot and system server profiles.
+      // TODO(ngeoffray): Fix this so we still collect deoptimized and unused code.
+      code_cache_->SetGarbageCollectCode(false);
     }
 
     // Resume JIT compilation.
diff --git a/runtime/jit/jit.h b/runtime/jit/jit.h
index 15cdacb..4b81f71 100644
--- a/runtime/jit/jit.h
+++ b/runtime/jit/jit.h
@@ -21,6 +21,7 @@
 #include "base/macros.h"
 #include "base/mutex.h"
 #include "base/timing_logger.h"
+#include "handle.h"
 #include "jit/profile_saver_options.h"
 #include "obj_ptr.h"
 #include "thread_pool.h"
@@ -29,6 +30,7 @@
 
 class ArtMethod;
 class ClassLinker;
+class DexFile;
 class OatDexFile;
 struct RuntimeArgumentMap;
 union JValue;
@@ -36,6 +38,7 @@
 namespace mirror {
 class Object;
 class Class;
+class ClassLoader;
 }   // namespace mirror
 
 namespace jit {
@@ -304,9 +307,19 @@
   // Adjust state after forking.
   void PostZygoteFork();
 
-  // In case the boot classpath is not fully AOTed, add methods from the boot profile to the
-  // compilation queue.
-  void AddNonAotBootMethodsToQueue(Thread* self);
+  // Compile methods from the given profile. If `add_to_queue` is true, methods
+  // in the profile are added to the JIT queue. Otherwise they are compiled
+  // directly.
+  void CompileMethodsFromProfile(Thread* self,
+                                 const std::vector<const DexFile*>& dex_files,
+                                 const std::string& profile_path,
+                                 Handle<mirror::ClassLoader> class_loader,
+                                 bool add_to_queue);
+
+  // Register the dex files to the JIT. This is to perform any compilation/optimization
+  // at the point of loading the dex files.
+  void RegisterDexFiles(const std::vector<std::unique_ptr<const DexFile>>& dex_files,
+                        ObjPtr<mirror::ClassLoader> class_loader);
 
  private:
   Jit(JitCodeCache* code_cache, JitOptions* options);
diff --git a/runtime/native/dalvik_system_ZygoteHooks.cc b/runtime/native/dalvik_system_ZygoteHooks.cc
index 77ec795..de28c28 100644
--- a/runtime/native/dalvik_system_ZygoteHooks.cc
+++ b/runtime/native/dalvik_system_ZygoteHooks.cc
@@ -257,6 +257,7 @@
 
 static void ZygoteHooks_nativePostForkSystemServer(JNIEnv* env ATTRIBUTE_UNUSED,
                                                    jclass klass ATTRIBUTE_UNUSED) {
+  Runtime::Current()->SetSystemServer(true);
   // This JIT code cache for system server is created whilst the runtime is still single threaded.
   // System server has a window where it can create executable pages for this purpose, but this is
   // turned off after this hook. Consequently, the only JIT mode supported is the dual-view JIT
diff --git a/runtime/oat_file_manager.cc b/runtime/oat_file_manager.cc
index 0543644..ed55110 100644
--- a/runtime/oat_file_manager.cc
+++ b/runtime/oat_file_manager.cc
@@ -41,6 +41,7 @@
 #include "gc/scoped_gc_critical_section.h"
 #include "gc/space/image_space.h"
 #include "handle_scope-inl.h"
+#include "jit/jit.h"
 #include "jni/java_vm_ext.h"
 #include "jni/jni_internal.h"
 #include "mirror/class_loader.h"
@@ -645,6 +646,12 @@
     }
   }
 
+  if (Runtime::Current()->GetJit() != nullptr) {
+    ScopedObjectAccess soa(self);
+    Runtime::Current()->GetJit()->RegisterDexFiles(
+        dex_files, soa.Decode<mirror::ClassLoader>(class_loader));
+  }
+
   return dex_files;
 }
 
diff --git a/runtime/runtime.cc b/runtime/runtime.cc
index 62b38f1..426a5f1 100644
--- a/runtime/runtime.cc
+++ b/runtime/runtime.cc
@@ -229,6 +229,7 @@
       instruction_set_(InstructionSet::kNone),
       compiler_callbacks_(nullptr),
       is_zygote_(false),
+      is_system_server_(false),
       must_relocate_(false),
       is_concurrent_gc_enabled_(true),
       is_explicit_gc_disabled_(false),
diff --git a/runtime/runtime.h b/runtime/runtime.h
index a20aef4..6df9e3e 100644
--- a/runtime/runtime.h
+++ b/runtime/runtime.h
@@ -164,6 +164,14 @@
     return is_zygote_;
   }
 
+  bool IsSystemServer() const {
+    return is_system_server_;
+  }
+
+  void SetSystemServer(bool value) {
+    is_system_server_ = value;
+  }
+
   bool IsExplicitGcDisabled() const {
     return is_explicit_gc_disabled_;
   }
@@ -937,6 +945,7 @@
 
   CompilerCallbacks* compiler_callbacks_;
   bool is_zygote_;
+  bool is_system_server_;
   bool must_relocate_;
   bool is_concurrent_gc_enabled_;
   bool is_explicit_gc_disabled_;