Reland^3 "Revamp JIT GC."

This reverts commit b3d88b3f4b47b7635cbdab54dfdcbf866b9813ff.

Reason for revert: Do an atomic exchange when updating an entrypoint.

Change-Id: Idff00e87c2ef57d870975f7c63c06c672fd3279e

8 files changed, 287 insertions, 220 deletions

diff --git a/openjdkjvmti/ti_redefine.cc b/openjdkjvmti/ti_redefine.cc
index fc1828a56e..3d0c05734c 100644
--- a/openjdkjvmti/ti_redefine.cc
+++ b/openjdkjvmti/ti_redefine.cc
@@ -2550,6 +2550,9 @@ void Redefiner::ClassRedefinition::UpdateMethods(art::ObjPtr<art::mirror::Class>
   const art::DexFile& old_dex_file = mclass->GetDexFile();
   // Update methods.
   for (art::ArtMethod& method : mclass->GetDeclaredMethods(image_pointer_size)) {
+    // Reinitialize the method by calling `CopyFrom`. This ensures for example
+    // the entrypoint and the hotness are reset.
+    method.CopyFrom(&method, image_pointer_size);
     const art::dex::StringId* new_name_id = dex_file_->FindStringId(method.GetName());
     art::dex::TypeIndex method_return_idx =
         dex_file_->GetIndexForTypeId(*dex_file_->FindTypeId(method.GetReturnTypeDescriptorView()));
diff --git a/runtime/art_method.cc b/runtime/art_method.cc
index 40022c9d6d..11df3e6106 100644
--- a/runtime/art_method.cc
+++ b/runtime/art_method.cc
@@ -794,16 +794,18 @@ void ArtMethod::CopyFrom(ArtMethod* src, PointerSize image_pointer_size) {
   const void* entry_point = GetEntryPointFromQuickCompiledCodePtrSize(image_pointer_size);
   if (runtime->UseJitCompilation()) {
     if (runtime->GetJit()->GetCodeCache()->ContainsPc(entry_point)) {
-      SetEntryPointFromQuickCompiledCodePtrSize(
-          src->IsNative() ? GetQuickGenericJniStub() : GetQuickToInterpreterBridge(),
-          image_pointer_size);
+      SetNativePointer(EntryPointFromQuickCompiledCodeOffset(image_pointer_size),
+                       src->IsNative() ? GetQuickGenericJniStub() : GetQuickToInterpreterBridge(),
+                       image_pointer_size);
     }
   }
   ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
   if (interpreter::IsNterpSupported() && class_linker->IsNterpEntryPoint(entry_point)) {
     // If the entrypoint is nterp, it's too early to check if the new method
     // will support it. So for simplicity, use the interpreter bridge.
-    SetEntryPointFromQuickCompiledCodePtrSize(GetQuickToInterpreterBridge(), image_pointer_size);
+    SetNativePointer(EntryPointFromQuickCompiledCodeOffset(image_pointer_size),
+                     GetQuickToInterpreterBridge(),
+                     image_pointer_size);
   }
 
   // Clear the data pointer, it will be set if needed by the caller.
@@ -920,4 +922,33 @@ ALWAYS_INLINE static inline void DoGetAccessFlagsHelper(ArtMethod* method)
         method->GetDeclaringClass<kReadBarrierOption>()->IsErroneous());
 }
 
+template <typename T>
+const void* Exchange(uintptr_t ptr, uintptr_t new_value) {
+  std::atomic<T>* atomic_addr = reinterpret_cast<std::atomic<T>*>(ptr);
+  return reinterpret_cast<const void*>(
+      atomic_addr->exchange(dchecked_integral_cast<T>(new_value), std::memory_order_relaxed));
+}
+
+void ArtMethod::SetEntryPointFromQuickCompiledCodePtrSize(
+    const void* entry_point_from_quick_compiled_code, PointerSize pointer_size) {
+  const void* current_entry_point = GetEntryPointFromQuickCompiledCodePtrSize(pointer_size);
+  if (current_entry_point == entry_point_from_quick_compiled_code) {
+    return;
+  }
+
+  // Do an atomic exchange to avoid potentially unregistering JIT code twice.
+  MemberOffset offset = EntryPointFromQuickCompiledCodeOffset(pointer_size);
+  uintptr_t new_value = reinterpret_cast<uintptr_t>(entry_point_from_quick_compiled_code);
+  uintptr_t ptr = reinterpret_cast<uintptr_t>(this) + offset.Uint32Value();
+  const void* old_value = (pointer_size == PointerSize::k32)
+      ? Exchange<uint32_t>(ptr, new_value)
+      : Exchange<uint64_t>(ptr, new_value);
+
+  jit::Jit* jit = Runtime::Current()->GetJit();
+  if (jit != nullptr &&
+      jit->GetCodeCache()->ContainsPc(old_value)) {
+    jit->GetCodeCache()->AddZombieCode(this, old_value);
+  }
+}
+
 }  // namespace art
diff --git a/runtime/art_method.h b/runtime/art_method.h
index 96a204055b..bf1b4631c6 100644
--- a/runtime/art_method.h
+++ b/runtime/art_method.h
@@ -766,11 +766,7 @@ class EXPORT ArtMethod final {
   }
   ALWAYS_INLINE void SetEntryPointFromQuickCompiledCodePtrSize(
       const void* entry_point_from_quick_compiled_code, PointerSize pointer_size)
