Enable profile key updates via profman

Allow profman to update the profile key in an existing profile based on
new dex location.

This is needed in order to support profiles for dex files that might move
from their original profiling location (e.g. during install)

The matching [profile key <-> dex_file] is done based on the dex checksum
and the number of methods ids. If neither is a match then the profile key
is not updated. If the new profile key would collide with an existing key
(for a different dex) then the update will fail.

Test: profile_compilation_info_test
Bug: 30934496

Change-Id: Ic696b3f6fe9da2007421bf044d58a21c90fd9ee7
diff --git a/profman/profman.cc b/profman/profman.cc
index 9f3e3b6..ffc3c01 100644
--- a/profman/profman.cc
+++ b/profman/profman.cc
@@ -149,6 +149,10 @@
   UsageError("  --boot-image-sampled-method-threshold=<value>: minimum number of profiles a");
   UsageError("      non-hot method needs to be in order to be hot in the output profile. The");
   UsageError("      default is max int.");
+  UsageError("  --copy-and-update-profile-key: if present, profman will copy the profile from");
+  UsageError("      the file passed with --profile-fd(file) to the profile passed with");
+  UsageError("      --reference-profile-fd(file) and update at the same time the profile-key");
+  UsageError("      of entries corresponding to the apks passed with --apk(-fd).");
   UsageError("");
 
   exit(EXIT_FAILURE);
@@ -186,7 +190,8 @@
       test_profile_method_percerntage_(kDefaultTestProfileMethodPercentage),
       test_profile_class_percentage_(kDefaultTestProfileClassPercentage),
       test_profile_seed_(NanoTime()),
-      start_ns_(NanoTime()) {}
+      start_ns_(NanoTime()),
+      copy_and_update_profile_key_(false) {}
 
   ~ProfMan() {
     LogCompletionTime();
@@ -302,11 +307,13 @@
             "should only be used together");
     }
     ProfileAssistant::ProcessingResult result;
+
     if (profile_files_.empty()) {
       // The file doesn't need to be flushed here (ProcessProfiles will do it)
       // so don't check the usage.
       File file(reference_profile_file_fd_, false);
-      result = ProfileAssistant::ProcessProfiles(profile_files_fd_, reference_profile_file_fd_);
+      result = ProfileAssistant::ProcessProfiles(profile_files_fd_,
+                                                 reference_profile_file_fd_);
       CloseAllFds(profile_files_fd_, "profile_files_fd_");
     } else {
       result = ProfileAssistant::ProcessProfiles(profile_files_, reference_profile_file_);
@@ -314,7 +321,7 @@
     return result;
   }
 
