Revert^2 "Refactor code around JIT creation."

This reverts commit 763cd98161424cf19af2f113a6802f04860dcd6e.

Bug: 119800099

Reason for revert: Updated the DCHECKo take into account current state
of zygote not having a thread pool.

Change-Id: I1181ff85e7aebd062ee892548b80ab3de06a5ac7
diff --git a/runtime/class_linker.cc b/runtime/class_linker.cc
index ce7dfaf..bcc05c7 100644
--- a/runtime/class_linker.cc
+++ b/runtime/class_linker.cc
@@ -3145,7 +3145,7 @@
     return (jit == nullptr) || !jit->GetCodeCache()->ContainsPc(quick_code);
   }
 
-  if (runtime->IsNativeDebuggableZygoteOK()) {
+  if (runtime->IsNativeDebuggable()) {
     DCHECK(runtime->UseJitCompilation() && runtime->GetJit()->JitAtFirstUse());
     // If we are doing native debugging, ignore application's AOT code,
     // since we want to JIT it (at first use) with extra stackmaps for native
diff --git a/runtime/jit/jit.cc b/runtime/jit/jit.cc
index d67d9dc..4a3ef07 100644
--- a/runtime/jit/jit.cc
+++ b/runtime/jit/jit.cc
@@ -56,10 +56,12 @@
 // JIT compiler
 void* Jit::jit_library_handle_ = nullptr;
 void* Jit::jit_compiler_handle_ = nullptr;
-void* (*Jit::jit_load_)(bool*) = nullptr;
+void* (*Jit::jit_load_)(void) = nullptr;
 void (*Jit::jit_unload_)(void*) = nullptr;
 bool (*Jit::jit_compile_method_)(void*, ArtMethod*, Thread*, bool) = nullptr;
 void (*Jit::jit_types_loaded_)(void*, mirror::Class**, size_t count) = nullptr;
+bool (*Jit::jit_generate_debug_info_)(void*) = nullptr;
+void (*Jit::jit_update_options_)(void*) = nullptr;
 
 struct StressModeHelper {
   DECLARE_RUNTIME_DEBUG_FLAG(kSlowMode);
@@ -179,20 +181,21 @@
     LOG(WARNING) << "Not creating JIT: library not loaded";
     return nullptr;
   }
-  bool will_generate_debug_symbols = false;
-  jit_compiler_handle_ = (jit_load_)(&will_generate_debug_symbols);
+  jit_compiler_handle_ = (jit_load_)();
   if (jit_compiler_handle_ == nullptr) {
     LOG(WARNING) << "Not creating JIT: failed to allocate a compiler";
     return nullptr;
   }
   std::unique_ptr<Jit> jit(new Jit(code_cache, options));
-  jit->generate_debug_info_ = will_generate_debug_symbols;
 
+  // If the code collector is enabled, check if that still holds:
   // With 'perf', we want a 1-1 mapping between an address and a method.
   // We aren't able to keep method pointers live during the instrumentation method entry trampoline
   // so we will just disable jit-gc if we are doing that.
-  code_cache->SetGarbageCollectCode(!jit->generate_debug_info_ &&
-      !Runtime::Current()->GetInstrumentation()->AreExitStubsInstalled());
+  if (code_cache->GetGarbageCollectCode()) {
+    code_cache->SetGarbageCollectCode(!jit_generate_debug_info_(jit_compiler_handle_) &&
+        !Runtime::Current()->GetInstrumentation()->AreExitStubsInstalled());
+  }
 
   VLOG(jit) << "JIT created with initial_capacity="
       << PrettySize(options->GetCodeCacheInitialCapacity())
@@ -200,13 +203,21 @@
       << ", compile_threshold=" << options->GetCompileThreshold()
       << ", profile_saver_options=" << options->GetProfileSaverOptions();
 
-  jit->CreateThreadPool();
-
   // Notify native debugger about the classes already loaded before the creation of the jit.
   jit->DumpTypeInfoForLoadedTypes(Runtime::Current()->GetClassLinker());
   return jit.release();
 }
 
+template <typename T>
+bool Jit::LoadSymbol(T* address, const char* name, std::string* error_msg) {
+  *address = reinterpret_cast<T>(dlsym(jit_library_handle_, name));
+  if (*address == nullptr) {
+    *error_msg = std::string("JIT couldn't find ") + name + std::string(" entry point");
+    return false;
+  }
+  return true;
+}
+
 bool Jit::LoadCompilerLibrary(std::string* error_msg) {
   jit_library_handle_ = dlopen(
       kIsDebugBuild ? "libartd-compiler.so" : "libart-compiler.so", RTLD_NOW);
@@ -216,31 +227,16 @@
     *error_msg = oss.str();
     return false;
   }
-  jit_load_ = reinterpret_cast<void* (*)(bool*)>(dlsym(jit_library_handle_, "jit_load"));
-  if (jit_load_ == nullptr) {
+  bool all_resolved = true;
+  all_resolved = all_resolved && LoadSymbol(&jit_load_, "jit_load", error_msg);
+  all_resolved = all_resolved && LoadSymbol(&jit_unload_, "jit_unload", error_msg);
+  all_resolved = all_resolved && LoadSymbol(&jit_compile_method_, "jit_compile_method", error_msg);
+  all_resolved = all_resolved && LoadSymbol(&jit_types_loaded_, "jit_types_loaded", error_msg);
+  all_resolved = all_resolved && LoadSymbol(&jit_update_options_, "jit_update_options", error_msg);
+  all_resolved = all_resolved &&
+      LoadSymbol(&jit_generate_debug_info_, "jit_generate_debug_info", error_msg);
+  if (!all_resolved) {
     dlclose(jit_library_handle_);
-    *error_msg = "JIT couldn't find jit_load entry point";
-    return false;
-  }
-  jit_unload_ = reinterpret_cast<void (*)(void*)>(
-      dlsym(jit_library_handle_, "jit_unload"));
-  if (jit_unload_ == nullptr) {
-    dlclose(jit_library_handle_);
-    *error_msg = "JIT couldn't find jit_unload entry point";
-    return false;
-  }
-  jit_compile_method_ = reinterpret_cast<bool (*)(void*, ArtMethod*, Thread*, bool)>(
-      dlsym(jit_library_handle_, "jit_compile_method"));
-  if (jit_compile_method_ == nullptr) {
-    dlclose(jit_library_handle_);
-    *error_msg = "JIT couldn't find jit_compile_method entry point";
-    return false;
-  }
-  jit_types_loaded_ = reinterpret_cast<void (*)(void*, mirror::Class**, size_t)>(
-      dlsym(jit_library_handle_, "jit_types_loaded"));
-  if (jit_types_loaded_ == nullptr) {
-    dlclose(jit_library_handle_);
-    *error_msg = "JIT couldn't find jit_types_loaded entry point";
     return false;
   }
   return true;
@@ -296,7 +292,11 @@
 }
 
 void Jit::CreateThreadPool() {
-  // There is a DCHECK in the 'AddSamples' method to ensure the tread pool
+  if (Runtime::Current()->IsSafeMode()) {
+    // Never create the pool in safe mode.
+    return;
+  }
+  // There is a DCHECK in the 'AddSamples' method to ensure the thread pool
   // is not null when we instrument.
 
   // We need peers as we may report the JIT thread, e.g., in the debugger.
@@ -375,7 +375,7 @@
     return;
   }
   jit::Jit* jit = Runtime::Current()->GetJit();
-  if (jit->generate_debug_info_) {
+  if (jit_generate_debug_info_(jit->jit_compiler_handle_)) {
     DCHECK(jit->jit_types_loaded_ != nullptr);
     jit->jit_types_loaded_(jit->jit_compiler_handle_, &type, 1);
   }
@@ -390,7 +390,7 @@
     std::vector<mirror::Class*> classes_;
   };
 
-  if (generate_debug_info_) {
+  if (jit_generate_debug_info_(jit_compiler_handle_)) {
     ScopedObjectAccess so(Thread::Current());
 
     CollectClasses visitor;
@@ -630,8 +630,11 @@
 
 void Jit::AddSamples(Thread* self, ArtMethod* method, uint16_t count, bool with_backedges) {
   if (thread_pool_ == nullptr) {
-    // Should only see this when shutting down.
-    DCHECK(Runtime::Current()->IsShuttingDown(self));
+    // Should only see this when shutting down, starting up, or in zygote, which doesn't
+    // have a thread pool.
+    DCHECK(Runtime::Current()->IsShuttingDown(self) ||
+           !Runtime::Current()->IsFinishedStarting() ||
+           Runtime::Current()->IsZygote());
     return;
   }
   if (IgnoreSamplesForMethod(method)) {
@@ -795,5 +798,15 @@
   }
 }
 
+void Jit::PostForkChildAction() {
+  // At this point, the compiler options have been adjusted to the particular configuration
+  // of the forked child. Parse them again.
+  jit_update_options_(jit_compiler_handle_);
+
+  // Adjust the status of code cache collection: the status from zygote was to not collect.
+  code_cache_->SetGarbageCollectCode(!jit_generate_debug_info_(jit_compiler_handle_) &&
+      !Runtime::Current()->GetInstrumentation()->AreExitStubsInstalled());
+}
+
 }  // namespace jit
 }  // namespace art
diff --git a/runtime/jit/jit.h b/runtime/jit/jit.h
index 46b0762..e12b032 100644
--- a/runtime/jit/jit.h
+++ b/runtime/jit/jit.h
@@ -100,10 +100,6 @@
     return use_jit_compilation_;
   }
 
