Merge "Add internal APIs for system services to augment storage stats."

2 files changed, 105 insertions, 0 deletions

diff --git a/services/usage/java/com/android/server/usage/StorageStatsManagerInternal.java b/services/usage/java/com/android/server/usage/StorageStatsManagerInternal.java
new file mode 100644
index 000000000000..a5325482605d
--- /dev/null
+++ b/services/usage/java/com/android/server/usage/StorageStatsManagerInternal.java
@@ -0,0 +1,52 @@
+/*
+ * Copyright (C) 2020 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.usage;
+
+import android.annotation.NonNull;
+import android.annotation.UserIdInt;
+import android.content.pm.PackageStats;
+
+/**
+ * StorageStatsManager local system service interface.
+ *
+ * Only for use within the system server.
+ */
+public abstract class StorageStatsManagerInternal {
+    /**
+     * Class used to augment {@link PackageStats} with the data stored by the system on
+     * behalf of apps in system specific directories
+     * ({@link android.os.Environment#getDataSystemDirectory},
+     * {@link android.os.Environment#getDataSystemCeDirectory}, etc).
+     */
+    public interface StorageStatsAugmenter {
+        void augmentStatsForPackage(@NonNull PackageStats stats,
+                @NonNull String packageName, @UserIdInt int userId,
+                @NonNull String callingPackage);
+        void augmentStatsForUid(@NonNull PackageStats stats, int uid,
+                @NonNull String callingPackage);
+        void augmentStatsForUser(@NonNull PackageStats stats, @UserIdInt int userId,
+                @NonNull String callingPackage);
+    }
+
+    /**
+     * Register a {@link StorageStatsAugmenter}.
+     *
+     * @param augmenter the {@link StorageStatsAugmenter} object to be registered.
+     * @param tag the identifier to be used for debugging in logs/trace.
+     */
+    public abstract void registerStorageStatsAugmenter(@NonNull StorageStatsAugmenter augmenter,
+            @NonNull String tag);
+}
diff --git a/services/usage/java/com/android/server/usage/StorageStatsService.java b/services/usage/java/com/android/server/usage/StorageStatsService.java
index 531a93117752..18b640ff6bf5 100644
--- a/services/usage/java/com/android/server/usage/StorageStatsService.java
+++ b/services/usage/java/com/android/server/usage/StorageStatsService.java
@@ -18,6 +18,7 @@ package com.android.server.usage;
 
 import static com.android.internal.util.ArrayUtils.defeatNullable;
 import static com.android.server.pm.PackageManagerService.PLATFORM_PACKAGE_NAME;
+import static com.android.server.usage.StorageStatsManagerInternal.StorageStatsAugmenter;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -46,6 +47,7 @@ import android.os.Message;
 import android.os.ParcelableException;
 import android.os.StatFs;
 import android.os.SystemProperties;
+import android.os.Trace;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.os.storage.CrateInfo;
@@ -58,6 +60,7 @@ import android.text.TextUtils;
 import android.text.format.DateUtils;
 import android.util.ArrayMap;
 import android.util.DataUnit;
+import android.util.Pair;
 import android.util.Slog;
 import android.util.SparseLongArray;
 
@@ -77,6 +80,8 @@ import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.function.Consumer;
 
 public class StorageStatsService extends IStorageStatsManager.Stub {
     private static final String TAG = "StorageStatsService";
@@ -111,6 +116,9 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
     private final Installer mInstaller;
     private final H mHandler;
 
+    private final CopyOnWriteArrayList<Pair<String, StorageStatsAugmenter>>
+            mStorageStatsAugmenters = new CopyOnWriteArrayList<>();
+
     public StorageStatsService(Context context) {
         mContext = Preconditions.checkNotNull(context);
         mAppOps = Preconditions.checkNotNull(context.getSystemService(AppOpsManager.class));
@@ -139,6 +147,8 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
                 }
             }
         });
+
+        LocalServices.addService(StorageStatsManagerInternal.class, new LocalService());
     }
 
     private void invalidateMounts() {
@@ -300,6 +310,12 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
             } catch (InstallerException e) {
                 throw new ParcelableException(new IOException(e.getMessage()));
             }
+            if (volumeUuid == StorageManager.UUID_PRIVATE_INTERNAL) {
+                forEachStorageStatsAugmenter((storageStatsAugmenter) -> {
+                    storageStatsAugmenter.augmentStatsForPackage(stats,
+                            packageName, userId, callingPackage);
+                }, "queryStatsForPackage");
+            }
             return translate(stats);
         }
     }
@@ -353,6 +369,12 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
         } catch (InstallerException e) {
             throw new ParcelableException(new IOException(e.getMessage()));
         }
+
+        if (volumeUuid == StorageManager.UUID_PRIVATE_INTERNAL) {
+            forEachStorageStatsAugmenter((storageStatsAugmenter) -> {
+                storageStatsAugmenter.augmentStatsForUid(stats, uid, callingPackage);
+            }, "queryStatsForUid");
+        }
         return translate(stats);
     }
 
@@ -379,6 +401,12 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
         } catch (InstallerException e) {
             throw new ParcelableException(new IOException(e.getMessage()));
         }
+
+        if (volumeUuid == StorageManager.UUID_PRIVATE_INTERNAL) {
+            forEachStorageStatsAugmenter((storageStatsAugmenter) -> {
+                storageStatsAugmenter.augmentStatsForUser(stats, userId, callingPackage);
+            }, "queryStatsForUser");
+        }
         return translate(stats);
     }
 
@@ -705,4 +733,29 @@ public class StorageStatsService extends IStorageStatsManager.Stub {
             throw new ParcelableException(new IOException(e.getMessage()));
         }
     }
+
+    void forEachStorageStatsAugmenter(@NonNull Consumer<StorageStatsAugmenter> consumer,
+                @NonNull String queryTag) {
+        for (int i = 0, count = mStorageStatsAugmenters.size(); i < count; ++i) {
+            final Pair<String, StorageStatsAugmenter> pair = mStorageStatsAugmenters.get(i);
+            final String augmenterTag = pair.first;
+            final StorageStatsAugmenter storageStatsAugmenter = pair.second;
+
+            Trace.traceBegin(Trace.TRACE_TAG_SYSTEM_SERVER, queryTag + ":" + augmenterTag);
+            try {
+                consumer.accept(storageStatsAugmenter);
+            } finally {
+                Trace.traceEnd(Trace.TRACE_TAG_SYSTEM_SERVER);
+            }
+        }
+    }
+
+    private class LocalService extends StorageStatsManagerInternal {
+        @Override
+        public void registerStorageStatsAugmenter(
+                @NonNull StorageStatsAugmenter storageStatsAugmenter,
+                @NonNull String tag) {
+            mStorageStatsAugmenters.add(Pair.create(tag, storageStatsAugmenter));
+        }
+    }
 }