Store zygote compiled code in a shared map.

- Store flags in the ArtMethod to know that it is compiled by the zygote
- Query the map when entering a zygote compiled method from the
interpreter.

Bug: 119800099
Test: boots
Change-Id: Ib1a38266573e28d371034d02d6bb83f9b8b2e317
diff --git a/runtime/jit/jit-inl.h b/runtime/jit/jit-inl.h
index 80324ad..210a039 100644
--- a/runtime/jit/jit-inl.h
+++ b/runtime/jit/jit-inl.h
@@ -35,6 +35,9 @@
                             ArtMethod* method,
                             uint16_t samples,
                             bool with_backedges) {
+  if (method->IsZygoteCompiled()) {
+    return;
+  }
   if (Jit::ShouldUsePriorityThreadWeight(self)) {
     samples *= PriorityThreadWeight();
   }
diff --git a/runtime/jit/jit.cc b/runtime/jit/jit.cc
index 382b333..d497a52 100644
--- a/runtime/jit/jit.cc
+++ b/runtime/jit/jit.cc
@@ -268,6 +268,12 @@
   }
 
   JitMemoryRegion* region = GetCodeCache()->GetCurrentRegion();
+  if (osr && GetCodeCache()->IsSharedRegion(*region)) {
+    VLOG(jit) << "JIT not osr compiling "
+              << method->PrettyMethod()
+              << " due to using shared region";
+    return false;
+  }
 
   // If we get a request to compile a proxy method, we pass the actual Java method
   // of that proxy method, as the compiler does not expect a proxy method.
@@ -640,8 +646,11 @@
     ScopedNullHandle<mirror::ClassLoader> null_handle;
     // We add to the queue for zygote so that we can fork processes in-between
     // compilations.
-    runtime->GetJit()->CompileMethodsFromProfile(
+    uint32_t added_to_queue = runtime->GetJit()->CompileMethodsFromProfile(
         self, boot_class_path, profile_file, null_handle, /* add_to_queue= */ true);
+
+    JitCodeCache* code_cache = runtime->GetJit()->GetCodeCache();
+    code_cache->GetZygoteMap()->Initialize(added_to_queue);
   }
 
   void Finalize() override {
@@ -738,7 +747,7 @@
   }
 }
 
-void Jit::CompileMethodsFromProfile(
+uint32_t Jit::CompileMethodsFromProfile(
     Thread* self,
     const std::vector<const DexFile*>& dex_files,
     const std::string& profile_file,
@@ -747,7 +756,7 @@
 
   if (profile_file.empty()) {
     LOG(WARNING) << "Expected a profile file in JIT zygote mode";
-    return;
+    return 0u;
   }
 
   std::string error_msg;
@@ -757,18 +766,19 @@
   // Return early if we're unable to obtain a lock on the profile.
   if (profile.get() == nullptr) {
     LOG(ERROR) << "Cannot lock profile: " << error_msg;
-    return;
+    return 0u;
   }
 
   ProfileCompilationInfo profile_info;
   if (!profile_info.Load(profile->Fd())) {
     LOG(ERROR) << "Could not load profile file";
-    return;
+    return 0u;
   }
   ScopedObjectAccess soa(self);
   StackHandleScope<1> hs(self);
   MutableHandle<mirror::DexCache> dex_cache = hs.NewHandle<mirror::DexCache>(nullptr);
   ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
+  uint32_t added_to_queue = 0u;
   for (const DexFile* dex_file : dex_files) {
     if (LocationIsOnRuntimeModule(dex_file->GetLocation().c_str())) {
       // The runtime module jars are already preopted.
@@ -818,16 +828,21 @@
                 "Lcom/android/internal/os/ZygoteServer;")) {
           CompileMethod(method, self, /* baseline= */ false, /* osr= */ false, /* prejit= */ true);
         } else {
+          ++added_to_queue;
           thread_pool_->AddTask(self,
               new JitCompileTask(method, JitCompileTask::TaskKind::kPreCompile));
+          if (Runtime::Current()->IsZygote()) {
+            method->SetZygoteCompiled();
+          }
         }
       }
     }
   }