-  bool RWXMemoryAllowed() const {
-    return rwx_memory_allowed_;
-  }
-
   void SetUseJitCompilation(bool b) {
     use_jit_compilation_ = b;
   }
@@ -125,10 +121,6 @@
     compile_threshold_ = 0;
   }
 
-  void SetRWXMemoryAllowed(bool rwx_allowed) {
-    rwx_memory_allowed_ = rwx_allowed;
-  }
-
  private:
   bool use_jit_compilation_;
   size_t code_cache_initial_capacity_;
@@ -140,7 +132,6 @@
   uint16_t invoke_transition_weight_;
   bool dump_info_on_shutdown_;
   int thread_pool_pthread_priority_;
-  bool rwx_memory_allowed_;
   ProfileSaverOptions profile_saver_options_;
 
   JitOptions()
@@ -153,8 +144,7 @@
         priority_thread_weight_(0),
         invoke_transition_weight_(0),
         dump_info_on_shutdown_(false),
-        thread_pool_pthread_priority_(kJitPoolThreadPthreadDefaultPriority),
-        rwx_memory_allowed_(true) {}
+        thread_pool_pthread_priority_(kJitPoolThreadPthreadDefaultPriority) {}
 
   DISALLOW_COPY_AND_ASSIGN(JitOptions);
 };
