Add bug fix flag to get ART managed file stats only when the stats collection is requested.

Flag:  android.app.usage.get_app_art_managed_bytes

Test: manually run StorageStats cts tests.

Bug: 395548922

Change-Id: Ie785b76b959a19a45e29e470ec02535a3b890f8a

4 files changed, 139 insertions, 9 deletions

diff --git a/core/java/android/app/usage/IStorageStatsManager.aidl b/core/java/android/app/usage/IStorageStatsManager.aidl
index b5036da33a95..390300360c8f 100644
--- a/core/java/android/app/usage/IStorageStatsManager.aidl
+++ b/core/java/android/app/usage/IStorageStatsManager.aidl
@@ -30,6 +30,7 @@ interface IStorageStatsManager {
     long getCacheBytes(String volumeUuid, String callingPackage);
     long getCacheQuotaBytes(String volumeUuid, int uid, String callingPackage);
     StorageStats queryStatsForPackage(String volumeUuid, String packageName, int userId, String callingPackage);
+    StorageStats queryArtManagedStats(String packageName, int userId, int uid);
     StorageStats queryStatsForUid(String volumeUuid, int uid, String callingPackage);
     StorageStats queryStatsForUser(String volumeUuid, int userId, String callingPackage);
     ExternalStorageStats queryExternalStatsForUser(String volumeUuid, int userId, String callingPackage);
diff --git a/core/java/android/app/usage/StorageStats.java b/core/java/android/app/usage/StorageStats.java
index 7bfaef4e9857..d01b423d67f2 100644
--- a/core/java/android/app/usage/StorageStats.java
+++ b/core/java/android/app/usage/StorageStats.java
@@ -22,7 +22,10 @@ import android.annotation.IntDef;
 import android.content.Context;
 import android.os.Parcel;
 import android.os.Parcelable;
+import android.os.RemoteException;
+import android.os.ServiceManager;
 import android.os.UserHandle;
+import android.util.Log;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -34,6 +37,9 @@ import java.lang.annotation.RetentionPolicy;
  * @see StorageStatsManager
  */
 public final class StorageStats implements Parcelable {
+    /** @hide */ public String packageName;
+    /** @hide */ public int userHandle;
+    /** @hide */ public int uid;
     /** @hide */ public long codeBytes;
     /** @hide */ public long dataBytes;
     /** @hide */ public long cacheBytes;
@@ -130,6 +136,14 @@ public final class StorageStats implements Parcelable {
     @Retention(RetentionPolicy.SOURCE)
     public @interface AppDataType {}
 
+    private static final String TAG = "StorageStats";
+
+    /**
+     * artStatsFetched is only applicable when
+     * Flags.getAppArtManagedBytes() is true;
+     */
+    private boolean artStatsFetched;
+
     /**
      * Return the size of app. This includes {@code APK} files, optimized
      * compiler output, and unpacked native libraries.
@@ -157,9 +171,9 @@ public final class StorageStats implements Parcelable {
     @FlaggedApi(Flags.FLAG_GET_APP_BYTES_BY_DATA_TYPE_API)
     public long getAppBytesByDataType(@AppDataType int dataType) {
         switch (dataType) {
-          case APP_DATA_TYPE_FILE_TYPE_DEXOPT_ARTIFACT: return dexoptBytes;
-          case APP_DATA_TYPE_FILE_TYPE_REFERENCE_PROFILE: return refProfBytes;
-          case APP_DATA_TYPE_FILE_TYPE_CURRENT_PROFILE: return curProfBytes;
+          case APP_DATA_TYPE_FILE_TYPE_DEXOPT_ARTIFACT: return getDexoptBytes();
+          case APP_DATA_TYPE_FILE_TYPE_REFERENCE_PROFILE: return getRefProfBytes();
+          case APP_DATA_TYPE_FILE_TYPE_CURRENT_PROFILE: return getCurProfBytes();
           case APP_DATA_TYPE_FILE_TYPE_APK: return apkBytes;
           case APP_DATA_TYPE_LIB: return libBytes;
           case APP_DATA_TYPE_FILE_TYPE_DM: return dmBytes;
@@ -215,6 +229,9 @@ public final class StorageStats implements Parcelable {
 
     /** {@hide} */
     public StorageStats(Parcel in) {
+        this.packageName = in.readString8();
+        this.userHandle = in.readInt();
+        this.uid = in.readInt();
         this.codeBytes = in.readLong();
         this.dataBytes = in.readLong();
         this.cacheBytes = in.readLong();
@@ -234,6 +251,9 @@ public final class StorageStats implements Parcelable {
 
     @Override
     public void writeToParcel(Parcel dest, int flags) {
+        dest.writeString8(packageName);
+        dest.writeInt(userHandle);
+        dest.writeInt(uid);
         dest.writeLong(codeBytes);
         dest.writeLong(dataBytes);
         dest.writeLong(cacheBytes);
@@ -257,4 +277,42 @@ public final class StorageStats implements Parcelable {
             return new StorageStats[size];
         }
     };
+
+    private void getArtManagedStats() {
+      try {
+        IStorageStatsManager storageStatsManagerService;
+        // Fetch art stats only if it is not already fetched.
+        if (Flags.getAppArtManagedBytes() && !artStatsFetched) {
+          android.os.IBinder binder = ServiceManager.getService("storagestats");
+          storageStatsManagerService = IStorageStatsManager.Stub.asInterface(binder);
+
+          StorageStats newStats =
+              storageStatsManagerService.queryArtManagedStats(packageName, userHandle, uid);
+
+          dexoptBytes = newStats.dexoptBytes;
+          curProfBytes = newStats.curProfBytes;
+          refProfBytes = newStats.refProfBytes;
+
+          artStatsFetched = true;
+        }
+      } catch (RemoteException e) {
+          Log.e(TAG, "Failed to get art stats", e);
+          e.rethrowFromSystemServer();
+      }
+    }
+
+    private long getDexoptBytes() {
+      getArtManagedStats();
+      return dexoptBytes;
+    }
+
+    private long getCurProfBytes() {
+      getArtManagedStats();
+      return curProfBytes;
+    }
+
+    private long getRefProfBytes() {
+      getArtManagedStats();
+      return refProfBytes;
+    }
 }
diff --git a/core/java/android/app/usage/flags.aconfig b/core/java/android/app/usage/flags.aconfig
index 04c36867271c..520284993841 100644
--- a/core/java/android/app/usage/flags.aconfig
+++ b/core/java/android/app/usage/flags.aconfig
@@ -49,6 +49,16 @@ flag {
 }
 
 flag {
+    name: "get_app_art_managed_bytes"
+    namespace: "system_performance"
+    description: "Bug fixing flag for optional collection of app ART managed file stats"
+    bug: "395548922"
+    metadata {
+      purpose: PURPOSE_BUGFIX
+    }
+}
+
+flag {
     name: "disable_idle_check"
     namespace: "backstage_power"
     description: "disable idle check for USER_SYSTEM during boot up"
diff --git a/services/usage/java/com/android/server/usage/StorageStatsService.java b/services/usage/java/com/android/server/usage/StorageStatsService.java
index e9da53a8a899..5097231b468f 100644
--- a/services/usage/java/com/android/server/usage/StorageStatsService.java
+++ b/services/usage/java/com/android/server/usage/StorageStatsService.java
@@ -352,7 +352,53 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
     }
 
     @Override
-    public StorageStats queryStatsForPackage(String volumeUuid, String packageName, int userId,
+    public StorageStats queryArtManagedStats(String packageName, int userId, int uid) {
+        if (userId != UserHandle.getCallingUserId()) {
+            mContext.enforceCallingOrSelfPermission(
+                    android.Manifest.permission.INTERACT_ACROSS_USERS, TAG);
+        }
+
+        ApplicationInfo appInfo;
+        if (!TextUtils.isEmpty(packageName)) {
+          try {
+            appInfo = mPackage.getApplicationInfoAsUser(packageName,
+                    PackageManager.MATCH_UNINSTALLED_PACKAGES, userId);
+          } catch (NameNotFoundException e) {
+            throw new ParcelableException(e);
+          }
+          uid = appInfo.uid;
+          if (defeatNullable(mPackage.getPackagesForUid(appInfo.uid)).length > 1) {
+              // Multiple packages, skip
+              return translate(new PackageStats(TAG));
+          }
+        }
+
+        int callingUid = Binder.getCallingUid();
+        String callingPackage = mPackage.getNameForUid(callingUid);
+        if (Binder.getCallingUid() != uid) {
+            enforceStatsPermission(Binder.getCallingUid(), callingPackage);
+        }
+
+        final String[] packageNames = defeatNullable(mPackage.getPackagesForUid(uid));
+        final PackageStats stats = new PackageStats(TAG);
+        for (int i = 0; i < packageNames.length; i++) {
+            try {
+                appInfo = mPackage.getApplicationInfoAsUser(packageNames[i],
+                        PackageManager.MATCH_UNINSTALLED_PACKAGES, userId);
+                if (appInfo.isSystemApp() && !appInfo.isUpdatedSystemApp()) {
+                    // We don't count code baked into system image
+                } else {
+                    computeAppArtStats(stats, packageNames[i]);
+                }
+            } catch (NameNotFoundException e) {
+                throw new ParcelableException(e);
+            }
+        }
+        return translate(stats);
+    }
+
+    @Override
+    public  StorageStats queryStatsForPackage(String volumeUuid, String packageName, int userId,
             String callingPackage) {
         if (userId != UserHandle.getCallingUserId()) {
             mContext.enforceCallingOrSelfPermission(
@@ -378,9 +424,10 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
             callerHasStatsPermission = true;
         }
 
+        StorageStats storageStats;
         if (defeatNullable(mPackage.getPackagesForUid(appInfo.uid)).length == 1) {
             // Only one package inside UID means we can fast-path
-            return queryStatsForUid(volumeUuid, appInfo.uid, callingPackage);
+            storageStats = queryStatsForUid(volumeUuid, appInfo.uid, callingPackage);
         } else {
             // Multiple packages means we need to go manual
             final int appId = UserHandle.getAppId(appInfo.uid);
@@ -411,12 +458,16 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
                             packageName, userHandle, callerHasStatsPermission);
                 }, "queryStatsForPackage");
             }
-            return translate(stats);
+            storageStats = translate(stats);
         }
+        storageStats.packageName = packageName;
+        storageStats.userHandle = userId;
+        return storageStats;
     }
 
     @Override
-    public StorageStats queryStatsForUid(String volumeUuid, int uid, String callingPackage) {
+    public StorageStats queryStatsForUid(String volumeUuid, int uid,
+            String callingPackage) {
         final int userId = UserHandle.getUserId(uid);
         final int appId = UserHandle.getAppId(uid);
 
@@ -481,7 +532,10 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
                 storageStatsAugmenter.augmentStatsForUid(stats, uid, callerHasStatsPermission);
             }, "queryStatsForUid");
         }
-        return translate(stats);
+        StorageStats storageStats = translate(stats);
+        storageStats.userHandle = userId;
+        storageStats.uid = uid;
+        return storageStats;
     }
 
     @Override
@@ -592,8 +646,9 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
         }
     }
 
-    private static StorageStats translate(PackageStats stats) {
+    private StorageStats translate(PackageStats stats) {
         final StorageStats res = new StorageStats();
+        res.userHandle = stats.userHandle;
         res.codeBytes = stats.codeSize + stats.externalCodeSize;
         res.dataBytes = stats.dataSize + stats.externalDataSize;
         res.cacheBytes = stats.cacheSize + stats.externalCacheSize;
@@ -967,6 +1022,12 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
         stats.dmSize += getFileBytesInDir(srcDir, ".dm");
         stats.libSize += getDirBytes(new File(sourceDirName + "/lib/"));
 
+        if (!Flags.getAppArtManagedBytes()) {
+            computeAppArtStats(stats, packageName);
+        }
+    }
+
+    private void computeAppArtStats(PackageStats stats, String packageName) {
         // Get dexopt, current profle and reference profile sizes.
         ArtManagedFileStats artManagedFileStats;
         try (var snapshot = getPackageManagerLocal().withFilteredSnapshot()) {