+  return added_to_queue;
 }
 
 static bool IgnoreSamplesForMethod(ArtMethod* method) REQUIRES_SHARED(Locks::mutator_lock_) {
-  if (method->IsClassInitializer() || !method->IsCompilable()) {
+  if (method->IsClassInitializer() || !method->IsCompilable() || method->IsZygoteCompiled()) {
     // We do not want to compile such methods.
     return true;
   }
@@ -956,6 +971,14 @@
     return;
   }
 
+  if (UNLIKELY(method->IsZygoteCompiled())) {
+    const void* code_ptr = code_cache_->GetZygoteMap()->GetCodeFor(method);
+    if (code_ptr != nullptr) {
+      Runtime::Current()->GetInstrumentation()->UpdateMethodsCode(method, code_ptr);
+      return;
+    }
+  }
+
   ProfilingInfo* profiling_info = method->GetProfilingInfo(kRuntimePointerSize);
   // Update the entrypoint if the ProfilingInfo has one. The interpreter will call it
   // instead of interpreting the method. We don't update it for instrumentation as the entrypoint
@@ -1033,12 +1056,11 @@
       !Runtime::Current()->GetInstrumentation()->AreExitStubsInstalled());
 
   if (thread_pool_ != nullptr) {
-    if (!is_system_server) {
-      // Remove potential tasks that have been inherited from the zygote.
-      // 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()) {
+    // Remove potential tasks that have been inherited from the zygote.
+    thread_pool_->RemoveAllTasks(Thread::Current());
+    if (is_system_server &&
+        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.
diff --git a/runtime/jit/jit.h b/runtime/jit/jit.h
index d272a18..41876ca 100644
--- a/runtime/jit/jit.h
+++ b/runtime/jit/jit.h
@@ -313,11 +313,12 @@
   // 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);
+  // Return the number of methods added to the queue.
+  uint32_t 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.
diff --git a/runtime/jit/jit_code_cache.cc b/runtime/jit/jit_code_cache.cc
index 960018c..a29201c 100644
--- a/runtime/jit/jit_code_cache.cc
+++ b/runtime/jit/jit_code_cache.cc
@@ -222,6 +222,7 @@
 JitCodeCache::JitCodeCache()
     : is_weak_access_enabled_(true),
       inline_cache_cond_("Jit inline cache condition variable", *Locks::jit_lock_),
+      zygote_map_(&shared_region_),
       lock_cond_("Jit code cache condition variable", *Locks::jit_lock_),
       collection_in_progress_(false),
       last_collection_increased_code_cache_(false),
@@ -266,6 +267,9 @@
         return true;
       }
     }
+    if (zygote_map_.ContainsMethod(method)) {
+      return true;
+    }
   }
   return false;
 }
@@ -747,7 +751,11 @@
         }
       }
     } else {
-      method_code_map_.Put(code_ptr, method);
+      if (method->IsZygoteCompiled() && IsSharedRegion(*region)) {
+        zygote_map_.Put(code_ptr, method);
+      } else {
+        method_code_map_.Put(code_ptr, method);
+      }
       if (osr) {
         number_of_osr_compilations_++;
         osr_code_map_.Put(method, code_ptr);
@@ -1351,6 +1359,12 @@
       return nullptr;
     }
   } else {
+    if (shared_region_.IsInExecSpace(reinterpret_cast<const void*>(pc))) {
+      const void* code_ptr = zygote_map_.GetCodeFor(method, pc);
+      if (code_ptr != nullptr) {
+        return OatQuickMethodHeader::FromCodePointer(code_ptr);
+      }
+    }
     auto it = method_code_map_.lower_bound(reinterpret_cast<const void*>(pc));
     if (it != method_code_map_.begin()) {
       --it;
@@ -1700,6 +1714,12 @@
       osr_code_map_.erase(it);
     }
   }
