Snap for 11794311 from f91d21c0b0e4a03102eaa6cfae0f500f8139bc40 to build-tools-releasebuild-tools-release

Change-Id: I50acbe7a0d282ed413555bdc74807df641aca6b4

16 files changed, 507 insertions, 119 deletions

diff --git a/artd/artd.cc b/artd/artd.cc
index 6154c8207b..095e78e495 100644
--- a/artd/artd.cc
+++ b/artd/artd.cc
@@ -117,6 +117,7 @@ using ::android::base::ReadFileToString;
 using ::android::base::Result;
 using ::android::base::Split;
 using ::android::base::StringReplace;
+using ::android::base::Tokenize;
 using ::android::base::WriteStringToFd;
 using ::android::base::WriteStringToFile;
 using ::android::fs_mgr::FstabEntry;
@@ -129,6 +130,7 @@ using ::art::tools::NonFatal;
 using ::ndk::ScopedAStatus;
 
 using TmpProfilePath = ProfilePath::TmpProfilePath;
+using WritableProfilePath = ProfilePath::WritableProfilePath;
 
 constexpr const char* kServiceName = "artd";
 constexpr const char* kPreRebootServiceName = "artd_pre_reboot";
@@ -1370,6 +1372,61 @@ ScopedAStatus Artd::getProfileSize(const ProfilePath& in_profile, int64_t* _aidl
   return ScopedAStatus::ok();
 }
 