-  void OpenApkFilesFromLocations(std::vector<std::unique_ptr<const DexFile>>* dex_files) {
+  void OpenApkFilesFromLocations(std::vector<std::unique_ptr<const DexFile>>* dex_files) const {
     bool use_apk_fd_list = !apks_fd_.empty();
     if (use_apk_fd_list) {
       // Get the APKs from the collection of FDs.
@@ -1070,6 +1077,42 @@
     return !test_profile_.empty();
   }
 
+  bool ShouldCopyAndUpdateProfileKey() const {
+    return copy_and_update_profile_key_;
+  }
+
+  bool CopyAndUpdateProfileKey() const {
+    // Validate that at least one profile file was passed, as well as a reference profile.
+    if (!(profile_files_.size() == 1 ^ profile_files_fd_.size() == 1)) {
+      Usage("Only one profile file should be specified.");
+    }
+    if (reference_profile_file_.empty() && !FdIsValid(reference_profile_file_fd_)) {
+      Usage("No reference profile file specified.");
+    }
+
+    if (apk_files_.empty() && apks_fd_.empty()) {
+      Usage("No apk files specified");
+    }
+
+    bool use_fds = profile_files_fd_.size() == 1;
+
+    ProfileCompilationInfo profile;
+    // Do not clear if invalid. The input might be an archive.
+    if (profile.Load(profile_files_[0], /*clear_if_invalid*/ false)) {
+      // Open the dex files to look up classes and methods.
+      std::vector<std::unique_ptr<const DexFile>> dex_files;
+      OpenApkFilesFromLocations(&dex_files);
+      if (!profile.UpdateProfileKeys(dex_files)) {
+        return false;
+      }
+      return use_fds
+        ? profile.Save(reference_profile_file_fd_)
+        : profile.Save(reference_profile_file_, /*bytes_written*/ nullptr);
+    } else {
+      return false;
+    }
+  }
+
  private:
   static void ParseFdForCollection(const StringPiece& option,
                                    const char* arg_name,
@@ -1114,6 +1157,7 @@
   uint16_t test_profile_class_percentage_;
   uint32_t test_profile_seed_;
   uint64_t start_ns_;
+  bool copy_and_update_profile_key_;
 };
 
 // See ProfileAssistant::ProcessingResult for return codes.
diff --git a/runtime/jit/profile_compilation_info.cc b/runtime/jit/profile_compilation_info.cc
index 33fa0d6..4bf2895 100644
--- a/runtime/jit/profile_compilation_info.cc
+++ b/runtime/jit/profile_compilation_info.cc
@@ -2030,4 +2030,28 @@
   return memcmp(buffer, kProfileMagic, byte_count) == 0;
 }
 
+bool ProfileCompilationInfo::UpdateProfileKeys(
+      const std::vector<std::unique_ptr<const DexFile>>& dex_files) {
+  for (const std::unique_ptr<const DexFile>& dex_file : dex_files) {
+    for (DexFileData* dex_data : info_) {
+      if (dex_data->checksum == dex_file->GetLocationChecksum()
+          && dex_data->num_method_ids == dex_file->NumMethodIds()) {
+        std::string new_profile_key = GetProfileDexFileKey(dex_file->GetLocation());
+        if (dex_data->profile_key != new_profile_key) {
+          if (profile_key_map_.find(new_profile_key) != profile_key_map_.end()) {
+            // We can't update the key if the new key belongs to a different dex file.
+            LOG(ERROR) << "Cannot update profile key to " << new_profile_key
+                << " because the new key belongs to another dex file.";
+            return false;
+          }
+          profile_key_map_.erase(dex_data->profile_key);
+          profile_key_map_.Put(new_profile_key, dex_data->profile_index);
+          dex_data->profile_key = new_profile_key;
+        }
+      }
+    }
+  }
+  return true;
+}
+
 }  // namespace art
diff --git a/runtime/jit/profile_compilation_info.h b/runtime/jit/profile_compilation_info.h
index 29a4c11..350ce9e 100644
--- a/runtime/jit/profile_compilation_info.h
+++ b/runtime/jit/profile_compilation_info.h
@@ -416,6 +416,17 @@
   // Return true if the fd points to a profile file.
   bool IsProfileFile(int fd);
 
+  // Update the profile keys corresponding to the given dex files based on their current paths.
+  // This method allows fix-ups in the profile for dex files that might have been renamed.
+  // The new profile key will be constructed based on the current dex location.
+  //
+  // The matching [profile key <-> dex_file] is done based on the dex checksum and the number of
+  // methods ids. If neither is a match then the profile key is not updated.
+  //
+  // If the new profile key would collide with an existing key (for a different dex)
+  // the method returns false. Otherwise it returns true.
+  bool UpdateProfileKeys(const std::vector<std::unique_ptr<const DexFile>>& dex_files);
+
  private:
   enum ProfileLoadStatus {
     kProfileLoadWouldOverwiteData,
diff --git a/runtime/jit/profile_compilation_info_test.cc b/runtime/jit/profile_compilation_info_test.cc
index 6ce9bcb..bb5bc88 100644
--- a/runtime/jit/profile_compilation_info_test.cc
+++ b/runtime/jit/profile_compilation_info_test.cc
@@ -22,6 +22,7 @@
 #include "class_linker-inl.h"
 #include "common_runtime_test.h"
 #include "dex/dex_file.h"
+#include "dex/dex_file_loader.h"
 #include "handle_scope-inl.h"
 #include "jit/profile_compilation_info.h"
 #include "linear_alloc.h"
@@ -1040,4 +1041,89 @@
   ASSERT_FALSE(loaded_info.Load(GetFd(zip)));
 }
 
