Use the StorageStatsManager in FileCollector.

This should vastly improve the speed of the FileCollector.

Change-Id: I7a878a0622bbd6c758fb1d36125414d8b025e709
Fixes: 35807386
Test: Existing tests continue to pass.

3 files changed, 69 insertions, 27 deletions

diff --git a/services/core/java/com/android/server/storage/DiskStatsLoggingService.java b/services/core/java/com/android/server/storage/DiskStatsLoggingService.java
index 4035adedafe1..1783dccd0b06 100644
--- a/services/core/java/com/android/server/storage/DiskStatsLoggingService.java
+++ b/services/core/java/com/android/server/storage/DiskStatsLoggingService.java
@@ -73,7 +73,6 @@ public class DiskStatsLoggingService extends JobService {
         final int userId = UserHandle.myUserId();
         UserEnvironment environment = new UserEnvironment(userId);
         LogRunnable task = new LogRunnable();
-        task.setRootDirectory(environment.getExternalStorageDirectory());
         task.setDownloadsDirectory(
                 environment.getExternalStoragePublicDirectory(Environment.DIRECTORY_DOWNLOADS));
         task.setSystemSize(FileCollector.getSystemSize(this));
@@ -127,14 +126,10 @@ public class DiskStatsLoggingService extends JobService {
         private JobParameters mParams;
         private AppCollector mCollector;
         private File mOutputFile;
-        private File mRootDirectory;
         private File mDownloadsDirectory;
+        private Context mContext;
         private long mSystemSize;
 
-        public void setRootDirectory(File file) {
-            mRootDirectory = file;
-        }
-
         public void setDownloadsDirectory(File file) {
             mDownloadsDirectory = file;
         }
@@ -151,14 +146,25 @@ public class DiskStatsLoggingService extends JobService {
             mSystemSize = size;
         }
 
+        public void setContext(Context context) {
+            mContext = context;
+        }
+
         public void setJobService(JobService jobService, JobParameters params) {
             mJobService = jobService;
             mParams = params;
         }
 
         public void run() {
-            FileCollector.MeasurementResult mainCategories =
-                    FileCollector.getMeasurementResult(mRootDirectory);
+            FileCollector.MeasurementResult mainCategories;
+            try {
+                mainCategories = FileCollector.getMeasurementResult(mContext);
+            } catch (IllegalStateException e) {
+                // This can occur if installd has an issue.
+                Log.e(TAG, "Error while measuring storage", e);
+                finishJob(true);
+                return;
+            }
             FileCollector.MeasurementResult downloads =
                     FileCollector.getMeasurementResult(mDownloadsDirectory);
 
@@ -168,12 +174,10 @@ public class DiskStatsLoggingService extends JobService {
                 needsReschedule = false;
                 logToFile(mainCategories, downloads, stats, mSystemSize);
             } else {
-                Log.w("TAG", "Timed out while fetching package stats.");
+                Log.w(TAG, "Timed out while fetching package stats.");
             }
 
-            if (mJobService != null) {
-                mJobService.jobFinished(mParams, needsReschedule);
-            }
+            finishJob(needsReschedule);
         }
 
         private void logToFile(MeasurementResult mainCategories, MeasurementResult downloads,
@@ -187,5 +191,11 @@ public class DiskStatsLoggingService extends JobService {
                 Log.e(TAG, "Exception while writing opportunistic disk file cache.", e);
             }
         }
+
+        private void finishJob(boolean needsReschedule) {
+            if (mJobService != null) {
+                mJobService.jobFinished(mParams, needsReschedule);
+            }
+        }
     }
 }
\ No newline at end of file
diff --git a/services/core/java/com/android/server/storage/FileCollector.java b/services/core/java/com/android/server/storage/FileCollector.java
index 90f9f1391679..04f25a484c79 100644
--- a/services/core/java/com/android/server/storage/FileCollector.java
+++ b/services/core/java/com/android/server/storage/FileCollector.java
@@ -17,8 +17,11 @@
 package com.android.server.storage;
 
 import android.annotation.IntDef;
+import android.app.usage.ExternalStorageStats;
+import android.app.usage.StorageStatsManager;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.os.UserHandle;
 import android.os.storage.StorageManager;
 import android.os.storage.VolumeInfo;
 import android.util.ArrayMap;
