[res] Speed up AssetManager pointer locking

AssetManager needs to lock (promote()) apk asset weak pointers
to use them in pretty much any operation, and often mulitple
times for the same object.

This CL introduces a concept of an 'operation' in AssetManager,
and adds a cache for the locked assets that are retained until
the end of that operation. This way we only need to lock each
assets object exactly once, losing pretty much no performance.

Benchmarks all returned to the pre-memory-safe state values.

Bug: 280951693
Bug: 197260547
Bug: 276922628
Test: UTs + native benchmarks + java benchmarks + boot
Change-Id: I26ef88df4f6267b5e8b89f1588f2382db74030c0

1 files changed, 94 insertions, 35 deletions

diff --git a/libs/androidfw/AssetManager2.cpp b/libs/androidfw/AssetManager2.cpp
index 94dbfb5db1d7..fd5d07f22419 100644
--- a/libs/androidfw/AssetManager2.cpp
+++ b/libs/androidfw/AssetManager2.cpp
@@ -112,7 +112,15 @@ bool AssetManager2::SetApkAssets(std::initializer_list<ApkAssetsPtr> apk_assets,
 }
 
 void AssetManager2::BuildDynamicRefTable(ApkAssetsList apk_assets) {
-  apk_assets_.assign(apk_assets.begin(), apk_assets.end());
+  auto op = StartOperation();
+
+  apk_assets_.resize(apk_assets.size());
+  for (size_t i = 0; i != apk_assets.size(); ++i) {
+    apk_assets_[i].first = apk_assets[i];
+    // Let's populate the locked assets right away as we're going to need them here later.
+    apk_assets_[i].second = apk_assets[i];
+  }
+
   package_groups_.clear();
   package_ids_.fill(0xff);
 
@@ -124,7 +132,7 @@ void AssetManager2::BuildDynamicRefTable(ApkAssetsList apk_assets) {
   // Overlay resources are not directly referenced by an application so their resource ids
   // can change throughout the application's lifetime. Assign overlay package ids last.
   std::vector<const ApkAssets*> sorted_apk_assets;
-  sorted_apk_assets.reserve(apk_assets_.size());
+  sorted_apk_assets.reserve(apk_assets.size());
   for (auto& asset : apk_assets) {
     sorted_apk_assets.push_back(asset.get());
   }
@@ -250,9 +258,10 @@ void AssetManager2::BuildDynamicRefTable(ApkAssetsList apk_assets) {
 void AssetManager2::DumpToLog() const {
   LOG(INFO) << base::StringPrintf("AssetManager2(this=%p)", this);
 
+  auto op = StartOperation();
   std::string list;
-  for (const auto& apk_assets : apk_assets_) {
-    auto assets = apk_assets.promote();
+  for (size_t i = 0; i < apk_assets_.size(); ++i) {
+    const auto& assets = GetApkAssets(i);
     base::StringAppendF(&list, "%s,", assets ? assets->GetDebugName().c_str() : "nullptr");
   }
   LOG(INFO) << "ApkAssets: " << list;
@@ -290,7 +299,8 @@ const ResStringPool* AssetManager2::GetStringPoolForCookie(ApkAssetsCookie cooki
   if (cookie < 0 || static_cast<size_t>(cookie) >= apk_assets_.size()) {
     return nullptr;
   }
-  auto assets = apk_assets_[cookie].promote();
+  auto op = StartOperation();
+  const auto& assets = GetApkAssets(cookie);
   return assets ? assets->GetLoadedArsc()->GetStringPool() : nullptr;
 }
 
@@ -341,9 +351,10 @@ const std::unordered_map<std::string, std::string>*
 
 bool AssetManager2::GetOverlayablesToString(android::StringPiece package_name,
                                             std::string* out) const {
+  auto op = StartOperation();
   uint8_t package_id = 0U;
-  for (const auto& apk_assets : apk_assets_) {
-    auto assets = apk_assets.promote();
+  for (size_t i = 0; i != apk_assets_.size(); ++i) {
+    const auto& assets = GetApkAssets(i);
     if (!assets) {
       continue;
     }
@@ -400,10 +411,14 @@ bool AssetManager2::GetOverlayablesToString(android::StringPiece package_name,
 }
 
 bool AssetManager2::ContainsAllocatedTable() const {
-  return std::find_if(apk_assets_.begin(), apk_assets_.end(), [](auto&& assets_weak) {
-           auto assets = assets_weak.promote();
-           return assets && assets->IsTableAllocated();
-         }) != apk_assets_.end();
+  auto op = StartOperation();
+  for (size_t i = 0; i != apk_assets_.size(); ++i) {
+    const auto& assets = GetApkAssets(i);
+    if (assets && assets->IsTableAllocated()) {
+      return true;
+    }
+  }
+  return false;
 }
 
 void AssetManager2::SetConfiguration(const ResTable_config& configuration) {
@@ -428,8 +443,9 @@ std::set<AssetManager2::ApkAssetsPtr> AssetManager2::GetNonSystemOverlays() cons
     }
 
     if (!found_system_package) {
+      auto op = StartOperation();
       for (const ConfiguredOverlay& overlay : package_group.overlays_) {
-        if (auto asset = apk_assets_[overlay.cookie].promote()) {
+        if (const auto& asset = GetApkAssets(overlay.cookie)) {
           non_system_overlays.insert(std::move(asset));
         }
       }
@@ -442,6 +458,8 @@ std::set<AssetManager2::ApkAssetsPtr> AssetManager2::GetNonSystemOverlays() cons
 base::expected<std::set<ResTable_config>, IOError> AssetManager2::GetResourceConfigurations(
     bool exclude_system, bool exclude_mipmap) const {
   ATRACE_NAME("AssetManager::GetResourceConfigurations");
+  auto op = StartOperation();
+
   const auto non_system_overlays =
       exclude_system ? GetNonSystemOverlays() : std::set<ApkAssetsPtr>();
 
@@ -455,7 +473,7 @@ base::expected<std::set<ResTable_config>, IOError> AssetManager2::GetResourceCon
         }
         if (!non_system_overlays.empty()) {
           // Exclude overlays that target only system resources.
-          auto apk_assets = apk_assets_[package_group.cookies_[i]].promote();
+          const auto& apk_assets = GetApkAssets(package_group.cookies_[i]);
           if (apk_assets && apk_assets->IsOverlay() &&
               non_system_overlays.find(apk_assets) == non_system_overlays.end()) {
             continue;
@@ -475,6 +493,8 @@ base::expected<std::set<ResTable_config>, IOError> AssetManager2::GetResourceCon
 std::set<std::string> AssetManager2::GetResourceLocales(bool exclude_system,
                                                         bool merge_equivalent_languages) const {
   ATRACE_NAME("AssetManager::GetResourceLocales");
+  auto op = StartOperation();
+
   std::set<std::string> locales;
   const auto non_system_overlays =
       exclude_system ? GetNonSystemOverlays() : std::set<ApkAssetsPtr>();
@@ -488,7 +508,7 @@ std::set<std::string> AssetManager2::GetResourceLocales(bool exclude_system,
         }
         if (!non_system_overlays.empty()) {
           // Exclude overlays that target only system resources.
-          auto apk_assets = apk_assets_[package_group.cookies_[i]].promote();
+          const auto& apk_assets = GetApkAssets(package_group.cookies_[i]);
           if (apk_assets && apk_assets->IsOverlay() &&
               non_system_overlays.find(apk_assets) == non_system_overlays.end()) {
             continue;
@@ -516,13 +536,14 @@ std::unique_ptr<Asset> AssetManager2::Open(const std::string& filename, ApkAsset
 
 std::unique_ptr<AssetDir> AssetManager2::OpenDir(const std::string& dirname) const {
   ATRACE_NAME("AssetManager::OpenDir");
+  auto op = StartOperation();
 
   std::string full_path = "assets/" + dirname;
   auto files = util::make_unique<SortedVector<AssetDir::FileInfo>>();
 
   // Start from the back.
-  for (auto iter = apk_assets_.rbegin(); iter != apk_assets_.rend(); ++iter) {
-    auto apk_assets = iter->promote();
+  for (size_t i = apk_assets_.size(); i > 0; --i) {
+    const auto& apk_assets = GetApkAssets(i - 1);
     if (!apk_assets || apk_assets->IsOverlay()) {
       continue;
     }
@@ -551,8 +572,9 @@ std::unique_ptr<AssetDir> AssetManager2::OpenDir(const std::string& dirname) con
 std::unique_ptr<Asset> AssetManager2::OpenNonAsset(const std::string& filename,
                                                    Asset::AccessMode mode,
                                                    ApkAssetsCookie* out_cookie) const {
-  for (int32_t i = apk_assets_.size() - 1; i >= 0; i--) {
-    const auto assets = apk_assets_[i].promote();
+  auto op = StartOperation();
+  for (size_t i = apk_assets_.size(); i > 0; i--) {
+    const auto& assets = GetApkAssets(i - 1);
     // Prevent RRO from modifying assets and other entries accessed by file
     // path. Explicitly asking for a path in a given package (denoted by a
     // cookie) is still OK.
@@ -581,7 +603,8 @@ std::unique_ptr<Asset> AssetManager2::OpenNonAsset(const std::string& filename,
   if (cookie < 0 || static_cast<size_t>(cookie) >= apk_assets_.size()) {
     return {};
   }
-  auto assets = apk_assets_[cookie].promote();
+  auto op = StartOperation();
+  const auto& assets = GetApkAssets(cookie);
   return assets ? assets->GetAssetsProvider()->Open(filename, mode) : nullptr;
 }
 
@@ -595,6 +618,8 @@ base::expected<FindEntryResult, NullOrIOError> AssetManager2::FindEntry(
     last_resolution_.resid = resid;
   }
 
+  auto op = StartOperation();
+
   // Might use this if density_override != 0.
   ResTable_config density_override_config;
 
@@ -631,7 +656,7 @@ base::expected<FindEntryResult, NullOrIOError> AssetManager2::FindEntry(
 
   bool overlaid = false;
   if (!stop_at_first_match && !ignore_configuration) {
-    auto assets = apk_assets_[result->cookie].promote();
+    const auto& assets = GetApkAssets(result->cookie);
     if (!assets) {
       ALOGE("Found expired ApkAssets #%d for resource ID 0x%08x.", result->cookie, resid);
       return base::unexpected(std::nullopt);
@@ -892,8 +917,10 @@ std::string AssetManager2::GetLastResourceResolution() const {
     return {};
   }
 
+  auto op = StartOperation();
+
   const uint32_t resid = last_resolution_.resid;
-  auto assets = apk_assets_[cookie].promote();
+  const auto& assets = GetApkAssets(cookie);
   const auto package =
       assets ? assets->GetLoadedArsc()->GetPackageById(get_package_id(resid)) : nullptr;
 
@@ -926,7 +953,7 @@ std::string AssetManager2::GetLastResourceResolution() const {
       continue;
     }
     const auto prefix = kStepStrings[int(step.type) - int(Resolution::Step::Type::INITIAL)];
-    auto assets = apk_assets_[step.cookie].promote();
+    const auto& assets = GetApkAssets(step.cookie);
     log_stream << "\n\t" << prefix << ": " << (assets ? assets->GetDebugName() : "<null>")
                << " #" << step.cookie;
     if (!step.config_name.isEmpty()) {
@@ -1458,6 +1485,37 @@ void AssetManager2::ForEachPackage(base::function_ref<bool(const std::string&, u
   }
 }
 
+AssetManager2::ScopedOperation AssetManager2::StartOperation() const {
+  ++number_of_running_scoped_operations_;
+  return ScopedOperation(*this);
+}
+
+void AssetManager2::FinishOperation() const {
+  if (number_of_running_scoped_operations_ < 1) {
+    ALOGW("Invalid FinishOperation() call when there's none happening");
+    return;
+  }
+  if (--number_of_running_scoped_operations_ == 0) {
+    for (auto&& [_, assets] : apk_assets_) {
+      assets.clear();
+    }
+  }
+}
+
+const AssetManager2::ApkAssetsPtr& AssetManager2::GetApkAssets(ApkAssetsCookie cookie) const {
+  DCHECK(number_of_running_scoped_operations_ > 0) << "Must have an operation running";
+
+  if (cookie < 0 || cookie >= apk_assets_.size()) {
+    static const ApkAssetsPtr empty{};
+    return empty;
+  }
+  auto& [wptr, res] = apk_assets_[cookie];
+  if (!res) {
+    res = wptr.promote();
+  }
+  return res;
+}
+
 Theme::Theme(AssetManager2* asset_manager) : asset_manager_(asset_manager) {
 }
 
@@ -1590,22 +1648,16 @@ base::expected<std::monostate, IOError> Theme::SetTo(const Theme& source) {
     using SourceToDestinationRuntimePackageMap = std::unordered_map<int, int>;
     std::unordered_map<ApkAssetsCookie, SourceToDestinationRuntimePackageMap> src_asset_cookie_id_map;
 
-    // Determine which ApkAssets are loaded in both theme AssetManagers.
-    const auto& src_assets = source.asset_manager_->GetApkAssets();
-    const auto& dest_assets = asset_manager_->GetApkAssets();
-    std::vector<AssetManager2::ApkAssetsPtr> promoted_src_assets;
-    promoted_src_assets.reserve(src_assets.size());
-    for (const auto& src_asset : src_assets) {
-      promoted_src_assets.emplace_back(src_asset.promote());
-    }
+    auto op_src = source.asset_manager_->StartOperation();
+    auto op_dst = asset_manager_->StartOperation();
 
-    for (size_t j = 0; j < dest_assets.size(); j++) {
-      auto dest_asset = dest_assets[j].promote();
-      if (!dest_asset) {
+    for (size_t i = 0; i < source.asset_manager_->GetApkAssetsCount(); i++) {
+      const auto& src_asset = source.asset_manager_->GetApkAssets(i);
+      if (!src_asset) {
         continue;
       }
-      for (size_t i = 0; i < promoted_src_assets.size(); i++) {
-        const auto& src_asset = promoted_src_assets[i];
+      for (int j = 0; j < asset_manager_->GetApkAssetsCount(); j++) {
+        const auto& dest_asset = asset_manager_->GetApkAssets(j);
         if (src_asset != dest_asset) {
           // ResourcesManager caches and reuses ApkAssets when the same apk must be present in
           // multiple AssetManagers. Two ApkAssets point to the same version of the same resources
@@ -1731,4 +1783,11 @@ void Theme::Dump() const {
   }
 }
 
+AssetManager2::ScopedOperation::ScopedOperation(const AssetManager2& am) : am_(am) {
+}
+
+AssetManager2::ScopedOperation::~ScopedOperation() {
+  am_.FinishOperation();
+}
+
 }  // namespace android