Call JitCodeCache::FreeAllMethodHeaders under JIT lock.

This ensures that both ArtMethods and the code are removed
from the JIT atomically (otherwise other thread can interfere).

I am not sure if this is partial issue, but keeping the
allocated code and methods in sync makes sanity check happier.

Bug: 151137723
Test: test.py -b -r --host --jit --64
Change-Id: I8580a9112c86b277631c0b2f9dd245fd55a73559
diff --git a/runtime/jit/jit_code_cache.cc b/runtime/jit/jit_code_cache.cc
index 0c3b18b..75ce1c0 100644
--- a/runtime/jit/jit_code_cache.cc
+++ b/runtime/jit/jit_code_cache.cc
@@ -484,7 +484,6 @@
   // first since once we do FreeCode() below, the memory can be reused
   // so it's possible for the same method_header to start representing
   // different compile code.
-  MutexLock mu(Thread::Current(), *Locks::jit_lock_);
   {
     MutexLock mu2(Thread::Current(), *Locks::cha_lock_);
     Runtime::Current()->GetClassLinker()->GetClassHierarchyAnalysis()
@@ -551,8 +550,8 @@
         ++it;
       }
     }
+    FreeAllMethodHeaders(method_headers);
   }
-  FreeAllMethodHeaders(method_headers);
 }
 
 bool JitCodeCache::IsWeakAccessEnabled(Thread* self) const {
@@ -1213,8 +1212,8 @@
         it = method_code_map_.erase(it);
       }
     }
+    FreeAllMethodHeaders(method_headers);
   }
-  FreeAllMethodHeaders(method_headers);
 }
 
 bool JitCodeCache::GetGarbageCollectCode() {
diff --git a/runtime/jit/jit_code_cache.h b/runtime/jit/jit_code_cache.h
index 1bbe2a4..9ef1e4f 100644
--- a/runtime/jit/jit_code_cache.h
+++ b/runtime/jit/jit_code_cache.h
@@ -431,7 +431,7 @@
   // 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::jit_lock_)
       REQUIRES(!Locks::cha_lock_);
 
   // Removes method from the cache. The caller must ensure that all threads