@@ -295,6 +285,9 @@
   // Start JIT threads.
   void Start();
 
+  // Transition to a zygote child state.
+  void PostForkChildAction();
+
  private:
   Jit(JitCodeCache* code_cache, JitOptions* options);
 
@@ -303,13 +296,13 @@
   // JIT compiler
   static void* jit_library_handle_;
   static void* jit_compiler_handle_;
-  static void* (*jit_load_)(bool*);
+  static void* (*jit_load_)(void);
   static void (*jit_unload_)(void*);
   static bool (*jit_compile_method_)(void*, ArtMethod*, Thread*, bool);
   static void (*jit_types_loaded_)(void*, mirror::Class**, size_t count);
-
-  // Whether we should generate debug info when compiling.
-  bool generate_debug_info_;
+  static void (*jit_update_options_)(void*);
+  static bool (*jit_generate_debug_info_)(void*);
+  template <typename T> static bool LoadSymbol(T*, const char* symbol, std::string* error_msg);
 
   // JIT resources owned by runtime.
   jit::JitCodeCache* const code_cache_;
diff --git a/runtime/jit/jit_code_cache.cc b/runtime/jit/jit_code_cache.cc
index 1701ca8..de262a8 100644
--- a/runtime/jit/jit_code_cache.cc
+++ b/runtime/jit/jit_code_cache.cc
@@ -64,6 +64,11 @@
 static constexpr size_t kCodeSizeLogThreshold = 50 * KB;
 static constexpr size_t kStackMapSizeLogThreshold = 50 * KB;
 
+// Data cache will be half of the capacity
+// Code cache will be the other half of the capacity.
+// TODO: Make this variable?
+static constexpr size_t kCodeAndDataCapacityDivider = 2;
+
 static constexpr int kProtR = PROT_READ;
 static constexpr int kProtRW = PROT_READ | PROT_WRITE;
 static constexpr int kProtRWX = PROT_READ | PROT_WRITE | PROT_EXEC;
@@ -183,69 +188,45 @@
   std::vector<ArtMethod*> methods_;
 };
 