-      REQUIRES_SHARED(Locks::mutator_lock_) {
-    SetNativePointer(EntryPointFromQuickCompiledCodeOffset(pointer_size),
-                     entry_point_from_quick_compiled_code,
-                     pointer_size);
-  }
+      REQUIRES_SHARED(Locks::mutator_lock_);
 
   static constexpr MemberOffset DataOffset(PointerSize pointer_size) {
     return MemberOffset(PtrSizedFieldsOffset(pointer_size) + OFFSETOF_MEMBER(
@@ -1192,6 +1188,8 @@ class EXPORT ArtMethod final {
   // Used by GetName and GetNameView to share common code.
   const char* GetRuntimeMethodName() REQUIRES_SHARED(Locks::mutator_lock_);
 
+  friend class RuntimeImageHelper;  // For SetNativePointer.
+
   DISALLOW_COPY_AND_ASSIGN(ArtMethod);  // Need to use CopyFrom to deal with 32 vs 64 bits.
 };
 
diff --git a/runtime/jit/jit_code_cache.cc b/runtime/jit/jit_code_cache.cc
index eae56a45b3..e6db9817df 100644
--- a/runtime/jit/jit_code_cache.cc
+++ b/runtime/jit/jit_code_cache.cc
@@ -342,6 +342,8 @@ const void* JitCodeCache::GetSavedEntryPointOfPreCompiledMethod(ArtMethod* metho
       auto it = saved_compiled_methods_map_.find(method);
       if (it != saved_compiled_methods_map_.end()) {
         code_ptr = it->second;
+        // Now that we're using the saved entrypoint, remove it from the saved map.
+        saved_compiled_methods_map_.erase(it);
       }
     }
     if (code_ptr != nullptr) {
@@ -488,33 +490,36 @@ void JitCodeCache::FreeAllMethodHeaders(
         ->RemoveDependentsWithMethodHeaders(method_headers);
   }
 
-  ScopedCodeCacheWrite scc(private_region_);
-  for (const OatQuickMethodHeader* method_header : method_headers) {
-    FreeCodeAndData(method_header->GetCode());
-  }
+  {
+    ScopedCodeCacheWrite scc(private_region_);
+    for (const OatQuickMethodHeader* method_header : method_headers) {
+      FreeCodeAndData(method_header->GetCode());
+    }
 
-  // We have potentially removed a lot of debug info. Do maintenance pass to save space.
-  RepackNativeDebugInfoForJit();
+    // We have potentially removed a lot of debug info. Do maintenance pass to save space.
+    RepackNativeDebugInfoForJit();
+  }
 
   // Check that the set of compiled methods exactly matches native debug information.
   // Does not check zygote methods since they can change concurrently.
   if (kIsDebugBuild && !Runtime::Current()->IsZygote()) {
     std::map<const void*, ArtMethod*> compiled_methods;
+    std::set<const void*> debug_info;
     VisitAllMethods([&](const void* addr, ArtMethod* method) {
       if (!IsInZygoteExecSpace(addr)) {
         CHECK(addr != nullptr && method != nullptr);
         compiled_methods.emplace(addr, method);
       }
     });
-    std::set<const void*> debug_info;
     ForEachNativeDebugSymbol([&](const void* addr, size_t, const char* name) {
       addr = AlignDown(addr, GetInstructionSetInstructionAlignment(kRuntimeISA));  // Thumb-bit.
-      CHECK(debug_info.emplace(addr).second) << "Duplicate debug info: " << addr << " " << name;
+      bool res = debug_info.emplace(addr).second;
+      CHECK(res) << "Duplicate debug info: " << addr << " " << name;
       CHECK_EQ(compiled_methods.count(addr), 1u) << "Extra debug info: " << addr << " " << name;
     });
     if (!debug_info.empty()) {  // If debug-info generation is enabled.
       for (const auto& [addr, method] : compiled_methods) {
-        CHECK_EQ(debug_info.count(addr), 1u) << "No debug info: " << method->PrettyMethod();
+        CHECK_EQ(debug_info.count(addr), 1u) << "Mising debug info";
       }
       CHECK_EQ(compiled_methods.size(), debug_info.size());
     }
@@ -529,52 +534,67 @@ void JitCodeCache::RemoveMethodsIn(Thread* self, const LinearAlloc& alloc) {
   // for entries in this set. And it's more efficient to iterate through
   // the CHA dependency map just once with an unordered_set.
   std::unordered_set<OatQuickMethodHeader*> method_headers;
+  MutexLock mu(self, *Locks::jit_lock_);
+  // We do not check if a code cache GC is in progress, as this method comes
+  // with the classlinker_classes_lock_ held, and suspending ourselves could
+  // lead to a deadlock.
   {
-    MutexLock mu(self, *Locks::jit_lock_);
-    // We do not check if a code cache GC is in progress, as this method comes
-    // with the classlinker_classes_lock_ held, and suspending ourselves could
-    // lead to a deadlock.
-    {
-      for (auto it = jni_stubs_map_.begin(); it != jni_stubs_map_.end();) {
-        it->second.RemoveMethodsIn(alloc);
-        if (it->second.GetMethods().empty()) {
-          method_headers.insert(OatQuickMethodHeader::FromCodePointer(it->second.GetCode()));
-          it = jni_stubs_map_.erase(it);
-        } else {
-          it->first.UpdateShorty(it->second.GetMethods().front());
-          ++it;
-        }
+    for (auto it = jni_stubs_map_.begin(); it != jni_stubs_map_.end();) {
+      it->second.RemoveMethodsIn(alloc);
+      if (it->second.GetMethods().empty()) {
+        method_headers.insert(OatQuickMethodHeader::FromCodePointer(it->second.GetCode()));
+        it = jni_stubs_map_.erase(it);
+      } else {
+        it->first.UpdateShorty(it->second.GetMethods().front());
+        ++it;
       }
-      for (auto it = method_code_map_.begin(); it != method_code_map_.end();) {
-        if (alloc.ContainsUnsafe(it->second)) {
-          method_headers.insert(OatQuickMethodHeader::FromCodePointer(it->first));
-          VLOG(jit) << "JIT removed " << it->second->PrettyMethod() << ": " << it->first;
-          it = method_code_map_.erase(it);
-        } else {
-          ++it;
-        }
+    }
+    for (auto it = zombie_jni_code_.begin(); it != zombie_jni_code_.end();) {
+      if (alloc.ContainsUnsafe(*it)) {
+        it = zombie_jni_code_.erase(it);
+      } else {
+        ++it;
       }
     }
-    for (auto it = osr_code_map_.begin(); it != osr_code_map_.end();) {
-      if (alloc.ContainsUnsafe(it->first)) {
-        // Note that the code has already been pushed to method_headers in the loop
-        // above and is going to be removed in FreeCode() below.
-        it = osr_code_map_.erase(it);
+    for (auto it = processed_zombie_jni_code_.begin(); it != processed_zombie_jni_code_.end();) {
+      if (alloc.ContainsUnsafe(*it)) {
+        it = processed_zombie_jni_code_.erase(it);
       } else {
         ++it;
       }
     }
-    for (auto it = profiling_infos_.begin(); it != profiling_infos_.end();) {
-      ProfilingInfo* info = it->second;
-      if (alloc.ContainsUnsafe(info->GetMethod())) {
-        private_region_.FreeWritableData(reinterpret_cast<uint8_t*>(info));
-        it = profiling_infos_.erase(it);
+    for (auto it = method_code_map_.begin(); it != method_code_map_.end();) {
+      if (alloc.ContainsUnsafe(it->second)) {
+        method_headers.insert(OatQuickMethodHeader::FromCodePointer(it->first));
+        VLOG(jit) << "JIT removed " << it->second->PrettyMethod() << ": " << it->first;
+        it = method_code_map_.erase(it);
+        zombie_code_.erase(it->first);
+        processed_zombie_code_.erase(it->first);
       } else {
         ++it;
       }
     }
-    FreeAllMethodHeaders(method_headers);
   }
+  for (auto it = osr_code_map_.begin(); it != osr_code_map_.end();) {
+    DCHECK(!ContainsElement(zombie_code_, it->second));
+    if (alloc.ContainsUnsafe(it->first)) {
+      // Note that the code has already been pushed to method_headers in the loop
+      // above and is going to be removed in FreeCode() below.
+      it = osr_code_map_.erase(it);
+    } else {
+      ++it;
+    }
+  }
+  for (auto it = profiling_infos_.begin(); it != profiling_infos_.end();) {
+    ProfilingInfo* info = it->second;
+    if (alloc.ContainsUnsafe(info->GetMethod())) {
+      private_region_.FreeWritableData(reinterpret_cast<uint8_t*>(info));
+      it = profiling_infos_.erase(it);
+    } else {
+      ++it;
+    }
+  }
+  FreeAllMethodHeaders(method_headers);
 }
 
 bool JitCodeCache::IsWeakAccessEnabled(Thread* self) const {
@@ -640,18 +660,6 @@ static void ClearMethodCounter(ArtMethod* method, bool was_warm)
   method->UpdateCounter(/* new_samples= */ 1);
 }
 
-void JitCodeCache::WaitForPotentialCollectionToCompleteRunnable(Thread* self) {
-  while (collection_in_progress_) {
-    Locks::jit_lock_->Unlock(self);
-    {
-      ScopedThreadSuspension sts(self, ThreadState::kSuspended);
-      MutexLock mu(self, *Locks::jit_lock_);
-      WaitForPotentialCollectionToComplete(self);
-    }
-    Locks::jit_lock_->Lock(self);
-  }
-}
-
 bool JitCodeCache::Commit(Thread* self,
                           JitMemoryRegion* region,
                           ArtMethod* method,
@@ -679,9 +687,6 @@ bool JitCodeCache::Commit(Thread* self,
   OatQuickMethodHeader* method_header = nullptr;
   {
     MutexLock mu(self, *Locks::jit_lock_);
-    // We need to make sure that there will be no jit-gcs going on and wait for any ongoing one to
-    // finish.
-    WaitForPotentialCollectionToCompleteRunnable(self);
     const uint8_t* code_ptr = region->CommitCode(reserved_code, code, stack_map_data);
     if (code_ptr == nullptr) {
       return false;
@@ -781,11 +786,6 @@ bool JitCodeCache::Commit(Thread* self,
             method, method_header->GetEntryPoint());
       }
     }
-    if (collection_in_progress_) {
-      // We need to update the live bitmap if there is a GC to ensure it sees this new
-      // code.
-      GetLiveBitmap()->AtomicTestAndSet(FromCodeToAllocation(code_ptr));
-    }
     VLOG(jit)
         << "JIT added (kind=" << compilation_kind << ") "
         << ArtMethod::PrettyMethod(method) << "@" << method
@@ -854,6 +854,7 @@ bool JitCodeCache::RemoveMethodLocked(ArtMethod* method, bool release_memory) {
           FreeCodeAndData(it->second.GetCode());
         }
         jni_stubs_map_.erase(it);
+        zombie_jni_code_.erase(method);
       } else {
         it->first.UpdateShorty(it->second.GetMethods().front());
       }
@@ -984,7 +985,6 @@ bool JitCodeCache::Reserve(Thread* self,
     {
       ScopedThreadSuspension sts(self, ThreadState::kSuspended);
       MutexLock mu(self, *Locks::jit_lock_);
-      WaitForPotentialCollectionToComplete(self);
       ScopedCodeCacheWrite ccw(*region);
       code = region->AllocateCode(code_size);
       data = region->AllocateData(data_size);
@@ -1001,8 +1001,8 @@ bool JitCodeCache::Reserve(Thread* self,
                 << PrettySize(data_size);
       return false;
     }
-    // Run a code cache collection and try again.
-    GarbageCollectCache(self);
+    // Increase the capacity and try again.
+    IncreaseCodeCacheCapacity(self);
   }
 
   *reserved_code = ArrayRef<const uint8_t>(code, code_size);
@@ -1080,11 +1080,6 @@ class MarkCodeClosure final : public Closure {
   Barrier* const barrier_;
 };
 
-void JitCodeCache::NotifyCollectionDone(Thread* self) {
-  collection_in_progress_ = false;
-  lock_cond_.Broadcast(self);
-}
-
 void JitCodeCache::MarkCompiledCodeOnThreadStacks(Thread* self) {
   Barrier barrier(0);
   size_t threads_running_checkpoint = 0;
@@ -1102,86 +1097,116 @@ bool JitCodeCache::IsAtMaxCapacity() const {
   return private_region_.GetCurrentCapacity() == private_region_.GetMaxCapacity();
 }
 
-void JitCodeCache::GarbageCollectCache(Thread* self) {
+void JitCodeCache::IncreaseCodeCacheCapacity(Thread* self) {
+  ScopedThreadSuspension sts(self, ThreadState::kSuspended);
+  MutexLock mu(self, *Locks::jit_lock_);
+  // Wait for a potential collection, as the size of the bitmap used by that collection
+  // is of the current capacity.
+  WaitForPotentialCollectionToComplete(self);
+  private_region_.IncreaseCodeCacheCapacity();
+}
+
+void JitCodeCache::RemoveUnmarkedCode(Thread* self) {
   ScopedTrace trace(__FUNCTION__);
-  // Wait for an existing collection, or let everyone know we are starting one.
-  {
-    ScopedThreadSuspension sts(self, ThreadState::kSuspended);
-    MutexLock mu(self, *Locks::jit_lock_);
-    if (!garbage_collect_code_) {
-      private_region_.IncreaseCodeCacheCapacity();
-      return;
-    } else if (WaitForPotentialCollectionToComplete(self)) {
-      return;
+  std::unordered_set<OatQuickMethodHeader*> method_headers;
+  ScopedDebugDisallowReadBarriers sddrb(self);
+  MutexLock mu(self, *Locks::jit_lock_);
+  // Iterate over all zombie code and remove entries that are not marked.
+  for (auto it = processed_zombie_code_.begin(); it != processed_zombie_code_.end();) {
+    const void* code_ptr = *it;
+    uintptr_t allocation = FromCodeToAllocation(code_ptr);
+    DCHECK(!IsInZygoteExecSpace(code_ptr));
+    if (GetLiveBitmap()->Test(allocation)) {
+      ++it;
     } else {
-      number_of_collections_++;
-      live_bitmap_.reset(CodeCacheBitmap::Create(
-          "code-cache-bitmap",
-          reinterpret_cast<uintptr_t>(private_region_.GetExecPages()->Begin()),
-          reinterpret_cast<uintptr_t>(
-              private_region_.GetExecPages()->Begin() + private_region_.GetCurrentCapacity() / 2)));
-      collection_in_progress_ = true;
+      OatQuickMethodHeader* header = OatQuickMethodHeader::FromCodePointer(code_ptr);
+      method_headers.insert(header);
+      method_code_map_.erase(header->GetCode());
+      VLOG(jit) << "JIT removed " << *it;
+      it = processed_zombie_code_.erase(it);
+    }
+  }
+  for (auto it = processed_zombie_jni_code_.begin(); it != processed_zombie_jni_code_.end();) {
+    ArtMethod* method = *it;
+    auto stub = jni_stubs_map_.find(JniStubKey(method));
+    if (stub == jni_stubs_map_.end()) {
+      it = processed_zombie_jni_code_.erase(it);
+      continue;
+    }
+    JniStubData& data = stub->second;
+    if (!data.IsCompiled() || !ContainsElement(data.GetMethods(), method)) {
+      it = processed_zombie_jni_code_.erase(it);
+    } else if (method->GetEntryPointFromQuickCompiledCode() ==
+            OatQuickMethodHeader::FromCodePointer(data.GetCode())->GetEntryPoint()) {
+      // The stub got reused for this method, remove ourselves from the zombie
+      // list.
+      it = processed_zombie_jni_code_.erase(it);
+    } else if (!GetLiveBitmap()->Test(FromCodeToAllocation(data.GetCode()))) {
+      data.RemoveMethod(method);
+      if (data.GetMethods().empty()) {
+        OatQuickMethodHeader* header = OatQuickMethodHeader::FromCodePointer(data.GetCode());
+        method_headers.insert(header);
+        CHECK(ContainsPc(header));
+        VLOG(jit) << "JIT removed native code of" << method->PrettyMethod();
+        jni_stubs_map_.erase(stub);
+      } else {
+        stub->first.UpdateShorty(stub->second.GetMethods().front());
+      }
+      it = processed_zombie_jni_code_.erase(it);
+    } else {
+      ++it;
     }
   }
+  FreeAllMethodHeaders(method_headers);
+}
 
-  TimingLogger logger("JIT code cache timing logger", true, VLOG_IS_ON(jit));
-  {
-    TimingLogger::ScopedTiming st("Code cache collection", &logger);
+void JitCodeCache::AddZombieCode(ArtMethod* method, const void* entry_point) {
+  CHECK(ContainsPc(entry_point));
+  CHECK(method->IsNative() || (method->GetEntryPointFromQuickCompiledCode() != entry_point));
+  const void* code_ptr = OatQuickMethodHeader::FromEntryPoint(entry_point)->GetCode();
+  if (!IsInZygoteExecSpace(code_ptr)) {
+    Thread* self = Thread::Current();
+    if (Locks::jit_lock_->IsExclusiveHeld(self)) {
+      AddZombieCodeInternal(method, code_ptr);
+    } else {
+      MutexLock mu(Thread::Current(), *Locks::jit_lock_);
+      AddZombieCodeInternal(method, code_ptr);
+    }
+  }
+}
 
-    VLOG(jit) << "Do code cache collection, code="
-              << PrettySize(CodeCacheSize())
-              << ", data=" << PrettySize(DataCacheSize());
 
-    DoCollection(self);
+class JitGcTask final : public Task {
+ public:
+  JitGcTask() {}
 
-    VLOG(jit) << "After code cache collection, code="
-              << PrettySize(CodeCacheSize())
-              << ", data=" << PrettySize(DataCacheSize());
+  void Run(Thread* self) override {
+    Runtime::Current()->GetJit()->GetCodeCache()->DoCollection(self);
+  }
 
-    {
-      MutexLock mu(self, *Locks::jit_lock_);
-      private_region_.IncreaseCodeCacheCapacity();
-      live_bitmap_.reset(nullptr);
-      NotifyCollectionDone(self);
-    }
+  void Finalize() override {
+    delete this;
   }
-  Runtime::Current()->GetJit()->AddTimingLogger(logger);
-}
+};
 
-void JitCodeCache::RemoveUnmarkedCode(Thread* self) {
-  ScopedTrace trace(__FUNCTION__);
-  ScopedDebugDisallowReadBarriers sddrb(self);
-  std::unordered_set<OatQuickMethodHeader*> method_headers;
-  {
-    MutexLock mu(self, *Locks::jit_lock_);
-    // Iterate over all compiled code and remove entries that are not marked.
-    for (auto it = jni_stubs_map_.begin(); it != jni_stubs_map_.end();) {
-      JniStubData* data = &it->second;
-      if (IsInZygoteExecSpace(data->GetCode()) ||
-          !data->IsCompiled() ||
-          GetLiveBitmap()->Test(FromCodeToAllocation(data->GetCode()))) {
-        ++it;
-      } else {
-        method_headers.insert(OatQuickMethodHeader::FromCodePointer(data->GetCode()));
-        for (ArtMethod* method : data->GetMethods()) {
-          VLOG(jit) << "JIT removed (JNI) " << method->PrettyMethod() << ": " << data->GetCode();
-        }
-        it = jni_stubs_map_.erase(it);
-      }
-    }
-    for (auto it = method_code_map_.begin(); it != method_code_map_.end();) {
-      const void* code_ptr = it->first;
-      uintptr_t allocation = FromCodeToAllocation(code_ptr);
-      if (IsInZygoteExecSpace(code_ptr) || GetLiveBitmap()->Test(allocation)) {
-        ++it;
-      } else {
-        OatQuickMethodHeader* header = OatQuickMethodHeader::FromCodePointer(code_ptr);
-        method_headers.insert(header);
-        VLOG(jit) << "JIT removed " << it->second->PrettyMethod() << ": " << it->first;
-        it = method_code_map_.erase(it);
-      }
+void JitCodeCache::AddZombieCodeInternal(ArtMethod* method, const void* code_ptr) {
+  if (method->IsNative()) {
+    CHECK(jni_stubs_map_.find(JniStubKey(method)) != jni_stubs_map_.end());
+    zombie_jni_code_.insert(method);
+  } else {
+    CHECK(!ContainsElement(zombie_code_, code_ptr));
+    zombie_code_.insert(code_ptr);
+  }
+  // Arbitrary threshold of number of zombie code before doing a GC.
+  static constexpr size_t kNumberOfZombieCodeThreshold = kIsDebugBuild ? 1 : 1000;
+  size_t number_of_code_to_delete =
+      zombie_code_.size() + zombie_jni_code_.size() + osr_code_map_.size();
+  if (number_of_code_to_delete >= kNumberOfZombieCodeThreshold) {
+    JitThreadPool* pool = Runtime::Current()->GetJit()->GetThreadPool();
+    if (pool != nullptr && !gc_task_scheduled_) {
+      gc_task_scheduled_ = true;
+      pool->AddTask(Thread::Current(), new JitGcTask());
     }
-    FreeAllMethodHeaders(method_headers);
   }
 }
 
@@ -1233,51 +1258,58 @@ void JitCodeCache::ResetHotnessCounter(ArtMethod* method, Thread* self) {
 
 void JitCodeCache::DoCollection(Thread* self) {
   ScopedTrace trace(__FUNCTION__);
+
   {
     ScopedDebugDisallowReadBarriers sddrb(self);
     MutexLock mu(self, *Locks::jit_lock_);
-    // Mark compiled code that are entrypoints of ArtMethods. Compiled code that is not
-    // an entry point is either:
-    // - an osr compiled code, that will be removed if not in a thread call stack.
-    // - discarded compiled code, that will be removed if not in a thread call stack.
-    for (const auto& entry : jni_stubs_map_) {
-      const JniStubData& data = entry.second;
-      const void* code_ptr = data.GetCode();
-      if (IsInZygoteExecSpace(code_ptr)) {
-        continue;
-      }
-      const OatQuickMethodHeader* method_header = OatQuickMethodHeader::FromCodePointer(code_ptr);
-      for (ArtMethod* method : data.GetMethods()) {
-        if (method_header->GetEntryPoint() == method->GetEntryPointFromQuickCompiledCode()) {
-          GetLiveBitmap()->AtomicTestAndSet(FromCodeToAllocation(code_ptr));
-          break;
-        }
-      }
-    }
-    for (const auto& it : method_code_map_) {
-      ArtMethod* method = it.second;
-      const void* code_ptr = it.first;
-      if (IsInZygoteExecSpace(code_ptr)) {
-        continue;
-      }
-      const OatQuickMethodHeader* method_header = OatQuickMethodHeader::FromCodePointer(code_ptr);
-      if (method_header->GetEntryPoint() == method->GetEntryPointFromQuickCompiledCode()) {
-        GetLiveBitmap()->AtomicTestAndSet(FromCodeToAllocation(code_ptr));
-      }
+    if (!garbage_collect_code_) {
+      return;
+    } else if (WaitForPotentialCollectionToComplete(self)) {
+      return;
     }
-
+    collection_in_progress_ = true;
+    number_of_collections_++;
+    live_bitmap_.reset(CodeCacheBitmap::Create(
+          "code-cache-bitmap",
+          reinterpret_cast<uintptr_t>(private_region_.GetExecPages()->Begin()),
+          reinterpret_cast<uintptr_t>(
+              private_region_.GetExecPages()->Begin() + private_region_.GetCurrentCapacity() / 2)));
+    processed_zombie_code_.insert(zombie_code_.begin(), zombie_code_.end());
+    zombie_code_.clear();
+    processed_zombie_jni_code_.insert(zombie_jni_code_.begin(), zombie_jni_code_.end());
+    zombie_jni_code_.clear();
     // Empty osr method map, as osr compiled code will be deleted (except the ones
     // on thread stacks).
+    for (auto it = osr_code_map_.begin(); it != osr_code_map_.end(); ++it) {
+      processed_zombie_code_.insert(it->second);
+    }
     osr_code_map_.clear();
   }
+  TimingLogger logger("JIT code cache timing logger", true, VLOG_IS_ON(jit));
+  {
+    TimingLogger::ScopedTiming st("Code cache collection", &logger);
+
+    {
+      ScopedObjectAccess soa(self);
+      // Run a checkpoint on all threads to mark the JIT compiled code they are running.
+      MarkCompiledCodeOnThreadStacks(self);
+
+      // Remove zombie code which hasn't been marked.
+      RemoveUnmarkedCode(self);
+    }
+
+    MutexLock mu(self, *Locks::jit_lock_);
+    live_bitmap_.reset(nullptr);
+    NotifyCollectionDone(self);
+  }
 
-  // Run a checkpoint on all threads to mark the JIT compiled code they are running.
-  MarkCompiledCodeOnThreadStacks(self);
+  Runtime::Current()->GetJit()->AddTimingLogger(logger);
+}
 
-  // At this point, mutator threads are still running, and entrypoints of methods can
-  // change. We do know they cannot change to a code cache entry that is not marked,
-  // therefore we can safely remove those entries.
-  RemoveUnmarkedCode(self);
+void JitCodeCache::NotifyCollectionDone(Thread* self) {
+  collection_in_progress_ = false;
+  gc_task_scheduled_ = false;
+  lock_cond_.Broadcast(self);
 }
 
 OatQuickMethodHeader* JitCodeCache::LookupMethodHeader(uintptr_t pc, ArtMethod* method) {
@@ -1381,7 +1413,7 @@ ProfilingInfo* JitCodeCache::AddProfilingInfo(Thread* self,
   }
 
   if (info == nullptr) {
-    GarbageCollectCache(self);
+    IncreaseCodeCacheCapacity(self);
     MutexLock mu(self, *Locks::jit_lock_);
     info = AddProfilingInfoInternal(self, method, inline_cache_entries, branch_cache_entries);
   }
@@ -1609,11 +1641,6 @@ bool JitCodeCache::NotifyCompilationOf(ArtMethod* method,
       // as well change them back as this stub shall not be collected anyway and this
       // can avoid a few expensive GenericJNI calls.
       data->UpdateEntryPoints(entrypoint);
-      if (collection_in_progress_) {
-        if (!IsInZygoteExecSpace(data->GetCode())) {
-          GetLiveBitmap()->AtomicTestAndSet(FromCodeToAllocation(data->GetCode()));
-        }
-      }
     }
     return new_compilation;
   } else {
diff --git a/runtime/jit/jit_code_cache.h b/runtime/jit/jit_code_cache.h
index 96fc7e2706..7de29d4024 100644
--- a/runtime/jit/jit_code_cache.h
+++ b/runtime/jit/jit_code_cache.h
@@ -285,11 +285,9 @@ class JitCodeCache {
       REQUIRES_SHARED(Locks::mutator_lock_)
       REQUIRES(!Locks::jit_lock_);
   void FreeLocked(JitMemoryRegion* region, const uint8_t* code, const uint8_t* data)
-      REQUIRES_SHARED(Locks::mutator_lock_)
       REQUIRES(Locks::jit_lock_);
 
-  // Perform a collection on the code cache.
-  EXPORT void GarbageCollectCache(Thread* self)
+  void IncreaseCodeCacheCapacity(Thread* self)
       REQUIRES(!Locks::jit_lock_)
       REQUIRES_SHARED(Locks::mutator_lock_);
 
@@ -423,9 +421,20 @@ class JitCodeCache {
                               Thread* self)
       REQUIRES_SHARED(Locks::mutator_lock_);
 
+  // NO_THREAD_SAFETY_ANALYSIS because we may be called with the JIT lock held
+  // or not. The implementation of this method handles the two cases.
+  void AddZombieCode(ArtMethod* method, const void* code_ptr) NO_THREAD_SAFETY_ANALYSIS;
+
+  EXPORT void DoCollection(Thread* self)
+      REQUIRES(!Locks::jit_lock_);
+
  private:
   JitCodeCache();
 
+  void AddZombieCodeInternal(ArtMethod* method, const void* code_ptr)
+      REQUIRES(Locks::jit_lock_)
+      REQUIRES_SHARED(Locks::mutator_lock_);
+
   ProfilingInfo* AddProfilingInfoInternal(Thread* self,
                                           ArtMethod* method,
                                           const std::vector<uint32_t>& inline_cache_entries,
@@ -433,20 +442,16 @@ class JitCodeCache {
       REQUIRES(Locks::jit_lock_)
       REQUIRES_SHARED(Locks::mutator_lock_);
 
-  // If a collection is in progress, wait for it to finish. Must be called with the mutator lock.
-  // The non-mutator lock version should be used if possible. This method will release then
-  // re-acquire the mutator lock.
-  void WaitForPotentialCollectionToCompleteRunnable(Thread* self)
-      REQUIRES(Locks::jit_lock_, !Roles::uninterruptible_) REQUIRES_SHARED(Locks::mutator_lock_);
-
   // If a collection is in progress, wait for it to finish. Return
   // whether the thread actually waited.
   bool WaitForPotentialCollectionToComplete(Thread* self)
-      REQUIRES(Locks::jit_lock_) REQUIRES(!Locks::mutator_lock_);
+      REQUIRES(Locks::jit_lock_) REQUIRES_SHARED(!Locks::mutator_lock_);
+
+  // Notify all waiting threads that a collection is done.
+  void NotifyCollectionDone(Thread* self) REQUIRES(Locks::jit_lock_);
 
   // Remove CHA dependents and underlying allocations for entries in `method_headers`.
   void FreeAllMethodHeaders(const std::unordered_set<OatQuickMethodHeader*>& method_headers)
-      REQUIRES_SHARED(Locks::mutator_lock_)
       REQUIRES(Locks::jit_lock_)
       REQUIRES(!Locks::cha_lock_);
 
@@ -462,8 +467,7 @@ class JitCodeCache {
 
   // Free code and data allocations for `code_ptr`.
   void FreeCodeAndData(const void* code_ptr)
-      REQUIRES(Locks::jit_lock_)
-      REQUIRES_SHARED(Locks::mutator_lock_);
+      REQUIRES(Locks::jit_lock_);
 
   // Number of bytes allocated in the code cache.
   size_t CodeCacheSize() REQUIRES(!Locks::jit_lock_);
@@ -477,16 +481,9 @@ class JitCodeCache {
   // Number of bytes allocated in the data cache.
   size_t DataCacheSizeLocked() REQUIRES(Locks::jit_lock_);
 
-  // Notify all waiting threads that a collection is done.
-  void NotifyCollectionDone(Thread* self) REQUIRES(Locks::jit_lock_);
-
   // Return whether the code cache's capacity is at its maximum.
   bool IsAtMaxCapacity() const REQUIRES(Locks::jit_lock_);
 
-  void DoCollection(Thread* self)
-      REQUIRES(!Locks::jit_lock_)
-      REQUIRES_SHARED(Locks::mutator_lock_);
-
   void RemoveUnmarkedCode(Thread* self)
       REQUIRES(!Locks::jit_lock_)
       REQUIRES_SHARED(Locks::mutator_lock_);
@@ -563,18 +560,28 @@ class JitCodeCache {
   // ProfilingInfo objects we have allocated.
   SafeMap<ArtMethod*, ProfilingInfo*> profiling_infos_ GUARDED_BY(Locks::jit_lock_);
 
+  // Zombie code and JNI methods to consider for collection.
+  std::set<const void*> zombie_code_ GUARDED_BY(Locks::jit_lock_);
+  std::set<ArtMethod*> zombie_jni_code_ GUARDED_BY(Locks::jit_lock_);
+
+  std::set<const void*> processed_zombie_code_ GUARDED_BY(Locks::jit_lock_);
+  std::set<ArtMethod*> processed_zombie_jni_code_ 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.
+  // Condition to wait on during collection and for accessing weak references in inline caches.
   ConditionVariable lock_cond_ GUARDED_BY(Locks::jit_lock_);
 
   // Whether there is a code cache collection in progress.
   bool collection_in_progress_ GUARDED_BY(Locks::jit_lock_);
 
+  // Whether a GC task is already scheduled.
+  bool gc_task_scheduled_ GUARDED_BY(Locks::jit_lock_);
+
   // Bitmap for collecting code and data.
   std::unique_ptr<CodeCacheBitmap> live_bitmap_;
 
diff --git a/runtime/runtime_image.cc b/runtime/runtime_image.cc
index 997ea2fde6..b134c79286 100644
--- a/runtime/runtime_image.cc
+++ b/runtime/runtime_image.cc
@@ -964,7 +964,9 @@ class RuntimeImageHelper {
           entrypoint = boot_jni_stub;
         }
       }
-      copy->SetEntryPointFromQuickCompiledCode(entrypoint);
+      copy->SetNativePointer(ArtMethod::EntryPointFromQuickCompiledCodeOffset(kRuntimePointerSize),
+                             entrypoint,
+                             kRuntimePointerSize);
 
       if (method->IsNative()) {
         StubType stub_type = method->IsCriticalNative()
diff --git a/test/667-jit-jni-stub/jit_jni_stub_test.cc b/test/667-jit-jni-stub/jit_jni_stub_test.cc
index 2b8e14c03d..52f467e38b 100644
--- a/test/667-jit-jni-stub/jit_jni_stub_test.cc
+++ b/test/667-jit-jni-stub/jit_jni_stub_test.cc
@@ -39,9 +39,15 @@ extern "C" JNIEXPORT
 void Java_Main_jitGc(JNIEnv*, jclass) {
   CHECK(Runtime::Current()->GetJit() != nullptr);
   jit::JitCodeCache* cache = Runtime::Current()->GetJit()->GetCodeCache();
-  ScopedObjectAccess soa(Thread::Current());
-  cache->InvalidateAllCompiledCode();
-  cache->GarbageCollectCache(Thread::Current());
+  Thread* self = Thread::Current();
+  {
+    ScopedObjectAccess soa(self);
+    cache->InvalidateAllCompiledCode();
+  }
+  cache->DoCollection(self);
+  // Run a second time in case the first run was a no-op due to a concurrent JIT
+  // GC from the JIT thread.
+  cache->DoCollection(self);
 }
 
 extern "C" JNIEXPORT
diff --git a/test/667-jit-jni-stub/src/Main.java b/test/667-jit-jni-stub/src/Main.java
index c0829b5100..05039670d1 100644
--- a/test/667-jit-jni-stub/src/Main.java
+++ b/test/667-jit-jni-stub/src/Main.java
@@ -67,9 +67,6 @@ public class Main {
     assertFalse(hasJitCompiledEntrypoint(Main.class, "callThrough"));
     assertFalse(hasJitCompiledCode(Main.class, "callThrough"));
     callThrough(Main.class, "testMixedFramesOnStackStage2");
-    // We have just returned through the JIT-compiled JNI stub, so it must still
-    // be compiled (though not necessarily with the entrypoint pointing to it).
-    assertTrue(hasJitCompiledCode(Main.class, "callThrough"));
     // Though the callThrough() is on the stack, that frame is using the GenericJNI
     // and does not prevent the collection of the JNI stub.
     testStubCanBeCollected();
@@ -94,10 +91,6 @@ public class Main {
   }
 
   public static void testStubCanBeCollected() {
-    assertTrue(hasJitCompiledCode(Main.class, "callThrough"));
-    doJitGcsUntilFullJitGcIsScheduled();
-    assertFalse(hasJitCompiledEntrypoint(Main.class, "callThrough"));
-    assertTrue(hasJitCompiledCode(Main.class, "callThrough"));
     jitGc();  // JIT GC without callThrough() on the stack should collect the callThrough() stub.
     assertFalse(hasJitCompiledEntrypoint(Main.class, "callThrough"));
     assertFalse(hasJitCompiledCode(Main.class, "callThrough"));