Cache storage info into file.

System is able to query the storage info from AppSearch through
AppSearchStorageStatsAugmenter. Currently, AppSearchImpl is required to
be initialized for querying the storage info. So if system need to query
the storage info before AppSearch is connected/initialized, it would
introduce multiple AppSearchImpl instance initializations for user,
which may block the Setting-Storage UI thread.

In this CL, we cache the user's storage info into a file under their
credential encrypted system directory while initializing
AppSearchUserInstance for the user.
Next time, if system need to query the user's storage info before
AppSearch is intialized, it would try to read the cached info from the
file on disk.

Bug: 179160886
Test: UserStorageInfoTest
Change-Id: I1f760879858a4d667b734018f63e108084c54e1a
diff --git a/service/java/com/android/server/appsearch/AppSearchManagerService.java b/service/java/com/android/server/appsearch/AppSearchManagerService.java
index b49bbc5..4ec2640 100644
--- a/service/java/com/android/server/appsearch/AppSearchManagerService.java
+++ b/service/java/com/android/server/appsearch/AppSearchManagerService.java
@@ -86,6 +86,7 @@
 
 /**
  * The main service implementation which contains AppSearch's platform functionality.
+ *
  * @hide
  */
 public class AppSearchManagerService extends SystemService {
@@ -183,7 +184,6 @@
      * when a user is removed.
      *
      * @param userHandle The multi-user handle of the user that need to be removed.
-     *
      * @see android.os.Environment#getDataSystemCeDirectory
      */
     private void handleUserRemoved(@NonNull UserHandle userHandle) {
@@ -1500,7 +1500,6 @@
         }
     }
 