-JitCodeCache* JitCodeCache::Create(size_t initial_capacity,
-                                   size_t max_capacity,
-                                   bool used_only_for_profile_data,
-                                   bool rwx_memory_allowed,
-                                   std::string* error_msg) {
+bool JitCodeCache::InitializeMappings(bool rwx_memory_allowed,
+                                      bool is_zygote,
+                                      std::string* error_msg) {
   ScopedTrace trace(__PRETTY_FUNCTION__);
-  CHECK_GE(max_capacity, initial_capacity);
 
-  // We need to have 32 bit offsets from method headers in code cache which point to things
-  // in the data cache. If the maps are more than 4G apart, having multiple maps wouldn't work.
-  // Ensure we're below 1 GB to be safe.
-  if (max_capacity > 1 * GB) {
-    std::ostringstream oss;
-    oss << "Maxium code cache capacity is limited to 1 GB, "
-        << PrettySize(max_capacity) << " is too big";
-    *error_msg = oss.str();
-    return nullptr;
-  }
-
-  // Register for membarrier expedited sync core if JIT will be generating code.
-  if (!used_only_for_profile_data) {
-    if (art::membarrier(art::MembarrierCommand::kRegisterPrivateExpeditedSyncCore) != 0) {
-      // MEMBARRIER_CMD_PRIVATE_EXPEDITED_SYNC_CORE ensures that CPU instruction pipelines are
-      // flushed and it's used when adding code to the JIT. The memory used by the new code may
-      // have just been released and, in theory, the old code could still be in a pipeline.
-      VLOG(jit) << "Kernel does not support membarrier sync-core";
-    }
-  }
+  const size_t capacity = max_capacity_;
+  const size_t data_capacity = capacity / kCodeAndDataCapacityDivider;
+  const size_t exec_capacity = capacity - data_capacity;
 
   // File descriptor enabling dual-view mapping of code section.
   unique_fd mem_fd;
 
-  // Bionic supports memfd_create, but the call may fail on older kernels.
-  mem_fd = unique_fd(art::memfd_create("/jit-cache", /* flags= */ 0));
-  if (mem_fd.get() < 0) {
-    std::ostringstream oss;
-    oss << "Failed to initialize dual view JIT. memfd_create() error: " << strerror(errno);
-    if (!rwx_memory_allowed) {
-      // Without using RWX page permissions, the JIT can not fallback to single mapping as it
-      // requires tranitioning the code pages to RWX for updates.
-      *error_msg = oss.str();
-      return nullptr;
+  // Zygote shouldn't create a shared mapping for JIT, so we cannot use dual view
+  // for it.
+  if (!is_zygote) {
+    // Bionic supports memfd_create, but the call may fail on older kernels.
+    mem_fd = unique_fd(art::memfd_create("/jit-cache", /* flags= */ 0));
+    if (mem_fd.get() < 0) {
+      std::ostringstream oss;
+      oss << "Failed to initialize dual view JIT. memfd_create() error: " << strerror(errno);
+      if (!rwx_memory_allowed) {
+        // Without using RWX page permissions, the JIT can not fallback to single mapping as it
+        // requires tranitioning the code pages to RWX for updates.
+        *error_msg = oss.str();
+        return false;
+      }
+      VLOG(jit) << oss.str();
     }
-    VLOG(jit) << oss.str();
   }
 
-  if (mem_fd.get() >= 0 && ftruncate(mem_fd, max_capacity) != 0) {
+  if (mem_fd.get() >= 0 && ftruncate(mem_fd, capacity) != 0) {
     std::ostringstream oss;
     oss << "Failed to initialize memory file: " << strerror(errno);
     *error_msg = oss.str();
-    return nullptr;
+    return false;
   }
 
-  // Data cache will be half of the initial allocation.
-  // Code cache will be the other half of the initial allocation.
-  // TODO: Make this variable?
-
-  // Align both capacities to page size, as that's the unit mspaces use.
-  initial_capacity = RoundDown(initial_capacity, 2 * kPageSize);
-  max_capacity = RoundDown(max_capacity, 2 * kPageSize);
-  const size_t data_capacity = max_capacity / 2;
-  const size_t exec_capacity = used_only_for_profile_data ? 0 : max_capacity - data_capacity;
-  DCHECK_LE(data_capacity + exec_capacity, max_capacity);
+  std::string data_cache_name = is_zygote ? "zygote-data-code-cache" : "data-code-cache";
+  std::string exec_cache_name = is_zygote ? "zygote-jit-code-cache" : "jit-code-cache";
 
   std::string error_str;
   // Map name specific for android_os_Debug.cpp accounting.
@@ -285,7 +266,7 @@
         mem_fd,
         /* start= */ 0,
         /* low_4gb= */ true,
-        "data-code-cache",
+        data_cache_name.c_str(),
         &error_str);
   } else {
     // Single view of JIT code cache case. Create an initial mapping of data pages large enough
@@ -304,7 +285,7 @@
     // back to RX after the update.
     base_flags = MAP_PRIVATE | MAP_ANON;
     data_pages = MemMap::MapAnonymous(
-        "data-code-cache",
+        data_cache_name.c_str(),
         data_capacity + exec_capacity,
         kProtRW,
         /* low_4gb= */ true,
@@ -313,9 +294,9 @@
 
   if (!data_pages.IsValid()) {
     std::ostringstream oss;
-    oss << "Failed to create read write cache: " << error_str << " size=" << max_capacity;
+    oss << "Failed to create read write cache: " << error_str << " size=" << capacity;
     *error_msg = oss.str();
-    return nullptr;
+    return false;
   }
 
   MemMap exec_pages;
@@ -326,7 +307,7 @@
     // (for processes that cannot map WX pages). Otherwise, this region does not need to be
     // executable as there is no code in the cache yet.
     exec_pages = data_pages.RemapAtEnd(divider,
-                                       "jit-code-cache",
+                                       exec_cache_name.c_str(),
                                        kProtRX,
                                        base_flags | MAP_FIXED,
                                        mem_fd.get(),
@@ -334,21 +315,22 @@
                                        &error_str);
     if (!exec_pages.IsValid()) {
       std::ostringstream oss;
-      oss << "Failed to create read execute code cache: " << error_str << " size=" << max_capacity;
+      oss << "Failed to create read execute code cache: " << error_str << " size=" << capacity;
       *error_msg = oss.str();
-      return nullptr;
+      return false;
     }
 
     if (mem_fd.get() >= 0) {
       // For dual view, create the secondary view of code memory used for updating code. This view
       // is never executable.
+      std::string name = exec_cache_name + "-rw";
       non_exec_pages = MemMap::MapFile(exec_capacity,
                                        kProtR,
                                        base_flags,
                                        mem_fd,
                                        /* start= */ data_capacity,
                                        /* low_4GB= */ false,
-                                       "jit-code-cache-rw",
+                                       name.c_str(),
                                        &error_str);
       if (!non_exec_pages.IsValid()) {
         static const char* kFailedNxView = "Failed to map non-executable view of JIT code cache";
@@ -357,44 +339,77 @@
           VLOG(jit) << kFailedNxView;
         } else {
           *error_msg = kFailedNxView;
-          return nullptr;
+          return false;
         }
       }
     }
   } else {
     // Profiling only. No memory for code required.
-    DCHECK(used_only_for_profile_data);
   }
 
-  const size_t initial_data_capacity = initial_capacity / 2;
-  const size_t initial_exec_capacity =
-      (exec_capacity == 0) ? 0 : (initial_capacity - initial_data_capacity);
-
-  return new JitCodeCache(
-      std::move(data_pages),
-      std::move(exec_pages),
-      std::move(non_exec_pages),
-      initial_data_capacity,
-      initial_exec_capacity,
-      max_capacity);
+  data_pages_ = std::move(data_pages);
+  exec_pages_ = std::move(exec_pages);
+  non_exec_pages_ = std::move(non_exec_pages);
+  return true;
 }
 
-JitCodeCache::JitCodeCache(MemMap&& data_pages,
-                           MemMap&& exec_pages,
-                           MemMap&& non_exec_pages,
-                           size_t initial_data_capacity,
-                           size_t initial_exec_capacity,
-                           size_t max_capacity)
+JitCodeCache* JitCodeCache::Create(bool used_only_for_profile_data,
+                                   bool rwx_memory_allowed,
+                                   bool is_zygote,
+                                   std::string* error_msg) {
+  // Register for membarrier expedited sync core if JIT will be generating code.
+  if (!used_only_for_profile_data) {
+    if (art::membarrier(art::MembarrierCommand::kRegisterPrivateExpeditedSyncCore) != 0) {
+      // MEMBARRIER_CMD_PRIVATE_EXPEDITED_SYNC_CORE ensures that CPU instruction pipelines are
+      // flushed and it's used when adding code to the JIT. The memory used by the new code may
+      // have just been released and, in theory, the old code could still be in a pipeline.
+      VLOG(jit) << "Kernel does not support membarrier sync-core";
+    }
+  }
+
+  // Check whether the provided max capacity in options is below 1GB.
+  size_t max_capacity = Runtime::Current()->GetJITOptions()->GetCodeCacheMaxCapacity();
+  // We need to have 32 bit offsets from method headers in code cache which point to things
+  // in the data cache. If the maps are more than 4G apart, having multiple maps wouldn't work.
+  // Ensure we're below 1 GB to be safe.
+  if (max_capacity > 1 * GB) {
+    std::ostringstream oss;
+    oss << "Maxium code cache capacity is limited to 1 GB, "
+        << PrettySize(max_capacity) << " is too big";
+    *error_msg = oss.str();
+    return nullptr;
+  }
+
+  size_t initial_capacity = Runtime::Current()->GetJITOptions()->GetCodeCacheInitialCapacity();
+
+  std::unique_ptr<JitCodeCache> jit_code_cache(new JitCodeCache());
+
+  MutexLock mu(Thread::Current(), jit_code_cache->lock_);
+  jit_code_cache->InitializeState(initial_capacity, max_capacity);
+
+  // Zygote should never collect code to share the memory with the children.
+  if (is_zygote) {
+    jit_code_cache->SetGarbageCollectCode(false);
+  }
+
+  if (!jit_code_cache->InitializeMappings(rwx_memory_allowed, is_zygote, error_msg)) {
+    return nullptr;
+  }
+
+  jit_code_cache->InitializeSpaces();
+
+  VLOG(jit) << "Created jit code cache: initial capacity="
+            << PrettySize(initial_capacity)
+            << ", maximum capacity="
+            << PrettySize(max_capacity);
+
+  return jit_code_cache.release();
+}
+
+JitCodeCache::JitCodeCache()
     : lock_("Jit code cache", kJitCodeCacheLock),
       lock_cond_("Jit code cache condition variable", lock_),
       collection_in_progress_(false),
-      data_pages_(std::move(data_pages)),
-      exec_pages_(std::move(exec_pages)),
-      non_exec_pages_(std::move(non_exec_pages)),
-      max_capacity_(max_capacity),
-      current_capacity_(initial_exec_capacity + initial_data_capacity),
-      data_end_(initial_data_capacity),
-      exec_end_(initial_exec_capacity),
       last_collection_increased_code_cache_(false),
       garbage_collect_code_(true),
       used_memory_for_data_(0),
@@ -406,10 +421,31 @@
       histogram_code_memory_use_("Memory used for compiled code", 16),
       histogram_profiling_info_memory_use_("Memory used for profiling info", 16),
       is_weak_access_enabled_(true),
-      inline_cache_cond_("Jit inline cache condition variable", lock_) {
+      inline_cache_cond_("Jit inline cache condition variable", lock_),
+      zygote_data_pages_(),
+      zygote_exec_pages_(),
+      zygote_data_mspace_(nullptr),
+      zygote_exec_mspace_(nullptr) {
+}
 
-  DCHECK_GE(max_capacity, initial_exec_capacity + initial_data_capacity);
+void JitCodeCache::InitializeState(size_t initial_capacity, size_t max_capacity) {
+  CHECK_GE(max_capacity, initial_capacity);
+  CHECK(max_capacity <= 1 * GB) << "The max supported size for JIT code cache is 1GB";
+  // Align both capacities to page size, as that's the unit mspaces use.
+  initial_capacity = RoundDown(initial_capacity, 2 * kPageSize);
+  max_capacity = RoundDown(max_capacity, 2 * kPageSize);
 
+  data_pages_ = MemMap();
+  exec_pages_ = MemMap();
+  non_exec_pages_ = MemMap();
+  initial_capacity_ = initial_capacity;
+  max_capacity_ = max_capacity;
+  current_capacity_ = initial_capacity,
+  data_end_ = initial_capacity / kCodeAndDataCapacityDivider;
+  exec_end_ = initial_capacity - data_end_;
+}
+
+void JitCodeCache::InitializeSpaces() {
   // Initialize the data heap
   data_mspace_ = create_mspace_with_base(data_pages_.Begin(), data_end_, false /*locked*/);
   CHECK(data_mspace_ != nullptr) << "create_mspace_with_base (data) failed";
@@ -427,19 +463,14 @@
     CheckedCall(mprotect, "create code heap", code_heap->Begin(), code_heap->Size(), kProtRW);
     exec_mspace_ = create_mspace_with_base(code_heap->Begin(), exec_end_, false /*locked*/);
     CHECK(exec_mspace_ != nullptr) << "create_mspace_with_base (exec) failed";
-    SetFootprintLimit(current_capacity_);
+    SetFootprintLimit(initial_capacity_);
     // Protect pages containing heap metadata. Updates to the code heap toggle write permission to
     // perform the update and there are no other times write access is required.
     CheckedCall(mprotect, "protect code heap", code_heap->Begin(), code_heap->Size(), kProtR);
   } else {
     exec_mspace_ = nullptr;
-    SetFootprintLimit(current_capacity_);
+    SetFootprintLimit(initial_capacity_);
   }
-
-  VLOG(jit) << "Created jit code cache: initial data size="
-            << PrettySize(initial_data_capacity)
-            << ", initial code size="
-            << PrettySize(initial_exec_capacity);
 }
 
 JitCodeCache::~JitCodeCache() {}
@@ -1339,13 +1370,13 @@
 }
 
 void JitCodeCache::SetFootprintLimit(size_t new_footprint) {
-  size_t per_space_footprint = new_footprint / 2;
-  DCHECK(IsAlignedParam(per_space_footprint, kPageSize));
-  DCHECK_EQ(per_space_footprint * 2, new_footprint);
-  mspace_set_footprint_limit(data_mspace_, per_space_footprint);
+  size_t data_space_footprint = new_footprint / kCodeAndDataCapacityDivider;
+  DCHECK(IsAlignedParam(data_space_footprint, kPageSize));
+  DCHECK_EQ(data_space_footprint * kCodeAndDataCapacityDivider, new_footprint);
+  mspace_set_footprint_limit(data_mspace_, data_space_footprint);
   if (HasCodeMapping()) {
     ScopedCodeCacheWrite scc(this);
-    mspace_set_footprint_limit(exec_mspace_, per_space_footprint);
+    mspace_set_footprint_limit(exec_mspace_, new_footprint - data_space_footprint);
   }
 }
 
@@ -2064,5 +2095,33 @@
   histogram_profiling_info_memory_use_.PrintMemoryUse(os);
 }
 