+
+  // In case the method was compiled by the zygote, clear that information so we
+  // can recompile it ourselves.
+  if (method->IsZygoteCompiled()) {
+    method->ClearZygoteCompiled();
+  }
 }
 
 void JitCodeCache::Dump(std::ostream& os) {
@@ -1755,5 +1775,91 @@
   return Runtime::Current()->IsZygote() ? &shared_region_ : &private_region_;
 }
 
+void ZygoteMap::Initialize(uint32_t number_of_methods) {
+  MutexLock mu(Thread::Current(), *Locks::jit_lock_);
+  // Allocate for 40-80% capacity. This will offer OK lookup times, and termination
+  // cases.
+  size_t capacity = RoundUpToPowerOfTwo(number_of_methods * 100 / 80);
+  Entry* data = reinterpret_cast<Entry*>(region_->AllocateData(capacity * sizeof(Entry)));
+  if (data != nullptr) {
+    region_->FillData(data, capacity, Entry { nullptr, nullptr });
+    map_ = ArrayRef(data, capacity);
+  }
+}
+
+const void* ZygoteMap::GetCodeFor(ArtMethod* method, uintptr_t pc) const {
+  if (map_.empty()) {
+    return nullptr;
+  }
+
+  if (method == nullptr) {
+    // Do a linear search. This should only be used in debug builds.
+    CHECK(kIsDebugBuild);
+    for (const Entry& entry : map_) {
+      const void* code_ptr = entry.code_ptr;
+      if (code_ptr != nullptr) {
+        OatQuickMethodHeader* method_header = OatQuickMethodHeader::FromCodePointer(code_ptr);
+        if (method_header->Contains(pc)) {
+          return code_ptr;
+        }
+      }
+    }
+    return nullptr;
+  }
+
+  std::hash<ArtMethod*> hf;
+  size_t index = hf(method) & (map_.size() - 1u);
+  size_t original_index = index;
+  // Loop over the array: we know this loop terminates as we will either
+  // encounter the given method, or a null entry. Both terminate the loop.
+  // Note that the zygote may concurrently write new entries to the map. That's OK as the
+  // map is never resized.
+  while (true) {
+    const Entry& entry = map_[index];
+    if (entry.method == nullptr) {
+      // Not compiled yet.
+      return nullptr;
+    }
+    if (entry.method == method) {
+      if (entry.code_ptr == nullptr) {
+        // This is a race with the zygote which wrote the method, but hasn't written the
+        // code. Just bail and wait for the next time we need the method.
+        return nullptr;
+      }
+      if (pc != 0 && !OatQuickMethodHeader::FromCodePointer(entry.code_ptr)->Contains(pc)) {
+        return nullptr;
+      }
+      return entry.code_ptr;
+    }
+    index = (index + 1) & (map_.size() - 1);
+    DCHECK_NE(original_index, index);
+  }
+}
+
+void ZygoteMap::Put(const void* code, ArtMethod* method) {
+  if (map_.empty()) {
+    return;
+  }
+  CHECK(Runtime::Current()->IsZygote());
+  std::hash<ArtMethod*> hf;
+  size_t index = hf(method) & (map_.size() - 1);
+  size_t original_index = index;
+  // Because the size of the map is bigger than the number of methods that will
+  // be added, we are guaranteed to find a free slot in the array, and
+  // therefore for this loop to terminate.
+  while (true) {
+    Entry* entry = &map_[index];
+    if (entry->method == nullptr) {
+      // Note that readers can read this memory concurrently, but that's OK as
+      // we are writing pointers.
+      region_->WriteData(entry, Entry { method, code });
+      break;
+    }
+    index = (index + 1) & (map_.size() - 1);
+    DCHECK_NE(original_index, index);
+  }
+  DCHECK_EQ(GetCodeFor(method), code);
+}
+
 }  // namespace jit
 }  // namespace art
diff --git a/runtime/jit/jit_code_cache.h b/runtime/jit/jit_code_cache.h
index 88b440b..20b74d6 100644
--- a/runtime/jit/jit_code_cache.h
+++ b/runtime/jit/jit_code_cache.h
@@ -25,6 +25,7 @@
 #include <vector>
 
 #include "base/arena_containers.h"
