1 files changed, 95 insertions, 42 deletions

diff --git a/libs/androidfw/AssetManager2.cpp b/libs/androidfw/AssetManager2.cpp
index 3a7c91ef5b08..d33b592bddc6 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);
@@ -873,13 +898,7 @@ base::expected<FindEntryResult, NullOrIOError> AssetManager2::FindEntryInternal(
 }
 
 void AssetManager2::ResetResourceResolution() const {
-  last_resolution_.cookie = kInvalidCookie;
-  last_resolution_.resid = 0;
-  last_resolution_.steps.clear();
-  last_resolution_.type_string_ref = StringPoolRef();
-  last_resolution_.entry_string_ref = StringPoolRef();
-  last_resolution_.best_config_name.clear();
-  last_resolution_.best_package_name.clear();
+  last_resolution_ = Resolution{};
 }
 
 void AssetManager2::SetResourceResolutionLoggingEnabled(bool enabled) {
@@ -901,8 +920,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;
 
@@ -934,7 +955,7 @@ std::string AssetManager2::GetLastResourceResolution() const {
     if (prefix == kStepStrings.end()) {
       continue;
     }
-    auto assets = apk_assets_[step.cookie].promote();
+    const auto& assets = GetApkAssets(step.cookie);
     log_stream << "\n\t" << prefix->second << ": " 
                << (assets ? assets->GetDebugName() : "<null>") << " #" << step.cookie;
     if (!step.config_name.isEmpty()) {
@@ -1466,6 +1487,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) {
 }
 
@@ -1598,22 +1650,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
@@ -1739,4 +1785,11 @@ void Theme::Dump() const {
   }
 }
 
+AssetManager2::ScopedOperation::ScopedOperation(const AssetManager2& am) : am_(am) {
+}
+
+AssetManager2::ScopedOperation::~ScopedOperation() {
+  am_.FinishOperation();
+}
+
 }  // namespace android