+void JitCodeCache::PostForkChildAction(bool is_system_server, bool is_zygote) {
+  MutexLock mu(Thread::Current(), lock_);
+  // Currently, we don't expect any compilations from zygote.
+  CHECK_EQ(number_of_compilations_, 0u);
+  CHECK_EQ(number_of_osr_compilations_, 0u);
+  CHECK(jni_stubs_map_.empty());
+  CHECK(method_code_map_.empty());
+  CHECK(osr_code_map_.empty());
+
+  zygote_data_pages_ = std::move(data_pages_);
+  zygote_exec_pages_ = std::move(exec_pages_);
+  zygote_data_mspace_ = data_mspace_;
+  zygote_exec_mspace_ = exec_mspace_;
+
+  size_t initial_capacity = Runtime::Current()->GetJITOptions()->GetCodeCacheInitialCapacity();
+  size_t max_capacity = Runtime::Current()->GetJITOptions()->GetCodeCacheMaxCapacity();
+
+  InitializeState(initial_capacity, max_capacity);
+
+  std::string error_msg;
+  if (!InitializeMappings(/* rwx_memory_allowed= */ !is_system_server, is_zygote, &error_msg)) {
+    LOG(WARNING) << "Could not reset JIT state after zygote fork: " << error_msg;
+    return;
+  }
+
+  InitializeSpaces();
+}
+
 }  // namespace jit
 }  // namespace art
