Move location of where we clear thread pool tasks post fork.

Bug: 119800099
Test: boot jitzygote, system server methods get JIT compiled.

Change-Id: Ieb66b36aefad7a570ebfdfa296a2713c20d36f22
diff --git a/runtime/jit/jit.cc b/runtime/jit/jit.cc
index 4430c1b..b6d2687 100644
--- a/runtime/jit/jit.cc
+++ b/runtime/jit/jit.cc
@@ -1043,8 +1043,6 @@
 
 void Jit::PostForkChildAction(bool is_system_server, bool is_zygote) {
   if (is_zygote || Runtime::Current()->IsSafeMode()) {
-    // Remove potential tasks that have been inherited from the zygote.
-    thread_pool_->RemoveAllTasks(Thread::Current());
     // Delete the thread pool, we are not going to JIT.
     thread_pool_.reset(nullptr);
     return;
@@ -1058,8 +1056,6 @@
       !Runtime::Current()->GetInstrumentation()->AreExitStubsInstalled());
 
   if (thread_pool_ != nullptr) {
-    // Remove potential tasks that have been inherited from the zygote.
-    thread_pool_->RemoveAllTasks(Thread::Current());
     if (is_system_server &&
         Runtime::Current()->IsUsingApexBootImageLocation() &&
         UseJitCompilation()) {
diff --git a/runtime/jit/jit_code_cache.cc b/runtime/jit/jit_code_cache.cc
index a29201c..ffb0117 100644
--- a/runtime/jit/jit_code_cache.cc
+++ b/runtime/jit/jit_code_cache.cc
@@ -1742,7 +1742,18 @@
 }
 
 void JitCodeCache::PostForkChildAction(bool is_system_server, bool is_zygote) {
-  MutexLock mu(Thread::Current(), *Locks::jit_lock_);
+  Thread* self = Thread::Current();
+
+  // Remove potential tasks that have been inherited from the zygote.
+  // We do this now and not in Jit::PostForkChildAction, as system server calls
+  // JitCodeCache::PostForkChildAction first, and then does some code loading
+  // that may result in new JIT tasks that we want to keep.
+  ThreadPool* pool = Runtime::Current()->GetJit()->GetThreadPool();
+  if (pool != nullptr) {
+    pool->RemoveAllTasks(self);
+  }
+
+  MutexLock mu(self, *Locks::jit_lock_);
 
   // Reset potential writable MemMaps inherited from the zygote. We never want
   // to write to them.