+#include "base/array_ref.h"
 #include "base/atomic.h"
 #include "base/histogram.h"
 #include "base/macros.h"
@@ -78,6 +79,51 @@
 // of garbage collecting code.
 using CodeCacheBitmap = gc::accounting::MemoryRangeBitmap<kJitCodeAccountingBytes>;
 
+// Class abstraction over a map of ArtMethod -> compiled code, where the
+// ArtMethod are compiled by the zygote, and the map acts as a communication
+// channel between the zygote and the other processes.
+// For the zygote process, this map is the only map it is placing the compiled
+// code. JitCodeCache.method_code_map_ is empty.
+//
+// This map is writable only by the zygote, and readable by all children.
+class ZygoteMap {
+ public:
+  explicit ZygoteMap(JitMemoryRegion* region) : map_(), region_(region) {}
+
+  // Initialize the data structure so it can hold `number_of_methods` mappings.
+  // Note that the map is fixed size and never grows.
+  void Initialize(uint32_t number_of_methods) REQUIRES(!Locks::jit_lock_);
+
+  // Add the mapping method -> code.
+  void Put(const void* code, ArtMethod* method) REQUIRES(Locks::jit_lock_);
+
+  // Return the code pointer for the given method. If pc is not zero, check that
+  // the pc falls into that code range. Return null otherwise.
+  const void* GetCodeFor(ArtMethod* method, uintptr_t pc = 0) const;
+
+  // Return whether the map has associated code for the given method.
+  bool ContainsMethod(ArtMethod* method) const {
+    return GetCodeFor(method) != nullptr;
+  }
+
+ private:
+  struct Entry {
+    ArtMethod* method;
+    // Note we currently only allocate code in the low 4g, so we could just reserve 4 bytes
+    // for the code pointer. For simplicity and in the case we move to 64bit
+    // addresses for code, just keep it void* for now.
+    const void* code_ptr;
+  };
+
+  // The map allocated with `region_`.
+  ArrayRef<Entry> map_;
+
+  // The region in which the map is allocated.
+  JitMemoryRegion* const region_;
+
+  DISALLOW_COPY_AND_ASSIGN(ZygoteMap);
+};
+
 class JitCodeCache {
  public:
   static constexpr size_t kMaxCapacity = 64 * MB;
@@ -260,6 +306,9 @@
   bool GetGarbageCollectCodeUnsafe() const NO_THREAD_SAFETY_ANALYSIS {
     return garbage_collect_code_;
   }
+  ZygoteMap* GetZygoteMap() {
+    return &zygote_map_;
+  }
 
   // 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.
@@ -423,6 +472,10 @@
   // ProfilingInfo objects we have allocated.
   std::vector<ProfilingInfo*> profiling_infos_ GUARDED_BY(Locks::jit_lock_);
 
+  // Methods that the zygote has compiled and can be shared across processes
+  // forked from the zygote.
+  ZygoteMap zygote_map_;
+
   // -------------- JIT GC related data structures ----------------------- //
 
   // Condition to wait on during collection.
diff --git a/runtime/jit/jit_memory_region.h b/runtime/jit/jit_memory_region.h
index 0833380..d09c46a 100644
--- a/runtime/jit/jit_memory_region.h
+++ b/runtime/jit/jit_memory_region.h
@@ -111,6 +111,18 @@
     return exec_mspace_ != nullptr || data_mspace_ != nullptr;
   }
 
+  template <typename T>
+  void FillData(T* address, size_t n, const T& t)  REQUIRES(Locks::jit_lock_) {
+    std::fill_n(GetWritableDataAddress(address), n, t);
+  }
+
+  // Generic helper for writing abritrary data in the data portion of the
+  // region.
+  template <typename T>
+  void WriteData(T* address, const T& value) {
+    *GetWritableDataAddress(address) = value;
+  }
+
   bool HasDualCodeMapping() const {
     return non_exec_pages_.IsValid();
   }