diff --git a/runtime/jit/jit_code_cache.h b/runtime/jit/jit_code_cache.h
index a507563..7a838fd 100644
--- a/runtime/jit/jit_code_cache.h
+++ b/runtime/jit/jit_code_cache.h
@@ -89,10 +89,9 @@
 
   // Create the code cache with a code + data capacity equal to "capacity", error message is passed
   // in the out arg error_msg.
-  static JitCodeCache* Create(size_t initial_capacity,
-                              size_t max_capacity,
-                              bool used_only_for_profile_data,
+  static JitCodeCache* Create(bool used_only_for_profile_data,
                               bool rwx_memory_allowed,
+                              bool is_zygote,
                               std::string* error_msg);
   ~JitCodeCache();
 
@@ -262,14 +261,17 @@
       REQUIRES(!lock_)
       REQUIRES_SHARED(Locks::mutator_lock_);
 
+  void PostForkChildAction(bool is_system_server, bool is_zygote);
+
  private:
-  // Take ownership of maps.
-  JitCodeCache(MemMap&& data_pages,
-               MemMap&& exec_pages,
-               MemMap&& non_exec_pages,
-               size_t initial_data_capacity,
-               size_t initial_exec_capacity,
-               size_t max_capacity);
+  JitCodeCache();
+
+  void InitializeState(size_t initial_capacity, size_t max_capacity) REQUIRES(lock_);
+
+  bool InitializeMappings(bool rwx_memory_allowed, bool is_zygote, std::string* error_msg)
+      REQUIRES(lock_);
+
+  void InitializeSpaces() REQUIRES(lock_);
 
   // Internal version of 'CommitCode' that will not retry if the
   // allocation fails. Return null if the allocation fails.
@@ -421,6 +423,9 @@
   // ProfilingInfo objects we have allocated.
   std::vector<ProfilingInfo*> profiling_infos_ GUARDED_BY(lock_);
 
+  // The initial capacity in bytes this code cache starts with.
+  size_t initial_capacity_ GUARDED_BY(lock_);
+
   // The maximum capacity in bytes this code cache can go to.
   size_t max_capacity_ GUARDED_BY(lock_);
 
@@ -471,10 +476,19 @@
   // Condition to wait on for accessing inline caches.
   ConditionVariable inline_cache_cond_ GUARDED_BY(lock_);
 
+  // Mem map which holds zygote data (stack maps and profiling info).
+  MemMap zygote_data_pages_;
+  // Mem map which holds zygote code and has executable permission.
+  MemMap zygote_exec_pages_;
+  // The opaque mspace for allocating zygote data.
+  void* zygote_data_mspace_ GUARDED_BY(lock_);
+  // The opaque mspace for allocating zygote code.
+  void* zygote_exec_mspace_ GUARDED_BY(lock_);
+
   friend class art::JitJniStubTestHelper;
   friend class ScopedCodeCacheWrite;
 
-  DISALLOW_IMPLICIT_CONSTRUCTORS(JitCodeCache);
+  DISALLOW_COPY_AND_ASSIGN(JitCodeCache);
 };
 
 }  // namespace jit