-    // TODO(b/179160886): Cache the previous storage stats.
     private class AppSearchStorageStatsAugmenter implements StorageStatsAugmenter {
         @Override
         public void augmentStatsForPackageForUser(
@@ -1514,12 +1513,19 @@
 
             try {
                 verifyUserUnlocked(userHandle);
-                Context userContext = mContext.createContextAsUser(userHandle, /*flags=*/ 0);
                 AppSearchUserInstance instance =
-                        mAppSearchUserInstanceManager.getOrCreateUserInstance(
-                                userContext, userHandle, AppSearchConfig.getInstance(EXECUTOR));
-                stats.dataSize += instance.getAppSearchImpl()
-                        .getStorageInfoForPackage(packageName).getSizeBytes();
+                        mAppSearchUserInstanceManager.getUserInstanceOrNull(userHandle);
+                if (instance == null) {
+                    // augment storage info from file
+                    UserStorageInfo userStorageInfo =
+                            mAppSearchUserInstanceManager.getOrCreateUserStorageInfoInstance(
+                                    userHandle);
+                    stats.dataSize +=
+                            userStorageInfo.getSizeBytesForPackage(packageName);
+                } else {
+                    stats.dataSize += instance.getAppSearchImpl()
+                            .getStorageInfoForPackage(packageName).getSizeBytes();
+                }
             } catch (Throwable t) {
                 Log.e(
                         TAG,
@@ -1543,13 +1549,22 @@
                 if (packagesForUid == null) {
                     return;
                 }
-                Context userContext = mContext.createContextAsUser(userHandle, /*flags=*/ 0);
                 AppSearchUserInstance instance =
-                        mAppSearchUserInstanceManager.getOrCreateUserInstance(
-                                userContext, userHandle, AppSearchConfig.getInstance(EXECUTOR));
-                for (int i = 0; i < packagesForUid.length; i++) {
-                    stats.dataSize += instance.getAppSearchImpl()
-                            .getStorageInfoForPackage(packagesForUid[i]).getSizeBytes();
+                        mAppSearchUserInstanceManager.getUserInstanceOrNull(userHandle);
+                if (instance == null) {
+                    // augment storage info from file
+                    UserStorageInfo userStorageInfo =
+                            mAppSearchUserInstanceManager.getOrCreateUserStorageInfoInstance(
+                                    userHandle);
+                    for (int i = 0; i < packagesForUid.length; i++) {
+                        stats.dataSize += userStorageInfo.getSizeBytesForPackage(
+                                packagesForUid[i]);
+                    }
+                } else {
+                    for (int i = 0; i < packagesForUid.length; i++) {
+                        stats.dataSize += instance.getAppSearchImpl()
+                                .getStorageInfoForPackage(packagesForUid[i]).getSizeBytes();
+                    }
                 }
             } catch (Throwable t) {
                 Log.e(TAG, "Unable to augment storage stats for uid " + uid, t);
@@ -1567,19 +1582,24 @@
 
             try {
                 verifyUserUnlocked(userHandle);
-                List<PackageInfo> packagesForUser = mPackageManager.getInstalledPackagesAsUser(
-                        /*flags=*/0, userHandle.getIdentifier());
-                if (packagesForUser == null) {
-                    return;
-                }
-                Context userContext = mContext.createContextAsUser(userHandle, /*flags=*/ 0);
                 AppSearchUserInstance instance =
-                        mAppSearchUserInstanceManager.getOrCreateUserInstance(
-                                userContext, userHandle, AppSearchConfig.getInstance(EXECUTOR));
-                for (int i = 0; i < packagesForUser.size(); i++) {
-                    String packageName = packagesForUser.get(i).packageName;
-                    stats.dataSize += instance.getAppSearchImpl()
-                            .getStorageInfoForPackage(packageName).getSizeBytes();
+                        mAppSearchUserInstanceManager.getUserInstanceOrNull(userHandle);
+                if (instance == null) {
+                    // augment storage info from file
+                    UserStorageInfo userStorageInfo =
+                            mAppSearchUserInstanceManager.getOrCreateUserStorageInfoInstance(
+                                    userHandle);
+                    stats.dataSize += userStorageInfo.getTotalSizeBytes();
+                } else {
+                    List<PackageInfo> packagesForUser = mPackageManager.getInstalledPackagesAsUser(
+                            /*flags=*/0, userHandle.getIdentifier());
+                    if (packagesForUser != null) {
+                        for (int i = 0; i < packagesForUser.size(); i++) {
+                            String packageName = packagesForUser.get(i).packageName;
+                            stats.dataSize += instance.getAppSearchImpl()
+                                    .getStorageInfoForPackage(packageName).getSizeBytes();
+                        }
+                    }
                 }
             } catch (Throwable t) {
                 Log.e(TAG, "Unable to augment storage stats for " + userHandle, t);
diff --git a/service/java/com/android/server/appsearch/AppSearchUserInstanceManager.java b/service/java/com/android/server/appsearch/AppSearchUserInstanceManager.java
index 529f2b0..ca56325 100644
--- a/service/java/com/android/server/appsearch/AppSearchUserInstanceManager.java
+++ b/service/java/com/android/server/appsearch/AppSearchUserInstanceManager.java
@@ -17,6 +17,7 @@
 package com.android.server.appsearch;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.app.appsearch.exceptions.AppSearchException;
 import android.content.Context;
 import android.os.Environment;
@@ -50,6 +51,8 @@
 
     @GuardedBy("mInstancesLocked")
     private final Map<UserHandle, AppSearchUserInstance> mInstancesLocked = new ArrayMap<>();
+    @GuardedBy("mStorageInfoLocked")
+    private final Map<UserHandle, UserStorageInfo> mStorageInfoLocked = new ArrayMap<>();
 
     private AppSearchUserInstanceManager() {}
 
@@ -130,6 +133,9 @@
                 instance.getAppSearchImpl().close();
             }
         }
+        synchronized (mStorageInfoLocked) {
+            mStorageInfoLocked.remove(userHandle);
+        }
     }
 
     /**
@@ -160,6 +166,39 @@
     }
 
     /**
+     * Returns the initialized {@link AppSearchUserInstance} for the given user, or {@code null} if
+     * no such instance exists.
+     *
+     * @param userHandle The multi-user handle of the device user calling AppSearch
+     */
+    @Nullable
+    public AppSearchUserInstance getUserInstanceOrNull(@NonNull UserHandle userHandle) {
+        Objects.requireNonNull(userHandle);
+        synchronized (mInstancesLocked) {
+            return mInstancesLocked.get(userHandle);
+        }
+    }
+
+    /**
+     * Gets an {@link UserStorageInfo} for the given user.
+     *
+     * @param userHandle The multi-user handle of the device user
+     * @return An initialized {@link UserStorageInfo} for this user
+     */
+    @NonNull
+    public UserStorageInfo getOrCreateUserStorageInfoInstance(@NonNull UserHandle userHandle) {
+        Objects.requireNonNull(userHandle);
+        synchronized (mStorageInfoLocked) {
+            UserStorageInfo userStorageInfo = mStorageInfoLocked.get(userHandle);
+            if (userStorageInfo == null) {
+                userStorageInfo = new UserStorageInfo(getAppSearchDir(userHandle));
+                mStorageInfoLocked.put(userHandle, userStorageInfo);
+            }
+            return userStorageInfo;
+        }
+    }
+
+    /**
      * Returns the list of all {@link UserHandle}s.
      *
      * <p>It can return an empty list if there is no {@link AppSearchUserInstance} created yet.
@@ -197,6 +236,10 @@
                 VisibilityStoreImpl.create(appSearchImpl, userContext);
         long prepareVisibilityStoreLatencyEndMillis = SystemClock.elapsedRealtime();
 
+        // Update storage info file
+        UserStorageInfo userStorageInfo = getOrCreateUserStorageInfoInstance(userHandle);
+        userStorageInfo.updateStorageInfoFile(appSearchImpl);
+
         initStatsBuilder
                 .setTotalLatencyMillis(
                         (int) (SystemClock.elapsedRealtime() - totalLatencyStartMillis))
diff --git a/service/java/com/android/server/appsearch/UserStorageInfo.java b/service/java/com/android/server/appsearch/UserStorageInfo.java
new file mode 100644
index 0000000..d6e8598
--- /dev/null
+++ b/service/java/com/android/server/appsearch/UserStorageInfo.java
@@ -0,0 +1,152 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.appsearch;
+
+import static com.android.server.appsearch.external.localstorage.util.PrefixUtil.getPackageName;
+
+import android.annotation.NonNull;
+import android.util.ArrayMap;
+import android.util.Log;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.appsearch.external.localstorage.AppSearchImpl;
+
+import com.google.android.icing.proto.DocumentStorageInfoProto;
+import com.google.android.icing.proto.NamespaceStorageInfoProto;
+import com.google.android.icing.proto.StorageInfoProto;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.InputStream;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+
+/** Saves the storage info read from file for a user. */
+public final class UserStorageInfo {
+    public static final String STORAGE_INFO_FILE = "appsearch_storage";
+    private static final String TAG = "AppSearchUserStorage";
+    private final ReadWriteLock mReadWriteLock = new ReentrantReadWriteLock();
+    private final File mStorageInfoFile;
+
+    // Saves storage usage byte size for each package under the user, keyed by package name.
+    private Map<String, Long> mPackageStorageSizeMap;
+    // Saves storage usage byte size for all packages under the user.
+    private long mTotalStorageSizeBytes;
+
+    public UserStorageInfo(@NonNull File fileParentPath) {
+        Objects.requireNonNull(fileParentPath);
+        mStorageInfoFile = new File(fileParentPath, STORAGE_INFO_FILE);
+        readStorageInfoFromFile();
+    }
+
+    /**
+     * Updates storage info file with the latest storage info queried through
+     * {@link AppSearchImpl}.
+     */
+    public void updateStorageInfoFile(@NonNull AppSearchImpl appSearchImpl) {
+        Objects.requireNonNull(appSearchImpl);
+        mReadWriteLock.writeLock().lock();
+        try (FileOutputStream out = new FileOutputStream(mStorageInfoFile)) {
+            appSearchImpl.getRawStorageInfoProto().writeTo(out);
+        } catch (Throwable e) {
+            Log.w(TAG, "Failed to dump storage info into file", e);
+        } finally {
+            mReadWriteLock.writeLock().unlock();
+        }
+    }
+
+    /**
+     * Gets storage usage byte size for a package with a given package name.
+     *
+     * <p> Please note the storage info cached in file may be stale.
+     */
+    public long getSizeBytesForPackage(@NonNull String packageName) {
+        Objects.requireNonNull(packageName);
+        return mPackageStorageSizeMap.getOrDefault(packageName, 0L);
+    }
+
+    /**
+     * Gets total storage usage byte size for all packages under the user.
+     *
+     * <p> Please note the storage info cached in file may be stale.
+     */
+    public long getTotalSizeBytes() {
+        return mTotalStorageSizeBytes;
+    }
+
+    @VisibleForTesting
+    void readStorageInfoFromFile() {
+        if (mStorageInfoFile.exists()) {
+            mReadWriteLock.readLock().lock();
+            try (InputStream in = new FileInputStream(mStorageInfoFile)) {
+                StorageInfoProto storageInfo = StorageInfoProto.parseFrom(in);
+                mTotalStorageSizeBytes = storageInfo.getTotalStorageSize();
+                mPackageStorageSizeMap = calculatePackageStorageInfoMap(storageInfo);
+                return;
+            } catch (Throwable e) {
+                Log.w(TAG, "Failed to read storage info from file", e);
+            } finally {
+                mReadWriteLock.readLock().unlock();
+            }
+        }
+        mTotalStorageSizeBytes = 0;
+        mPackageStorageSizeMap = Collections.emptyMap();
+    }
+
+    /** Calculates storage usage byte size for packages from a {@link StorageInfoProto}. */
+    // TODO(b/198553756): Storage cache effort has created two copies of the storage
+    // calculation/interpolation logic.
+    @NonNull
+    @VisibleForTesting
+    Map<String, Long> calculatePackageStorageInfoMap(@NonNull StorageInfoProto storageInfo) {
+        Map<String, Long> packageStorageInfoMap = new ArrayMap<>();
+        if (storageInfo.hasDocumentStorageInfo()) {
+            DocumentStorageInfoProto documentStorageInfo = storageInfo.getDocumentStorageInfo();
+            List<NamespaceStorageInfoProto> namespaceStorageInfoList =
+                    documentStorageInfo.getNamespaceStorageInfoList();
+
+            Map<String, Integer> packageDocumentCountMap = new ArrayMap<>();
+            long totalDocuments = 0;
+            for (int i = 0; i < namespaceStorageInfoList.size(); i++) {
+                NamespaceStorageInfoProto namespaceStorageInfo = namespaceStorageInfoList.get(i);
+                String packageName = getPackageName(namespaceStorageInfo.getNamespace());
+                int namespaceDocuments = namespaceStorageInfo.getNumAliveDocuments()
+                        + namespaceStorageInfo.getNumExpiredDocuments();
+                totalDocuments += namespaceDocuments;
+                packageDocumentCountMap.put(packageName,
+                        packageDocumentCountMap.getOrDefault(packageName, 0)
+                                + namespaceDocuments);
+            }
+
+            long totalStorageSize = storageInfo.getTotalStorageSize();
+            for (Map.Entry<String, Integer> entry : packageDocumentCountMap.entrySet()) {
+                // Since we don't have the exact size of all the documents, we do an estimation.
+                // Note that while the total storage takes into account schema, index, etc. in
+                // addition to documents, we'll only calculate the percentage based on number of
+                // documents under packages.
+                packageStorageInfoMap.put(entry.getKey(),
+                        (long) (entry.getValue() * 1.0 / totalDocuments * totalStorageSize));
+            }
+        }
+        return Collections.unmodifiableMap(packageStorageInfoMap);
+    }
+}
diff --git a/testing/servicestests/src/com/android/server/appsearch/UserStorageInfoTest.java b/testing/servicestests/src/com/android/server/appsearch/UserStorageInfoTest.java
new file mode 100644
index 0000000..c79671f
--- /dev/null
+++ b/testing/servicestests/src/com/android/server/appsearch/UserStorageInfoTest.java
@@ -0,0 +1,123 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.appsearch;
+
+import static com.android.server.appsearch.UserStorageInfo.STORAGE_INFO_FILE;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.android.server.appsearch.icing.proto.DocumentStorageInfoProto;
+import com.android.server.appsearch.icing.proto.NamespaceStorageInfoProto;
+import com.android.server.appsearch.icing.proto.StorageInfoProto;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+
+public class UserStorageInfoTest {
+    @Rule public TemporaryFolder mTemporaryFolder = new TemporaryFolder();
+    private File mFileParentPath;
+    private UserStorageInfo mUserStorageInfo;
+
+    @Before
+    public void setUp() throws Exception {
+        mFileParentPath = mTemporaryFolder.newFolder();
+        mUserStorageInfo = new UserStorageInfo(mFileParentPath);
+    }
+
+    @Test
+    public void testReadStorageInfoFromFile() throws IOException {
+        NamespaceStorageInfoProto namespaceStorageInfo1 =
+                NamespaceStorageInfoProto.newBuilder()
+                        .setNamespace("pkg1$db/namespace")
+                        .setNumAliveDocuments(2)
+                        .setNumExpiredDocuments(1)
+                        .build();
+        NamespaceStorageInfoProto namespaceStorageInfo2 =
+                NamespaceStorageInfoProto.newBuilder()
+                        .setNamespace("pkg2$db/namespace")
+                        .setNumAliveDocuments(3)
+                        .setNumExpiredDocuments(3)
+                        .build();
+        DocumentStorageInfoProto documentStorageInfo =
+                DocumentStorageInfoProto.newBuilder()
+                        .setNumAliveDocuments(5)
+                        .setNumExpiredDocuments(4)
+                        .addNamespaceStorageInfo(namespaceStorageInfo1)
+                        .addNamespaceStorageInfo(namespaceStorageInfo2)
+                        .build();
+        StorageInfoProto storageInfo =
+                StorageInfoProto.newBuilder()
+                        .setDocumentStorageInfo(documentStorageInfo)
+                        .setTotalStorageSize(9)
+                        .build();
+        File storageInfoFile = new File(mFileParentPath, STORAGE_INFO_FILE);
+        try (FileOutputStream out = new FileOutputStream(storageInfoFile)) {
+            storageInfo.writeTo(out);
+        }
+
+        mUserStorageInfo.readStorageInfoFromFile();
+
+        assertThat(mUserStorageInfo.getTotalSizeBytes()).isEqualTo(
+                storageInfo.getTotalStorageSize());
+        // We calculate the package storage size based on number of documents a package has.
+        // Here, total storage size is 9. pkg1 has 3 docs and pkg2 has 6 docs. So storage size of
+        // pkg1 is 3. pkg2's storage size is 6.
+        assertThat(mUserStorageInfo.getSizeBytesForPackage("pkg1")).isEqualTo(3);
+        assertThat(mUserStorageInfo.getSizeBytesForPackage("pkg2")).isEqualTo(6);
+        assertThat(mUserStorageInfo.getSizeBytesForPackage("invalid_pkg")).isEqualTo(0);
+    }
+
+    @Test
+    public void testCalculatePackageStorageInfoMap() {
+        NamespaceStorageInfoProto namespaceStorageInfo1 =
+                NamespaceStorageInfoProto.newBuilder()
+                        .setNamespace("pkg1$db/namespace")
+                        .setNumAliveDocuments(2)
+                        .setNumExpiredDocuments(1)
+                        .build();
+        NamespaceStorageInfoProto namespaceStorageInfo2 =
+                NamespaceStorageInfoProto.newBuilder()
+                        .setNamespace("pkg2$db/namespace")
+                        .setNumAliveDocuments(3)
+                        .setNumExpiredDocuments(3)
+                        .build();
+        DocumentStorageInfoProto documentStorageInfo =
+                DocumentStorageInfoProto.newBuilder()
+                        .setNumAliveDocuments(5)
+                        .setNumExpiredDocuments(4)
+                        .addNamespaceStorageInfo(namespaceStorageInfo1)
+                        .addNamespaceStorageInfo(namespaceStorageInfo2)
+                        .build();
+        StorageInfoProto storageInfo =
+                StorageInfoProto.newBuilder()
+                        .setDocumentStorageInfo(documentStorageInfo)
+                        .setTotalStorageSize(9)
+                        .build();
+
+        // We calculate the package storage size based on number of documents a package has.
+        // Here, total storage size is 9. pkg1 has 3 docs and pkg2 has 6 docs. So storage size of
+        // pkg1 is 3. pkg2's storage size is 6.
+        assertThat(mUserStorageInfo.calculatePackageStorageInfoMap(storageInfo))
+                .containsExactly("pkg1", 3L, "pkg2", 6L);
+    }
+}