When generating a runtime image, load classes from reference profile.

This makes sure that we generate the best image.

Also add a couple of ScopedTrace to profile runtime app image
generation.

Test: test.py
Test: locally start an app, see image containing all classes.
Bug: 260557058
Change-Id: I1c2690c54a4c815ac575cada476f48948bd63131

5 files changed, 106 insertions, 2 deletions

diff --git a/libprofile/profile/profile_compilation_info.cc b/libprofile/profile/profile_compilation_info.cc
index b1371db16d..6b1aa0e87c 100644
--- a/libprofile/profile/profile_compilation_info.cc
+++ b/libprofile/profile/profile_compilation_info.cc
@@ -2180,6 +2180,16 @@ bool ProfileCompilationInfo::GetClassesAndMethods(
   return true;
 }
 
+const ArenaSet<dex::TypeIndex>* ProfileCompilationInfo::GetClasses(
+    const DexFile& dex_file,
+    const ProfileSampleAnnotation& annotation) const {
+  const DexFileData* dex_data = FindDexDataUsingAnnotations(&dex_file, annotation);
+  if (dex_data == nullptr) {
+    return nullptr;
+  }
+  return &dex_data->class_set;
+}
+
 bool ProfileCompilationInfo::SameVersion(const ProfileCompilationInfo& other) const {
   return memcmp(version_, other.version_, kProfileVersionSize) == 0;
 }
diff --git a/libprofile/profile/profile_compilation_info.h b/libprofile/profile/profile_compilation_info.h
index c74f478383..a6fd32ae22 100644
--- a/libprofile/profile/profile_compilation_info.h
+++ b/libprofile/profile/profile_compilation_info.h
@@ -593,6 +593,10 @@ class ProfileCompilationInfo {
       /*out*/std::set<uint16_t>* post_startup_method_method_set,
       const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone) const;
 
+  const ArenaSet<dex::TypeIndex>* GetClasses(
+      const DexFile& dex_file,
+      const ProfileSampleAnnotation& annotation = ProfileSampleAnnotation::kNone) const;
+
   // Returns true iff both profiles have the same version.
   bool SameVersion(const ProfileCompilationInfo& other) const;
 
diff --git a/runtime/app_info.cc b/runtime/app_info.cc
index 2dbbbf6a90..3fef5651b0 100644
--- a/runtime/app_info.cc
+++ b/runtime/app_info.cc
@@ -143,4 +143,17 @@ std::ostream& operator<<(std::ostream& os, AppInfo& rhs) {
   return os;
 }
 
+std::string AppInfo::GetPrimaryApkReferenceProfile() {
+  MutexLock mu(Thread::Current(), update_mutex_);
+
+  for (const auto& it : registered_code_locations_) {
+    const CodeLocationInfo& cli = it.second;
+    if (cli.code_type == CodeType::kPrimaryApk) {
+      return cli.ref_profile_path.value_or("");
+    }
+  }
+  return "";
+}
+
+
 }  // namespace art
diff --git a/runtime/app_info.h b/runtime/app_info.h
index 43e2ef320b..0b8132aa6f 100644
--- a/runtime/app_info.h
+++ b/runtime/app_info.h
@@ -68,15 +68,22 @@ class AppInfo {
                           const std::string& compilation_reason,
                           const std::string& odex_status);
 
-  // Extracts the optimization status of the primary apk into the given arguments.
+  // Extracts the optimization status of the primary APK into the given arguments.
   // If there are multiple primary APKs registed via RegisterAppInfo, the method
   // will assign the status of the first APK, sorted by the location name.
   //
-  // Assigns "unknown" if there is no primary apk or the optimization status was
+  // Assigns "unknown" if there is no primary APK or the optimization status was
   // not set via RegisterOdexStatus,
   void GetPrimaryApkOptimizationStatus(std::string* out_compiler_filter,
                                        std::string* out_compilation_reason);
 
+  // Returns the reference profile path of the primary APK.
+  // If there are multiple primary APKs registed via RegisterAppInfo, the method
+  // will return the profile of the first APK, sorted by the location name.
+  //
+  // Returns an empty string if there is no primary APK.
+  std::string GetPrimaryApkReferenceProfile();
+
   // Whether we've received a call to RegisterAppInfo.
   bool HasRegisteredAppInfo();
 
diff --git a/runtime/runtime_image.cc b/runtime/runtime_image.cc
index 6ad497548b..433852d12a 100644
--- a/runtime/runtime_image.cc
+++ b/runtime/runtime_image.cc
@@ -29,6 +29,7 @@
 #include "base/bit_utils.h"
 #include "base/file_utils.h"
 #include "base/length_prefixed_array.h"
+#include "base/scoped_flock.h"
 #include "base/stl_util.h"
 #include "base/systrace.h"
 #include "base/unix_file/fd_file.h"