diff --git a/runtime/native/dalvik_system_ZygoteHooks.cc b/runtime/native/dalvik_system_ZygoteHooks.cc
index 56e9094..530371d 100644
--- a/runtime/native/dalvik_system_ZygoteHooks.cc
+++ b/runtime/native/dalvik_system_ZygoteHooks.cc
@@ -29,6 +29,7 @@
 #include "debugger.h"
 #include "hidden_api.h"
 #include "jit/jit.h"
+#include "jit/jit_code_cache.h"
 #include "jni/java_vm_ext.h"
 #include "jni/jni_internal.h"
 #include "native_util.h"
@@ -292,7 +293,10 @@
   // System server has a window where it can create executable pages for this purpose, but this is
   // turned off after this hook. Consequently, the only JIT mode supported is the dual-view JIT
   // where one mapping is R->RW and the other is RX. Single view requires RX->RWX->RX.
-  Runtime::Current()->CreateJitCodeCache(/*rwx_memory_allowed=*/false);
+  if (Runtime::Current()->GetJit() != nullptr) {
+    Runtime::Current()->GetJit()->GetCodeCache()->PostForkChildAction(
+        /* is_system_server= */ true, /* is_zygote= */ false);
+  }
 }
 
 static void ZygoteHooks_nativePostForkChild(JNIEnv* env,
@@ -332,6 +336,15 @@
   }
 
   Runtime::Current()->GetHeap()->PostForkChildAction(thread);