+TEST_F(ProfileCompilationInfoTest, UpdateProfileKeyOk) {
+  std::vector<std::unique_ptr<const DexFile>> dex_files = OpenTestDexFiles("MultiDex");
+
+  ProfileCompilationInfo info;
+  for (const std::unique_ptr<const DexFile>& dex : dex_files) {
+    // Create the profile with a different location so that we can update it to the
+    // real dex location later.
+    std::string base_location = DexFileLoader::GetBaseLocation(dex->GetLocation());
+    std::string multidex_suffix = DexFileLoader::GetMultiDexSuffix(dex->GetLocation());
+    std::string old_name = base_location + "-old" + multidex_suffix;
+    info.AddMethodIndex(Hotness::kFlagHot,
+                        old_name,
+                        dex->GetLocationChecksum(),
+                        /* method_idx */ 0,
+                        dex->NumMethodIds());
+  }
+
+  // Update the profile keys based on the original dex files
+  ASSERT_TRUE(info.UpdateProfileKeys(dex_files));
+
+  // Verify that we find the methods when searched with the original dex files.
+  for (const std::unique_ptr<const DexFile>& dex : dex_files) {
+    std::unique_ptr<ProfileCompilationInfo::OfflineProfileMethodInfo> loaded_pmi =
+        info.GetMethod(dex->GetLocation(), dex->GetLocationChecksum(), /* method_idx */ 0);
+    ASSERT_TRUE(loaded_pmi != nullptr);
+  }
+}
+
+TEST_F(ProfileCompilationInfoTest, UpdateProfileKeyOkButNoUpdate) {
+  std::vector<std::unique_ptr<const DexFile>> dex_files = OpenTestDexFiles("MultiDex");
+
+  ProfileCompilationInfo info;
+  info.AddMethodIndex(Hotness::kFlagHot,
+                      "my.app",
+                      /* checksum */ 123,
+                      /* method_idx */ 0,
+                      /* num_method_ids */ 10);
+
+  // Update the profile keys based on the original dex files
+  ASSERT_TRUE(info.UpdateProfileKeys(dex_files));
+
+  // Verify that we did not perform any update and that we cannot find anything with the new
+  // location.
+  for (const std::unique_ptr<const DexFile>& dex : dex_files) {
+    std::unique_ptr<ProfileCompilationInfo::OfflineProfileMethodInfo> loaded_pmi =
+        info.GetMethod(dex->GetLocation(), dex->GetLocationChecksum(), /* method_idx */ 0);
+    ASSERT_TRUE(loaded_pmi == nullptr);
+  }
+
+  // Verify that we can find the original entry.
+  std::unique_ptr<ProfileCompilationInfo::OfflineProfileMethodInfo> loaded_pmi =
+        info.GetMethod("my.app", /* checksum */ 123, /* method_idx */ 0);
+  ASSERT_TRUE(loaded_pmi != nullptr);
+}
+
+TEST_F(ProfileCompilationInfoTest, UpdateProfileKeyFail) {
+  std::vector<std::unique_ptr<const DexFile>> dex_files = OpenTestDexFiles("MultiDex");
+
+
+  ProfileCompilationInfo info;
+  // Add all dex
+  for (const std::unique_ptr<const DexFile>& dex : dex_files) {
+    // Create the profile with a different location so that we can update it to the
+    // real dex location later.
+    std::string base_location = DexFileLoader::GetBaseLocation(dex->GetLocation());
+    std::string multidex_suffix = DexFileLoader::GetMultiDexSuffix(dex->GetLocation());
+    std::string old_name = base_location + "-old" + multidex_suffix;
+    info.AddMethodIndex(Hotness::kFlagHot,
+                        old_name,
+                        dex->GetLocationChecksum(),
+                        /* method_idx */ 0,
+                        dex->NumMethodIds());
+  }
+
+  // Add a method index using the location we want to rename to.
+  // This will cause the rename to fail because an existing entry would already have that name.
+  info.AddMethodIndex(Hotness::kFlagHot,
+                      dex_files[0]->GetLocation(),
+                      /* checksum */ 123,
+                      /* method_idx */ 0,
+                      dex_files[0]->NumMethodIds());
+
+  ASSERT_FALSE(info.UpdateProfileKeys(dex_files));
+}
+
 }  // namespace art