+ScopedAStatus Artd::commitPreRebootStagedFiles(
+    const std::vector<ArtifactsPath>& in_artifacts,
+    const std::vector<WritableProfilePath>& in_profiles) {
+  RETURN_FATAL_IF_PRE_REBOOT(options_);
+
+  std::vector<std::pair<std::string, std::string>> files_to_move;
+  std::vector<std::string> files_to_remove;
+
+  for (const ArtifactsPath& artifacts : in_artifacts) {
+    RETURN_FATAL_IF_ARG_IS_PRE_REBOOT(artifacts, "artifacts");
+
+    ArtifactsPath pre_reboot_artifacts = artifacts;
+    pre_reboot_artifacts.isPreReboot = true;
+
+    auto src_artifacts = std::make_unique<RawArtifactsPath>(
+        OR_RETURN_FATAL(BuildArtifactsPath(pre_reboot_artifacts)));
+    auto dst_artifacts =
+        std::make_unique<RawArtifactsPath>(OR_RETURN_FATAL(BuildArtifactsPath(artifacts)));
+
+    if (OS::FileExists(src_artifacts->oat_path.c_str())) {
+      files_to_move.emplace_back(src_artifacts->oat_path, dst_artifacts->oat_path);
+      files_to_move.emplace_back(src_artifacts->vdex_path, dst_artifacts->vdex_path);
+      if (OS::FileExists(src_artifacts->art_path.c_str())) {
+        files_to_move.emplace_back(src_artifacts->art_path, dst_artifacts->art_path);
+      } else {
+        files_to_remove.push_back(dst_artifacts->art_path);
+      }
+    }
+  }
+
+  for (const WritableProfilePath& profile : in_profiles) {
+    RETURN_FATAL_IF_ARG_IS_PRE_REBOOT(profile, "profiles");
+
+    WritableProfilePath pre_reboot_profile = profile;
+    PreRebootFlag(pre_reboot_profile) = true;
+
+    auto src_profile = std::make_unique<std::string>(
+        OR_RETURN_FATAL(BuildWritableProfilePath(pre_reboot_profile)));
+    auto dst_profile =
+        std::make_unique<std::string>(OR_RETURN_FATAL(BuildWritableProfilePath(profile)));
+
+    if (OS::FileExists(src_profile->c_str())) {
+      files_to_move.emplace_back(*src_profile, *dst_profile);
+    }
+  }
+
+  OR_RETURN_NON_FATAL(MoveAllOrAbandon(files_to_move, files_to_remove));
+
+  for (const auto& [src_path, dst_path] : files_to_move) {
+    LOG(INFO) << ART_FORMAT("Committed Pre-reboot staged file '{}' to '{}'", src_path, dst_path);
+  }
+
+  return ScopedAStatus::ok();
+}
+
 Result<void> Artd::Start() {
   OR_RETURN(SetLogVerbosity());
   MemMap::Init();
@@ -1538,8 +1595,7 @@ void Artd::AddCompilerConfigFlags(const std::string& instruction_set,
                      props_->GetOrEmpty("dalvik.vm.dex2oat-very-large"))
       .AddIfNonEmpty(
           "--resolve-startup-const-strings=%s",
-          props_->GetOrEmpty("persist.device_config.runtime.dex2oat_resolve_startup_strings",
-                             "dalvik.vm.dex2oat-resolve-startup-strings"));
+          props_->GetOrEmpty("dalvik.vm.dex2oat-resolve-startup-strings"));
 
   args.AddIf(dexopt_options.debuggable, "--debuggable")
       .AddIf(props_->GetBool("debug.generate-debug-info", /*default_value=*/false),
@@ -1589,6 +1645,11 @@ void Artd::AddPerfConfigFlags(PriorityClass priority_class,
   // It takes longer but reduces the memory footprint.
   dex2oat_args.AddIf(props_->GetBool("ro.config.low_ram", /*default_value=*/false),
                      "--compile-individually");
+
+  for (const std::string& flag :
+       Tokenize(props_->GetOrEmpty("dalvik.vm.dex2oat-flags"), /*delimiters=*/" ")) {
+    dex2oat_args.AddIfNonEmpty("%s", flag);
+  }
 }
 
 Result<int> Artd::ExecAndReturnCode(const std::vector<std::string>& args,
diff --git a/artd/artd.h b/artd/artd.h
index a1f3541fc0..37b6d0c42f 100644
--- a/artd/artd.h
+++ b/artd/artd.h
@@ -221,6 +221,11 @@ class Artd : public aidl::com::android::server::art::BnArtd {
   ndk::ScopedAStatus getProfileSize(const aidl::com::android::server::art::ProfilePath& in_profile,
                                     int64_t* _aidl_return) override;
 
+  ndk::ScopedAStatus commitPreRebootStagedFiles(
+      const std::vector<aidl::com::android::server::art::ArtifactsPath>& in_artifacts,
+      const std::vector<aidl::com::android::server::art::ProfilePath::WritableProfilePath>&
+          in_profiles) override;
+
   ndk::ScopedAStatus preRebootInit() override;
 
   ndk::ScopedAStatus validateDexPath(const std::string& in_dexFile,
diff --git a/artd/artd_test.cc b/artd/artd_test.cc
index fccf091464..858fd9909b 100644
--- a/artd/artd_test.cc
+++ b/artd/artd_test.cc
@@ -943,6 +943,8 @@ TEST_F(ArtdTest, dexoptFlagsFromSystemProps) {
   EXPECT_CALL(*mock_props_, GetProperty("ro.config.low_ram")).WillOnce(Return("1"));
   EXPECT_CALL(*mock_props_, GetProperty("dalvik.vm.appimageformat")).WillOnce(Return("imgfmt"));
   EXPECT_CALL(*mock_props_, GetProperty("dalvik.vm.boot-image")).WillOnce(Return("boot-image"));
+  EXPECT_CALL(*mock_props_, GetProperty("dalvik.vm.dex2oat-flags"))
+      .WillOnce(Return("--flag1 --flag2  --flag3"));
 
   EXPECT_CALL(*mock_exec_utils_,
               DoExecAndReturnCode(
@@ -962,7 +964,10 @@ TEST_F(ArtdTest, dexoptFlagsFromSystemProps) {
                                     Contains("--compile-individually"),
                                     Contains(Flag("--image-format=", "imgfmt")),
                                     Not(Contains("--force-jit-zygote")),
-                                    Contains(Flag("--boot-image=", "boot-image")))),
+                                    Contains(Flag("--boot-image=", "boot-image")),
+                                    Contains("--flag1"),
+                                    Contains("--flag2"),
+                                    Contains("--flag3"))),
                   _,
                   _))
       .WillOnce(Return(0));
@@ -2446,6 +2451,93 @@ TEST_F(ArtdTest, getProfileSize) {
   EXPECT_EQ(aidl_return, 1);
 }
 
+TEST_F(ArtdTest, commitPreRebootStagedFiles) {
+  CreateFile(android_data_ + "/dalvik-cache/arm64/system@app@Foo@Foo.apk@classes.dex.staged",
+             "new_odex_1");
+  CreateFile(android_data_ + "/dalvik-cache/arm64/system@app@Foo@Foo.apk@classes.vdex.staged",
+             "new_vdex_1");
+  CreateFile(android_data_ + "/dalvik-cache/arm64/system@app@Foo@Foo.apk@classes.art.staged",
+             "new_art_1");
+
+  CreateFile(android_data_ + "/dalvik-cache/arm64/system@app@Foo@Foo.apk@classes.dex",
+             "old_odex_1");
+  CreateFile(android_data_ + "/dalvik-cache/arm64/system@app@Foo@Foo.apk@classes.vdex",
+             "old_vdex_1");
+  CreateFile(android_data_ + "/dalvik-cache/arm64/system@app@Foo@Foo.apk@classes.art", "old_art_1");
+
+  CreateFile(android_data_ + "/app/com.android.foo/oat/arm64/base.odex", "old_odex_2");
+  CreateFile(android_data_ + "/app/com.android.foo/oat/arm64/base.vdex", "old_vdex_2");
+  CreateFile(android_data_ + "/app/com.android.foo/oat/arm64/base.art", "old_art_2");
+
+  CreateFile(android_data_ + "/app/com.android.foo/oat/arm64/base.odex.staged", "new_odex_2");
+  CreateFile(android_data_ + "/app/com.android.foo/oat/arm64/base.vdex.staged", "new_vdex_2");
+
+  CreateFile(android_data_ + "/app/com.android.foo/oat/arm/base.odex", "old_odex_3");
+  CreateFile(android_data_ + "/app/com.android.foo/oat/arm/base.vdex", "old_vdex_3");
+  CreateFile(android_data_ + "/app/com.android.foo/oat/arm/base.art", "old_art_3");
+
+  CreateFile(android_data_ + "/misc/profiles/ref/com.android.foo/primary.prof", "old_prof_1");
+  CreateFile(android_data_ + "/misc/profiles/ref/com.android.foo/primary.prof.staged",
+             "new_prof_1");
+
+  CreateFile(android_data_ + "/misc/profiles/ref/com.android.bar/primary.prof", "old_prof_2");
+
+  ASSERT_STATUS_OK(artd_->commitPreRebootStagedFiles(
+      {
+          // Has all new files. All old files should be replaced.
+          ArtifactsPath{
+              .dexPath = "/system/app/Foo/Foo.apk", .isa = "arm64", .isInDalvikCache = true},
+          // Has new files but not ".art" file. Old ".odex" and ".vdex" files should be replaced,
+          // and old ".art" file should be removed.
+          ArtifactsPath{.dexPath = android_data_ + "/app/com.android.foo/base.apk",
+                        .isa = "arm64",
+                        .isInDalvikCache = false},
+          // Has no new file. All old files should be kept.
+          ArtifactsPath{.dexPath = android_data_ + "/app/com.android.foo/base.apk",
+                        .isa = "arm",
+                        .isInDalvikCache = false},
+      },
+      {
+          // Has new file.
+          PrimaryRefProfilePath{.packageName = "com.android.foo", .profileName = "primary"},
+          // Has no new file.
+          PrimaryRefProfilePath{.packageName = "com.android.bar", .profileName = "primary"},
+      }));
+
+  CheckContent(android_data_ + "/dalvik-cache/arm64/system@app@Foo@Foo.apk@classes.dex",
+               "new_odex_1");
+  CheckContent(android_data_ + "/dalvik-cache/arm64/system@app@Foo@Foo.apk@classes.vdex",
+               "new_vdex_1");
+  CheckContent(android_data_ + "/dalvik-cache/arm64/system@app@Foo@Foo.apk@classes.art",
+               "new_art_1");
+
+  CreateFile(android_data_ + "/app/com.android.foo/oat/arm64/base.odex", "new_odex_2");
+  CreateFile(android_data_ + "/app/com.android.foo/oat/arm64/base.vdex", "new_vdex_2");
+  EXPECT_FALSE(std::filesystem::exists(android_data_ + "/app/com.android.foo/oat/arm64/base.art"));
+
+  CheckContent(android_data_ + "/app/com.android.foo/oat/arm/base.odex", "old_odex_3");
+  CheckContent(android_data_ + "/app/com.android.foo/oat/arm/base.vdex", "old_vdex_3");
+  CheckContent(android_data_ + "/app/com.android.foo/oat/arm/base.art", "old_art_3");
+
+  CheckContent(android_data_ + "/misc/profiles/ref/com.android.foo/primary.prof", "new_prof_1");
+
+  CheckContent(android_data_ + "/misc/profiles/ref/com.android.bar/primary.prof", "old_prof_2");
+
+  // All staged files are gone.
+  EXPECT_FALSE(std::filesystem::exists(
+      android_data_ + "/dalvik-cache/arm64/system@app@Foo@Foo.apk@classes.dex.staged"));
+  EXPECT_FALSE(std::filesystem::exists(
+      android_data_ + "/dalvik-cache/arm64/system@app@Foo@Foo.apk@classes.vdex.staged"));
+  EXPECT_FALSE(std::filesystem::exists(
+      android_data_ + "/dalvik-cache/arm64/system@app@Foo@Foo.apk@classes.art.staged"));
+  EXPECT_FALSE(
+      std::filesystem::exists(android_data_ + "/app/com.android.foo/oat/arm64/base.odex.staged"));
+  EXPECT_FALSE(
+      std::filesystem::exists(android_data_ + "/app/com.android.foo/oat/arm64/base.vdex.staged"));
+  EXPECT_FALSE(std::filesystem::exists(android_data_ +
+                                       "/misc/profiles/ref/com.android.foo/primary.prof.staged"));
+}
+
 class ArtdPreRebootTest : public ArtdTest {
  protected:
   void SetUp() override {
diff --git a/artd/binder/com/android/server/art/IArtd.aidl b/artd/binder/com/android/server/art/IArtd.aidl
index 7f0bc5690c..4532a8fb11 100644
--- a/artd/binder/com/android/server/art/IArtd.aidl
+++ b/artd/binder/com/android/server/art/IArtd.aidl
@@ -251,6 +251,22 @@ interface IArtd {
      */
     long getProfileSize(in com.android.server.art.ProfilePath profile);
 
+    /**
+     * Moves the staged files of the given artifacts and profiles to the permanent locations,
+     * replacing old files if they exist. Removes the staged files and restores the old files at
+     * best effort if any error occurs.
+     *
+     * This is intended to be called for a superset of the packages that we actually expect to have
+     * staged files, so missing files are expected.
+     *
+     * Not supported in Pre-reboot Dexopt mode.
+     *
+     * Throws fatal and non-fatal errors.
+     */
+    void commitPreRebootStagedFiles(
+            in List<com.android.server.art.ArtifactsPath> artifacts,
+            in List<com.android.server.art.ProfilePath.WritableProfilePath> profiles);
+
     // The methods below are only for Pre-reboot Dexopt and only supported in Pre-reboot Dexopt
     // mode.
 
diff --git a/artd/file_utils.cc b/artd/file_utils.cc
index 1415efbdb3..195daf6162 100644
--- a/artd/file_utils.cc
+++ b/artd/file_utils.cc
@@ -230,6 +230,20 @@ Result<void> MoveAllOrAbandon(
   return {};
 }
 
+android::base::Result<void> MoveAllOrAbandon(
+    const std::vector<std::pair<std::string, std::string>>& files_to_move,
+    const std::vector<std::string>& files_to_remove) {
+  std::vector<std::pair<std::string_view, std::string_view>> files_to_move_view;
+  std::vector<std::string_view> files_to_remove_view;
+  for (const auto& [src, dst] : files_to_move) {
+    files_to_move_view.emplace_back(src, dst);
+  }
+  for (const std::string& file : files_to_remove) {
+    files_to_remove_view.emplace_back(file);
+  }
+  return MoveAllOrAbandon(files_to_move_view, files_to_remove_view);
+}
+
 std::string NewFile::BuildTempPath(std::string_view final_path, const std::string& id) {
   return ART_FORMAT("{}.{}.tmp", final_path, id);
 }
diff --git a/artd/file_utils.h b/artd/file_utils.h
index a97f52c42a..77c7ffb19c 100644
--- a/artd/file_utils.h
+++ b/artd/file_utils.h
@@ -20,6 +20,7 @@
 #include <sys/types.h>
 
 #include <memory>
+#include <string>
 #include <string_view>
 #include <utility>
 #include <vector>
@@ -147,6 +148,11 @@ android::base::Result<void> MoveAllOrAbandon(
     const std::vector<std::pair<std::string_view, std::string_view>>& files_to_move,
     const std::vector<std::string_view>& files_to_remove = {});
 
+// Same as above, but takes `std::string`s.
+android::base::Result<void> MoveAllOrAbandon(
+    const std::vector<std::pair<std::string, std::string>>& files_to_move,
+    const std::vector<std::string>& files_to_remove = {});
+
 }  // namespace artd
 }  // namespace art
 
diff --git a/artd/path_utils.cc b/artd/path_utils.cc
index d3eb6d4dfe..7cd9413dea 100644
--- a/artd/path_utils.cc
+++ b/artd/path_utils.cc
@@ -240,18 +240,21 @@ Result<std::string> BuildSecondaryCurProfilePath(
       "{}/oat/{}.cur.prof", dex_path.parent_path().string(), dex_path.filename().string());
 }
 
-Result<std::string> BuildFinalProfilePath(const TmpProfilePath& tmp_profile_path) {
-  const WritableProfilePath& final_path = tmp_profile_path.finalPath;
-  switch (final_path.getTag()) {
+Result<std::string> BuildWritableProfilePath(const WritableProfilePath& profile_path) {
+  switch (profile_path.getTag()) {
     case WritableProfilePath::forPrimary:
-      return BuildPrimaryRefProfilePath(final_path.get<WritableProfilePath::forPrimary>());
+      return BuildPrimaryRefProfilePath(profile_path.get<WritableProfilePath::forPrimary>());
     case WritableProfilePath::forSecondary:
-      return BuildSecondaryRefProfilePath(final_path.get<WritableProfilePath::forSecondary>());
+      return BuildSecondaryRefProfilePath(profile_path.get<WritableProfilePath::forSecondary>());
       // No default. All cases should be explicitly handled, or the compilation will fail.
   }
   // This should never happen. Just in case we get a non-enumerator value.
   LOG(FATAL) << ART_FORMAT("Unexpected writable profile path type {}",
-                           fmt::underlying(final_path.getTag()));
+                           fmt::underlying(profile_path.getTag()));
+}
+
+Result<std::string> BuildFinalProfilePath(const TmpProfilePath& tmp_profile_path) {
+  return BuildWritableProfilePath(tmp_profile_path.finalPath);
 }
 
 Result<std::string> BuildTmpProfilePath(const TmpProfilePath& tmp_profile_path) {
diff --git a/artd/path_utils.h b/artd/path_utils.h
index a4a6e4f2a7..43a1043b42 100644
--- a/artd/path_utils.h
+++ b/artd/path_utils.h
@@ -78,6 +78,9 @@ android::base::Result<std::string> BuildSecondaryCurProfilePath(
     const aidl::com::android::server::art::ProfilePath::SecondaryCurProfilePath&
         secondary_cur_profile_path);
 
+android::base::Result<std::string> BuildWritableProfilePath(
+    const aidl::com::android::server::art::ProfilePath::WritableProfilePath& profile_path);
+
 android::base::Result<std::string> BuildFinalProfilePath(
     const aidl::com::android::server::art::ProfilePath::TmpProfilePath& tmp_profile_path);
 
diff --git a/dex2oat/driver/compiler_driver.cc b/dex2oat/driver/compiler_driver.cc
index 329c0d76a1..e830e86d90 100644
--- a/dex2oat/driver/compiler_driver.cc
+++ b/dex2oat/driver/compiler_driver.cc
@@ -1641,7 +1641,7 @@ class ParallelCompilationManager {
 static bool SkipClass(jobject class_loader, const DexFile& dex_file, ObjPtr<mirror::Class> klass)
     REQUIRES_SHARED(Locks::mutator_lock_) {
   DCHECK(klass != nullptr);
-  const DexFile& original_dex_file = *klass->GetDexCache()->GetDexFile();
+  const DexFile& original_dex_file = klass->GetDexFile();
   if (&dex_file != &original_dex_file) {
     if (class_loader == nullptr) {
       LOG(WARNING) << "Skipping class " << klass->PrettyDescriptor() << " from "
@@ -2029,10 +2029,10 @@ class VerifyClassVisitor : public CompilationVisitor {
           break;
         }
       }
-    } else if (&klass->GetDexFile() != &dex_file) {
+    } else if (SkipClass(jclass_loader, dex_file, klass.Get())) {
       // Skip a duplicate class (as the resolved class is from another, earlier dex file).
       return;  // Do not update state.
-    } else if (!SkipClass(jclass_loader, dex_file, klass.Get())) {
+    } else {
       CHECK(klass->IsResolved()) << klass->PrettyClass();
       failure_kind = class_linker->VerifyClass(soa.Self(),
                                                soa.Self()->GetVerifierDeps(),
@@ -2092,9 +2092,6 @@ class VerifyClassVisitor : public CompilationVisitor {
           DCHECK_EQ(failure_kind, verifier::FailureKind::kHardFailure);
         }
       }
-    } else {
-      // Make the skip a soft failure, essentially being considered as verify at runtime.
-      failure_kind = verifier::FailureKind::kSoftFailure;
     }
     verifier::VerifierDeps::MaybeRecordVerificationStatus(soa.Self()->GetVerifierDeps(),
                                                           dex_file,
@@ -2708,8 +2705,6 @@ static void CompileDexFile(CompilerDriver* driver,
       soa.Self()->ClearException();
       dex_cache = hs.NewHandle(class_linker->FindDexCache(soa.Self(), dex_file));
     } else if (SkipClass(jclass_loader, dex_file, klass.Get())) {
-      return;
-    } else if (&klass->GetDexFile() != &dex_file) {
       // Skip a duplicate class (as the resolved class is from another, earlier dex file).
       return;  // Do not update state.
     } else {
diff --git a/libartservice/service/java/com/android/server/art/AidlUtils.java b/libartservice/service/java/com/android/server/art/AidlUtils.java
index d42a6539c4..6531e61668 100644
--- a/libartservice/service/java/com/android/server/art/AidlUtils.java
+++ b/libartservice/service/java/com/android/server/art/AidlUtils.java
@@ -208,6 +208,19 @@ public final class AidlUtils {
     }
 
     @NonNull
+    public static WritableProfilePath toWritableProfilePath(@NonNull ProfilePath profile) {
+        switch (profile.getTag()) {
+            case ProfilePath.primaryRefProfilePath:
+                return WritableProfilePath.forPrimary(profile.getPrimaryRefProfilePath());
+            case ProfilePath.secondaryRefProfilePath:
+                return WritableProfilePath.forSecondary(profile.getSecondaryRefProfilePath());
+            default:
+                throw new IllegalStateException("ProfilePath tag " + profile.getTag()
+                        + " does not represent a writable type");
+        }
+    }
+
+    @NonNull
     public static String toString(@NonNull PrimaryRefProfilePath profile) {
         return String.format(
                 "PrimaryRefProfilePath[packageName = %s, profileName = %s, isPreReboot = %b]",
diff --git a/libartservice/service/java/com/android/server/art/ArtManagerLocal.java b/libartservice/service/java/com/android/server/art/ArtManagerLocal.java
index 276e6eb7f5..afe2ab0538 100644
--- a/libartservice/service/java/com/android/server/art/ArtManagerLocal.java
+++ b/libartservice/service/java/com/android/server/art/ArtManagerLocal.java
@@ -20,8 +20,10 @@ import static com.android.server.art.ArtFileManager.ProfileLists;
 import static com.android.server.art.ArtFileManager.UsableArtifactLists;
 import static com.android.server.art.ArtFileManager.WritableArtifactLists;
 import static com.android.server.art.DexMetadataHelper.DexMetadataInfo;
+import static com.android.server.art.DexUseManagerLocal.SecondaryDexInfo;
 import static com.android.server.art.PrimaryDexUtils.DetailedPrimaryDexInfo;
 import static com.android.server.art.PrimaryDexUtils.PrimaryDexInfo;
+import static com.android.server.art.ProfilePath.WritableProfilePath;
 import static com.android.server.art.ReasonMapping.BatchDexoptReason;
 import static com.android.server.art.ReasonMapping.BootReason;
 import static com.android.server.art.Utils.Abi;
@@ -39,7 +41,10 @@ import android.annotation.SystemApi;
 import android.annotation.SystemService;
 import android.app.job.JobInfo;
 import android.apphibernation.AppHibernationManager;
+import android.content.BroadcastReceiver;
 import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
 import android.os.Binder;
 import android.os.Build;
 import android.os.CancellationSignal;
@@ -58,6 +63,7 @@ import android.util.Pair;
 import androidx.annotation.RequiresApi;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.modules.utils.build.SdkLevel;
 import com.android.server.LocalManagerRegistry;
 import com.android.server.art.model.ArtFlags;
 import com.android.server.art.model.ArtManagedFileStats;
@@ -126,6 +132,8 @@ public final class ArtManagerLocal {
 
     @NonNull private final Injector mInjector;
 
+    private boolean mShouldCommitPreRebootStagedFiles = false;
+
     @Deprecated
     public ArtManagerLocal() {
         mInjector = new Injector();
@@ -864,12 +872,49 @@ public final class ArtManagerLocal {
             @Nullable @CallbackExecutor Executor progressCallbackExecutor,
             @Nullable Consumer<OperationProgress> progressCallback) {
         try (var snapshot = mInjector.getPackageManagerLocal().withFilteredSnapshot()) {
+            if ((bootReason.equals(ReasonMapping.REASON_BOOT_AFTER_OTA)
+                        || bootReason.equals(ReasonMapping.REASON_BOOT_AFTER_MAINLINE_UPDATE))
+                    && SdkLevel.isAtLeastV()) {
+                // The staged files have to be committed in two phases, one during boot, for primary
+                // dex files, and another after boot complete, for secondary dex files. We need to
+                // commit files for primary dex files early because apps will start using them as
+                // soon as the package manager is initialized. We need to wait until boot complete
+                // to commit files for secondary dex files because they are not decrypted before
+                // then.
+                mShouldCommitPreRebootStagedFiles = true;
+                commitPreRebootStagedFiles(snapshot, false /* forSecondary */);
+            }
             dexoptPackages(snapshot, bootReason, new CancellationSignal(), progressCallbackExecutor,
                     progressCallback != null ? Map.of(ArtFlags.PASS_MAIN, progressCallback) : null);
         }
     }
 
     /**
+     * Notifies this class that {@link Context#registerReceiver} is ready for use.
+     *
+     * Should be used by {@link DexUseManagerLocal} ONLY.
+     *
+     * @hide
+     */
+    @RequiresApi(Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
+    void systemReady() {
+        if (mShouldCommitPreRebootStagedFiles) {
+            mInjector.getContext().registerReceiver(new BroadcastReceiver() {
+                @Override
+                public void onReceive(Context context, Intent intent) {
+                    context.unregisterReceiver(this);
+                    if (!SdkLevel.isAtLeastV()) {
+                        throw new IllegalStateException("Broadcast receiver unexpectedly called");
+                    }
+                    try (var snapshot = mInjector.getPackageManagerLocal().withFilteredSnapshot()) {
+                        commitPreRebootStagedFiles(snapshot, true /* forSecondary */);
+                    }
+                }
+            }, new IntentFilter(Intent.ACTION_BOOT_COMPLETED));
+        }
+    }
+
+    /**
      * Notifies ART Service that there are apexes staged for installation on next reboot (see
      * <a href="https://source.android.com/docs/core/ota/apex#apex-manager">the update sequence of
      * an APEX</a>). ART Service may use this to schedule a pre-reboot dexopt job. This might change
@@ -1056,6 +1101,56 @@ public final class ArtManagerLocal {
         }
     }
 
+    /** @param forSecondary true for secondary dex files; false for primary dex files. */
+    @RequiresApi(Build.VERSION_CODES.VANILLA_ICE_CREAM)
+    private void commitPreRebootStagedFiles(
+            @NonNull PackageManagerLocal.FilteredSnapshot snapshot, boolean forSecondary) {
+        try {
+            // Because we don't know for which packages the Pre-reboot Dexopt job has generated
+            // staged files, we call artd for all dexoptable packages, which is a superset of the
+            // packages that we actually expect to have staged files.
+            for (PackageState pkgState : snapshot.getPackageStates().values()) {
+                if (!Utils.canDexoptPackage(pkgState, null /* appHibernationManager */)) {
+                    continue;
+                }
+                AndroidPackage pkg = Utils.getPackageOrThrow(pkgState);
+                var options = ArtFileManager.Options.builder()
+                                      .setForPrimaryDex(!forSecondary)
+                                      .setForSecondaryDex(forSecondary)
+                                      .setExcludeForObsoleteDexesAndLoaders(true)
+                                      .build();
+                List<ArtifactsPath> artifacts =
+                        mInjector.getArtFileManager()
+                                .getWritableArtifacts(pkgState, pkg, options)
+                                .artifacts();
+                List<WritableProfilePath> profiles = mInjector.getArtFileManager()
+                                                             .getProfiles(pkgState, pkg, options)
+                                                             .refProfiles()
+                                                             .stream()
+                                                             .map(AidlUtils::toWritableProfilePath)
+                                                             .collect(Collectors.toList());
+                try {
+                    // The artd method commits all files somewhat transactionally. Here, we are
+                    // committing files transactionally at the package level just for simplicity. In
+                    // fact, we only need transaction on the split level: the artifacts and the
+                    // profile of the same split must be committed transactionally. Consider the
+                    // case where the staged artifacts and profile have less methods than the active
+                    // ones generated by background dexopt, committing the artifacts while failing
+                    // to commit the profile can potentially cause a permanent performance
+                    // regression.
+                    mInjector.getArtd().commitPreRebootStagedFiles(artifacts, profiles);
+                } catch (ServiceSpecificException e) {
+                    Log.e(TAG,
+                            "Failed to commit Pre-reboot staged files for package '"
+                                    + pkgState.getPackageName() + "'",
+                            e);
+                }
+            }
+        } catch (RemoteException e) {
+            Utils.logArtdException(e);
+        }
+    }
+
     /**
      * Should be used by {@link BackgroundDexoptJobService} ONLY.
      *
diff --git a/libartservice/service/java/com/android/server/art/DexUseManagerLocal.java b/libartservice/service/java/com/android/server/art/DexUseManagerLocal.java
index c6394d3bd6..d5b0bab949 100644
--- a/libartservice/service/java/com/android/server/art/DexUseManagerLocal.java
+++ b/libartservice/service/java/com/android/server/art/DexUseManagerLocal.java
@@ -169,6 +169,7 @@ public class DexUseManagerLocal {
     /** Notifies dex use manager that {@link Context#registerReceiver} is ready for use. */
     public void systemReady() {
         Utils.check(!mInjector.isPreReboot());
+        mInjector.getArtManagerLocal().systemReady();
         // Save the data when the device is being shut down. The receiver is blocking, with a
         // 10s timeout.
         mInjector.getContext().registerReceiver(new BroadcastReceiver() {
@@ -1221,5 +1222,10 @@ public class DexUseManagerLocal {
             return Objects.requireNonNull(
                     LocalManagerRegistry.getManager(PackageManagerLocal.class));
         }
+
+        @NonNull
+        public ArtManagerLocal getArtManagerLocal() {
+            return Objects.requireNonNull(LocalManagerRegistry.getManager(ArtManagerLocal.class));
+        }
     }
 }
diff --git a/libartservice/service/javatests/com/android/server/art/ArtManagerLocalTest.java b/libartservice/service/javatests/com/android/server/art/ArtManagerLocalTest.java
index fa217ab89d..098c28b108 100644
--- a/libartservice/service/javatests/com/android/server/art/ArtManagerLocalTest.java
+++ b/libartservice/service/javatests/com/android/server/art/ArtManagerLocalTest.java
@@ -48,6 +48,9 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.apphibernation.AppHibernationManager;
+import android.content.BroadcastReceiver;
+import android.content.Context;
+import android.content.Intent;
 import android.os.CancellationSignal;
 import android.os.ParcelFileDescriptor;
 import android.os.Process;
@@ -83,6 +86,7 @@ import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 import org.junit.runners.Parameterized.Parameter;
 import org.junit.runners.Parameterized.Parameters;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 
 import java.io.File;
@@ -132,12 +136,14 @@ public class ArtManagerLocalTest {
     @Mock private StorageManager mStorageManager;
     @Mock private ArtdRefCache.Pin mArtdPin;
     @Mock private DexMetadataHelper.Injector mDexMetadataHelperInjector;
+    @Mock private Context mContext;
     private PackageState mPkgState1;
     private AndroidPackage mPkg1;
     private CheckedSecondaryDexInfo mPkg1SecondaryDexInfo1;
     private CheckedSecondaryDexInfo mPkg1SecondaryDexInfoNotFound;
     private Config mConfig;
     private DexMetadataHelper mDexMetadataHelper;
+    private ArgumentCaptor<BroadcastReceiver> mBroadcastReceiverCaptor;
 
     // True if the artifacts should be in dalvik-cache.
     @Parameter(0) public boolean mIsInDalvikCache;
@@ -173,6 +179,7 @@ public class ArtManagerLocalTest {
                 .when(mInjector.getArtFileManager())
                 .thenReturn(new ArtFileManager(mArtFileManagerInjector));
         lenient().when(mInjector.getDexMetadataHelper()).thenReturn(mDexMetadataHelper);
+        lenient().when(mInjector.getContext()).thenReturn(mContext);
 
         lenient().when(mArtFileManagerInjector.getArtd()).thenReturn(mArtd);
         lenient().when(mArtFileManagerInjector.getUserManager()).thenReturn(mUserManager);
@@ -186,6 +193,7 @@ public class ArtManagerLocalTest {
         lenient().when(SystemProperties.get(eq("pm.dexopt.install"))).thenReturn("speed-profile");
         lenient().when(SystemProperties.get(eq("pm.dexopt.bg-dexopt"))).thenReturn("speed-profile");
         lenient().when(SystemProperties.get(eq("pm.dexopt.first-boot"))).thenReturn("verify");
+        lenient().when(SystemProperties.get(eq("pm.dexopt.boot-after-ota"))).thenReturn("verify");
         lenient()
                 .when(SystemProperties.get(eq("pm.dexopt.boot-after-mainline-update")))
                 .thenReturn("verify");
@@ -269,6 +277,11 @@ public class ArtManagerLocalTest {
                 .when(mDexMetadataHelperInjector.openZipFile(any()))
                 .thenThrow(NoSuchFileException.class);
 
+        mBroadcastReceiverCaptor = ArgumentCaptor.forClass(BroadcastReceiver.class);
+        lenient()
+                .when(mContext.registerReceiver(mBroadcastReceiverCaptor.capture(), any()))
+                .thenReturn(mock(Intent.class));
+
         mArtManagerLocal = new ArtManagerLocal(mInjector);
     }
 
@@ -1330,6 +1343,56 @@ public class ArtManagerLocalTest {
         verify(mArtd, times(expectedGetProfileSizeCalls)).getProfileSize(any());
     }
 
+    @Test
+    public void testCommitPreRebootStagedFiles() throws Exception {
+        when(mSnapshot.getPackageStates()).thenReturn(Map.of(PKG_NAME_1, mPkgState1));
+
+        mArtManagerLocal.onBoot(ReasonMapping.REASON_BOOT_AFTER_OTA,
+                null /* progressCallbackExecutor */, null /* progressCallback */);
+
+        // It should commit files for primary dex files on boot.
+        verify(mArtd).commitPreRebootStagedFiles(
+                inAnyOrderDeepEquals(
+                        AidlUtils.buildArtifactsPathAsInput(
+                                "/somewhere/app/foo/base.apk", "arm64", mIsInDalvikCache),
+                        AidlUtils.buildArtifactsPathAsInput(
+                                "/somewhere/app/foo/base.apk", "arm", mIsInDalvikCache),
+                        AidlUtils.buildArtifactsPathAsInput(
+                                "/somewhere/app/foo/split_0.apk", "arm64", mIsInDalvikCache),
+                        AidlUtils.buildArtifactsPathAsInput(
+                                "/somewhere/app/foo/split_0.apk", "arm", mIsInDalvikCache)),
+                inAnyOrderDeepEquals(AidlUtils.toWritableProfilePath(
+                                             AidlUtils.buildProfilePathForPrimaryRefAsInput(
+                                                     PKG_NAME_1, "primary")),
+                        AidlUtils.toWritableProfilePath(
+                                AidlUtils.buildProfilePathForPrimaryRefAsInput(
+                                        PKG_NAME_1, "split_0.split"))));
+        verify(mArtd, times(1)).commitPreRebootStagedFiles(any(), any());
+
+        mArtManagerLocal.systemReady();
+
+        // It should not commit anything on system ready.
+        verify(mArtd, times(1)).commitPreRebootStagedFiles(any(), any());
+
+        mBroadcastReceiverCaptor.getValue().onReceive(mContext, mock(Intent.class));
+
+        // It should commit files for secondary dex files on boot complete.
+        verify(mArtd).commitPreRebootStagedFiles(
+                inAnyOrderDeepEquals(AidlUtils.buildArtifactsPathAsInput(
+                        "/data/user/0/foo/1.apk", "arm64", false /* isInDalvikCache */)),
+                inAnyOrderDeepEquals(AidlUtils.toWritableProfilePath(
+                        AidlUtils.buildProfilePathForSecondaryRefAsInput(
+                                "/data/user/0/foo/1.apk"))));
+        verify(mArtd, times(2)).commitPreRebootStagedFiles(any(), any());
+    }
+
+    @Test
+    public void testCommitPreRebootStagedFilesOnBootNotCalled() throws Exception {
+        mArtManagerLocal.systemReady();
+
+        verify(mContext, never()).registerReceiver(any(), any());
+    }
+
     private AndroidPackage createPackage(boolean multiSplit) {
         AndroidPackage pkg = mock(AndroidPackage.class);
 
diff --git a/libartservice/service/javatests/com/android/server/art/DexUseManagerTest.java b/libartservice/service/javatests/com/android/server/art/DexUseManagerTest.java
index 11810dfb35..ffc702f338 100644
--- a/libartservice/service/javatests/com/android/server/art/DexUseManagerTest.java
+++ b/libartservice/service/javatests/com/android/server/art/DexUseManagerTest.java
@@ -28,6 +28,7 @@ import static org.mockito.Mockito.argThat;
 import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.lenient;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.content.BroadcastReceiver;
@@ -97,6 +98,7 @@ public class DexUseManagerTest {
     @Mock private DexUseManagerLocal.Injector mInjector;
     @Mock private IArtd mArtd;
     @Mock private Context mContext;
+    @Mock private ArtManagerLocal mArtManagerLocal;
     private DexUseManagerLocal mDexUseManager;
     private String mCeDir;
     private String mDeDir;
@@ -167,6 +169,7 @@ public class DexUseManagerTest {
         lenient().when(mInjector.getContext()).thenReturn(mContext);
         lenient().when(mInjector.getAllPackageNames()).thenReturn(mPackageStates.keySet());
         lenient().when(mInjector.isPreReboot()).thenReturn(false);
+        lenient().when(mInjector.getArtManagerLocal()).thenReturn(mArtManagerLocal);
 
         mDexUseManager = new DexUseManagerLocal(mInjector);
         mDexUseManager.systemReady();
@@ -831,6 +834,11 @@ public class DexUseManagerTest {
                         true /* isUsedByOtherApps */, mDefaultFileVisibility));
     }
 
+    @Test
+    public void testSystemReady() {
+        verify(mArtManagerLocal).systemReady();
+    }
+
     @Test(expected = IllegalStateException.class)
     public void testPreRebootNoUpdate() throws Exception {
         when(mInjector.isPreReboot()).thenReturn(true);
diff --git a/runtime/jit/jit_code_cache.cc b/runtime/jit/jit_code_cache.cc
index 198ffcad25..378e6cba97 100644
--- a/runtime/jit/jit_code_cache.cc
+++ b/runtime/jit/jit_code_cache.cc
@@ -492,6 +492,7 @@ void JitCodeCache::FreeAllMethodHeaders(
   }
 
   {
+    MutexLock mu(Thread::Current(), *Locks::jit_lock_);
     ScopedCodeCacheWrite scc(private_region_);
     for (const OatQuickMethodHeader* method_header : method_headers) {
       FreeCodeAndData(method_header->GetCode());
@@ -506,18 +507,21 @@ void JitCodeCache::FreeAllMethodHeaders(
   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);
-      }
-    });
-    ForEachNativeDebugSymbol([&](const void* addr, size_t, const char* name) {
-      addr = AlignDown(addr, GetInstructionSetInstructionAlignment(kRuntimeISA));  // Thumb-bit.
-      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;
-    });
+    {
+      MutexLock mu(Thread::Current(), *Locks::jit_lock_);
+      VisitAllMethods([&](const void* addr, ArtMethod* method) {
+        if (!IsInZygoteExecSpace(addr)) {
+          CHECK(addr != nullptr && method != nullptr);
+          compiled_methods.emplace(addr, method);
+        }
+      });
+      ForEachNativeDebugSymbol([&](const void* addr, size_t, const char* name) {
+        addr = AlignDown(addr, GetInstructionSetInstructionAlignment(kRuntimeISA));  // Thumb-bit.
+        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 (auto it : compiled_methods) {
         CHECK_EQ(debug_info.count(it.first), 1u) << "Missing debug info";
@@ -535,66 +539,68 @@ 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.
   {
-    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;
+    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 = zombie_jni_code_.begin(); it != zombie_jni_code_.end();) {
-      if (alloc.ContainsUnsafe(*it)) {
-        it = zombie_jni_code_.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 = 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 = 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;
+        }
       }
     }
-    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);
+    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 = 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);
+    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;
       }
     }
   }
-  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);
 }
 
@@ -1107,53 +1113,56 @@ void JitCodeCache::RemoveUnmarkedCode(Thread* self) {
   ScopedTrace trace(__FUNCTION__);
   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 {
-      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());
+  {
+    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 {
+        OatQuickMethodHeader* header = OatQuickMethodHeader::FromCodePointer(code_ptr);
         method_headers.insert(header);
-        CHECK(ContainsPc(header));
-        VLOG(jit) << "JIT removed native code of" << method->PrettyMethod();
-        jni_stubs_map_.erase(stub);
+        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 {
-        stub->first.UpdateShorty(stub->second.GetMethods().front());
+        ++it;
       }
-      it = processed_zombie_jni_code_.erase(it);
-    } else {
-      ++it;
     }
   }
+
   FreeAllMethodHeaders(method_headers);
 }
 
diff --git a/runtime/jit/jit_code_cache.h b/runtime/jit/jit_code_cache.h
index 7de29d4024..5202c2c37e 100644
--- a/runtime/jit/jit_code_cache.h
+++ b/runtime/jit/jit_code_cache.h
@@ -452,7 +452,6 @@ class JitCodeCache {
 
   // Remove CHA dependents and underlying allocations for entries in `method_headers`.
   void FreeAllMethodHeaders(const std::unordered_set<OatQuickMethodHeader*>& method_headers)
-      REQUIRES(Locks::jit_lock_)
       REQUIRES(!Locks::cha_lock_);
 
   // Removes method from the cache. The caller must ensure that all threads