+  if (Runtime::Current()->GetJit() != nullptr) {
+    if (!is_system_server) {
+      // System server already called the JIT cache post fork action in `nativePostForkSystemServer`.
+      Runtime::Current()->GetJit()->GetCodeCache()->PostForkChildAction(
+          /* is_system_server= */ false, is_zygote);
+    }
+    // This must be called after EnableDebugFeatures.
+    Runtime::Current()->GetJit()->PostForkChildAction();
+  }
 
   // Update tracing.
   if (Trace::GetMethodTracingMode() != TracingMode::kTracingInactive) {
diff --git a/runtime/runtime.cc b/runtime/runtime.cc
index 19c1623..292a424 100644
--- a/runtime/runtime.cc
+++ b/runtime/runtime.cc
@@ -791,6 +791,8 @@
     if (!jit::Jit::LoadCompilerLibrary(&error_msg)) {
       LOG(WARNING) << "Failed to load JIT compiler with error " << error_msg;
     }
+    CreateJitCodeCache(/*rwx_memory_allowed=*/true);
+    CreateJit();
   }
 
   // Send the start phase event. We have to wait till here as this is when the main thread peer
@@ -894,15 +896,8 @@
     }
   }
 
-  if (jit_ == nullptr) {
-    // The system server's code cache was initialized specially. For other zygote forks or
-    // processes create it now.
-    if (!is_system_server) {
-      CreateJitCodeCache(/*rwx_memory_allowed=*/true);
-    }
-    // Note that when running ART standalone (not zygote, nor zygote fork),
-    // the jit may have already been created.
-    CreateJit();
+  if (jit_ != nullptr) {
+    jit_->CreateThreadPool();
   }
 
   // Create the thread pools.
@@ -2493,16 +2488,11 @@
     return;
   }
 
-  // SystemServer has execmem blocked by SELinux so can not use RWX page permissions after the
-  // cache initialized.
-  jit_options_->SetRWXMemoryAllowed(rwx_memory_allowed);
-
   std::string error_msg;
   bool profiling_only = !jit_options_->UseJitCompilation();
-  jit_code_cache_.reset(jit::JitCodeCache::Create(jit_options_->GetCodeCacheInitialCapacity(),
-                                                  jit_options_->GetCodeCacheMaxCapacity(),
-                                                  profiling_only,
-                                                  jit_options_->RWXMemoryAllowed(),
+  jit_code_cache_.reset(jit::JitCodeCache::Create(profiling_only,
+                                                  rwx_memory_allowed,
+                                                  IsZygote(),
                                                   &error_msg));
   if (jit_code_cache_.get() == nullptr) {
     LOG(WARNING) << "Failed to create JIT Code Cache: " << error_msg;
diff --git a/runtime/runtime.h b/runtime/runtime.h
index a696c28..1d58ad7 100644
--- a/runtime/runtime.h
+++ b/runtime/runtime.h
@@ -634,13 +634,6 @@
   void DeoptimizeBootImage();
 
   bool IsNativeDebuggable() const {
-    CHECK(!is_zygote_ || IsAotCompiler());
-    return is_native_debuggable_;
-  }
-
-  // Note: prefer not to use this method, but the checked version above. The separation exists
-  //       as the runtime state may change for a zygote child.
-  bool IsNativeDebuggableZygoteOK() const {
     return is_native_debuggable_;
   }
 
@@ -698,7 +691,6 @@
   double GetHashTableMaxLoadFactor() const;
 
   bool IsSafeMode() const {
-    CHECK(!is_zygote_);
     return safe_mode_;
   }