@@ -45,6 +46,7 @@
 #include "mirror/object_array.h"
 #include "mirror/string-inl.h"
 #include "oat.h"
+#include "profile/profile_compilation_info.h"
 #include "scoped_thread_state_change-inl.h"
 #include "vdex_file.h"
 
@@ -512,6 +514,7 @@ class RuntimeImageHelper {
 
   void EmitClasses(Thread* self, Handle<mirror::ObjectArray<mirror::Object>> dex_cache_array)
       REQUIRES_SHARED(Locks::mutator_lock_) {
+    ScopedTrace trace("Emit strings and classes");
     ArenaAllocator allocator(Runtime::Current()->GetArenaPool());
     ArenaSet<const DexFile*> dex_files(allocator.Adapter());
     for (int32_t i = 0; i < dex_cache_array->GetLength(); ++i) {
@@ -686,6 +689,7 @@ class RuntimeImageHelper {
   }
 
   void RelocateNativePointers() {
+    ScopedTrace relocate_native_pointers("Relocate native pointers");
     ScopedObjectAccess soa(Thread::Current());
     NativePointerVisitor visitor(this);
     for (auto it : classes_) {
@@ -849,7 +853,67 @@ class RuntimeImageHelper {
     return native_relocations_.find(ptr) != native_relocations_.end();
   }
 
+
+  static void LoadClassesFromReferenceProfile(
+      Thread* self,
+      const dchecked_vector<Handle<mirror::DexCache>>& dex_caches)
+          REQUIRES_SHARED(Locks::mutator_lock_) {
+    AppInfo* app_info = Runtime::Current()->GetAppInfo();
+    std::string profile_file = app_info->GetPrimaryApkReferenceProfile();
+
+    if (profile_file.empty()) {
+      return;
+    }
+
+    // Lock the file, it could be concurrently updated by the system. Don't block
+    // as this is app startup sensitive.
+    std::string error;
+    ScopedFlock profile =
+        LockedFile::Open(profile_file.c_str(), O_RDONLY, /*block=*/false, &error);
+
+    if (profile == nullptr) {
+      LOG(DEBUG) << "Couldn't lock the profile file " << profile_file << ": " << error;
+      return;
+    }
+
+    ProfileCompilationInfo profile_info(/* for_boot_image= */ false);
+
+    if (!profile_info.Load(profile->Fd())) {
+      LOG(DEBUG) << "Could not load profile file";
+      return;
+    }
+
+    StackHandleScope<1> hs(self);
+    Handle<mirror::ClassLoader> class_loader =
+        hs.NewHandle<mirror::ClassLoader>(dex_caches[0]->GetClassLoader());
+    ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
+    ScopedTrace loading_classes("Loading classes from profile");
+    for (auto dex_cache : dex_caches) {
+      const DexFile* dex_file = dex_cache->GetDexFile();
+      const ArenaSet<dex::TypeIndex>* class_types = profile_info.GetClasses(*dex_file);
+      if (class_types == nullptr) {
+        // This means the profile file did not reference the dex file, which is the case
+        // if there's no classes and methods of that dex file in the profile.
+        continue;
+      }
+
+      for (dex::TypeIndex idx : *class_types) {
+        // The index is greater or equal to NumTypeIds if the type is an extra
+        // descriptor, not referenced by the dex file.
+        if (idx.index_ < dex_file->NumTypeIds()) {
+          ObjPtr<mirror::Class> klass = class_linker->ResolveType(idx, dex_cache, class_loader);
+          if (klass == nullptr) {
+            self->ClearException();
+            LOG(DEBUG) << "Failed to preload " << dex_file->PrettyType(idx);
+            continue;
+          }
+        }
+      }
+    }
+  }
+
   bool WriteObjects(std::string* error_msg) {
+    ScopedTrace write_objects("Writing objects");
     ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
     ScopedObjectAccess soa(Thread::Current());
     VariableSizedHandleScope handles(soa.Self());
@@ -899,6 +963,11 @@ class RuntimeImageHelper {
       }
     }
 
+    // If classes referenced in the reference profile are not loaded, preload
+    // them. This makes sure we generate a good runtime app image, even if this
+    // current app run did not load all startup classes.
+    LoadClassesFromReferenceProfile(soa.Self(), dex_caches);
+
     // We store the checksums of the dex files used at runtime. These can be
     // different compared to the vdex checksums due to compact dex.
     std::vector<uint32_t> checksums(number_of_dex_files);
@@ -1228,6 +1297,7 @@ class RuntimeImageHelper {
                      dchecked_vector<Handle<mirror::DexCache>>& dex_caches,
                      VariableSizedHandleScope& handles)
       REQUIRES_SHARED(Locks::mutator_lock_) {
+    ScopedTrace trace("Find dex caches");
     DCHECK(dex_caches.empty());
     // Collect all dex caches.
     ClassLinker* class_linker = Runtime::Current()->GetClassLinker();