@@ -154,8 +157,29 @@ public class FileCollector {
     }
 
     /**
+     * Returns the file categorization result for the primary internal storage UUID.
+     *
+     * @param context
+     */
+    public static MeasurementResult getMeasurementResult(Context context) {
+        MeasurementResult result = new MeasurementResult();
+        StorageStatsManager ssm =
+                (StorageStatsManager) context.getSystemService(Context.STORAGE_STATS_SERVICE);
+        ExternalStorageStats stats =
+                ssm.queryExternalStatsForUser(
+                        StorageManager.UUID_PRIVATE_INTERNAL, UserHandle.of(context.getUserId()));
+        result.imagesSize = stats.getImageBytes();
+        result.videosSize = stats.getVideoBytes();
+        result.audioSize = stats.getAudioBytes();
+        result.miscSize =
+                stats.getTotalBytes() - result.imagesSize - result.videosSize - result.audioSize;
+        return result;
+    }
+
+    /**
      * Returns the size of a system for a given context. This is done by finding the difference
      * between the shared data and the total primary storage size.
+     *
      * @param context Context to use to get storage information.
      */
     public static long getSystemSize(Context context) {
diff --git a/services/tests/servicestests/src/com/android/server/storage/DiskStatsLoggingServiceTest.java b/services/tests/servicestests/src/com/android/server/storage/DiskStatsLoggingServiceTest.java
index 378908655cb0..81ce6069c92e 100644
--- a/services/tests/servicestests/src/com/android/server/storage/DiskStatsLoggingServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/storage/DiskStatsLoggingServiceTest.java
@@ -20,15 +20,18 @@ import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
-import static org.mockito.Mockito.eq;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
+import static org.mockito.Matchers.anyLong;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Matchers.isNull;
 import static org.mockito.Mockito.when;
 
 import android.app.job.JobService;
-import android.app.job.JobParameters;
+import android.app.usage.ExternalStorageStats;
+import android.app.usage.StorageStatsManager;
 import android.content.pm.PackageStats;
+import android.os.UserHandle;
 import android.test.AndroidTestCase;
+import android.util.Log;
 
 import com.android.server.storage.DiskStatsLoggingService.LogRunnable;
 
@@ -52,8 +55,10 @@ import java.util.ArrayList;
 public class DiskStatsLoggingServiceTest extends AndroidTestCase {
     @Rule public TemporaryFolder mTemporaryFolder;
     @Rule public TemporaryFolder mDownloads;
-    @Rule public TemporaryFolder mRootDirectory;
     @Mock private AppCollector mCollector;
+    @Mock private JobService mJobService;
+    @Mock private StorageStatsManager mSsm;
+    private ExternalStorageStats mStorageStats;
     private File mInputFile;
 
 
@@ -66,8 +71,10 @@ public class DiskStatsLoggingServiceTest extends AndroidTestCase {
         mInputFile = mTemporaryFolder.newFile();
         mDownloads = new TemporaryFolder();
         mDownloads.create();
-        mRootDirectory = new TemporaryFolder();
-        mRootDirectory.create();
+        mStorageStats = new ExternalStorageStats();
+        when(mSsm.queryExternalStatsForUser(isNull(String.class), any(UserHandle.class)))
+                .thenReturn(mStorageStats);
+        when(mJobService.getSystemService(anyString())).thenReturn(mSsm);
     }
 
     @Test
@@ -75,9 +82,9 @@ public class DiskStatsLoggingServiceTest extends AndroidTestCase {
         LogRunnable task = new LogRunnable();
         task.setAppCollector(mCollector);
         task.setDownloadsDirectory(mDownloads.getRoot());
-        task.setRootDirectory(mRootDirectory.getRoot());
         task.setLogOutputFile(mInputFile);
         task.setSystemSize(0L);
+        task.setContext(mJobService);
         task.run();
 
         JSONObject json = getJsonOutput();
@@ -99,10 +106,10 @@ public class DiskStatsLoggingServiceTest extends AndroidTestCase {
     public void testPopulatedLogTask() throws Exception {
         // Write data to directories.
         writeDataToFile(mDownloads.newFile(), "lol");
-        writeDataToFile(mRootDirectory.newFile("test.jpg"), "1234");
-        writeDataToFile(mRootDirectory.newFile("test.mp4"), "12345");
-        writeDataToFile(mRootDirectory.newFile("test.mp3"), "123456");
-        writeDataToFile(mRootDirectory.newFile("test.whatever"), "1234567");
+        mStorageStats.audioBytes = 6L;
+        mStorageStats.imageBytes = 4L;
+        mStorageStats.videoBytes = 5L;
+        mStorageStats.totalBytes = 22L;
 
         // Write apps.
         ArrayList<PackageStats> apps = new ArrayList<>();
@@ -110,15 +117,16 @@ public class DiskStatsLoggingServiceTest extends AndroidTestCase {
         testApp.dataSize = 5L;
         testApp.cacheSize = 55L;
         testApp.codeSize = 10L;
+        testApp.userHandle = UserHandle.USER_SYSTEM;
         apps.add(testApp);
-        when(mCollector.getPackageStats(anyInt())).thenReturn(apps);
+        when(mCollector.getPackageStats(anyLong())).thenReturn(apps);
 
         LogRunnable task = new LogRunnable();
         task.setAppCollector(mCollector);
         task.setDownloadsDirectory(mDownloads.getRoot());
-        task.setRootDirectory(mRootDirectory.getRoot());
         task.setLogOutputFile(mInputFile);
         task.setSystemSize(10L);
+        task.setContext(mJobService);
         task.run();
 
         JSONObject json = getJsonOutput();
@@ -143,9 +151,9 @@ public class DiskStatsLoggingServiceTest extends AndroidTestCase {
         LogRunnable task = new LogRunnable();
         task.setAppCollector(mCollector);
         task.setDownloadsDirectory(mDownloads.getRoot());
-        task.setRootDirectory(mRootDirectory.getRoot());
         task.setLogOutputFile(mInputFile);
         task.setSystemSize(10L);
+        task.setContext(mJobService);
         task.run();
 
         // No exception should be thrown.