[jitzygote] Remap boot boot image methods in zygote when single-threaded.

This avoids concurrent updates made by other threads while copying the
data.

Bug: 143569713
Bug: 119800099
Test: boots
Change-Id: Ife0141a45b8d8671ab0c5868ccf90b4799a9f5db
diff --git a/runtime/jit/jit.cc b/runtime/jit/jit.cc
index 25b5be5..e64908d 100644
--- a/runtime/jit/jit.cc
+++ b/runtime/jit/jit.cc
@@ -710,6 +710,8 @@
     }
   }
 
+  LOG(INFO) << "Successfully notified child processes on sharing boot image methods";
+
   // Mark that compilation of boot classpath is done, and memory can now be
   // shared. Other processes will pick up this information.
   code_cache_->GetZygoteMap()->SetCompilationState(ZygoteCompilationState::kNotifiedOk);
@@ -821,9 +823,9 @@
     }
 
     if (Runtime::Current()->IsZygote()) {
-      // Copy the boot image methods data to the mappings we created to share
-      // with the children.
-      Runtime::Current()->GetJit()->NotifyZygoteCompilationDone();
+      // Record that we are done compiling the profile.
+      Runtime::Current()->GetJit()->GetCodeCache()->GetZygoteMap()->SetCompilationState(
+          ZygoteCompilationState::kDone);
     }
   }
 
@@ -1624,6 +1626,15 @@
   if (thread_pool_ == nullptr) {
     return;
   }
+  if (Runtime::Current()->IsZygote() &&
+      code_cache_->GetZygoteMap()->IsCompilationDoneButNotNotified()) {
+    // Copy the boot image methods data to the mappings we created to share
+    // with the children. We do this here as we are the only thread running and
+    // we don't risk other threads concurrently updating the ArtMethod's.
+    CHECK_EQ(GetTaskCount(), 1);
+    NotifyZygoteCompilationDone();
+    CHECK(code_cache_->GetZygoteMap()->IsCompilationNotified());
+  }
   thread_pool_->CreateThreads();
 }
 
diff --git a/runtime/jit/jit_code_cache.h b/runtime/jit/jit_code_cache.h
index 637d5e3..58cf0e3 100644
--- a/runtime/jit/jit_code_cache.h
+++ b/runtime/jit/jit_code_cache.h
@@ -127,12 +127,17 @@
     region_->WriteData(compilation_state_, state);
   }
 
+  bool IsCompilationDoneButNotNotified() const {
+    return compilation_state_ != nullptr && *compilation_state_ == ZygoteCompilationState::kDone;
+  }
+
   bool IsCompilationNotified() const {
-    return *compilation_state_ > ZygoteCompilationState::kDone;
+    return compilation_state_ != nullptr && *compilation_state_ > ZygoteCompilationState::kDone;
   }
 
   bool CanMapBootImageMethods() const {
-    return *compilation_state_ == ZygoteCompilationState::kNotifiedOk;
+    return compilation_state_ != nullptr &&
+           *compilation_state_ == ZygoteCompilationState::kNotifiedOk;
   }
 
  private: