Merge "Use the StorageStatsManager in FileCollector." into oc-dev

3 files changed, 139 insertions, 28 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..9a08ac3b4eca 100644
--- a/services/core/java/com/android/server/storage/DiskStatsLoggingService.java
+++ b/services/core/java/com/android/server/storage/DiskStatsLoggingService.java
@@ -73,13 +73,13 @@ 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));
         task.setLogOutputFile(new File(DUMPSYS_CACHE_PATH));
         task.setAppCollector(collector);
         task.setJobService(this, params);
+        task.setContext(this);
         AsyncTask.execute(task);
         return true;
     }
@@ -106,7 +106,8 @@ public class DiskStatsLoggingService extends JobService {
     }
 
     private static boolean isCharging(Context context) {
-        BatteryManager batteryManager = context.getSystemService(BatteryManager.class);
+        BatteryManager batteryManager =
+                (BatteryManager) context.getSystemService(Context.BATTERY_SERVICE);
         if (batteryManager != null) {
             return batteryManager.isCharging();
         }
@@ -127,14 +128,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 +148,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 +176,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 +193,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..0c119a72e7f9 100644
--- a/services/core/java/com/android/server/storage/FileCollector.java
+++ b/services/core/java/com/android/server/storage/FileCollector.java
@@ -17,13 +17,17 @@
 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;
 
 import java.io.File;
+import java.io.IOException;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.Map;
@@ -154,15 +158,46 @@ 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 = null;
+        try {
+            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;
+        } catch (IOException e) {
+            throw new IllegalStateException("Could not query storage");
+        }
+
+        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) {
         PackageManager pm = context.getPackageManager();
         VolumeInfo primaryVolume = pm.getPrimaryStorageCurrentVolume();
 
-        StorageManager sm = context.getSystemService(StorageManager.class);
+        StorageManager sm = (StorageManager) context.getSystemService(Context.STORAGE_SERVICE);
         VolumeInfo shared = sm.findEmulatedForPrivate(primaryVolume);
         if (shared == null) {
             return 0;
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..375edf3cb5ae 100644
--- a/services/tests/servicestests/src/com/android/server/storage/DiskStatsLoggingServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/storage/DiskStatsLoggingServiceTest.java
@@ -20,16 +20,31 @@ 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.Matchers.anyLong;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Matchers.isNull;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
 
 import android.app.job.JobService;
-import android.app.job.JobParameters;
+import android.app.job.JobServiceEngine;
+import android.app.usage.ExternalStorageStats;
+import android.app.usage.StorageStatsManager;
+import android.content.Context;
+import android.content.pm.PackageManager;
 import android.content.pm.PackageStats;
+import android.os.BatteryManager;
+import android.os.UserHandle;
+import android.os.UserManager;
+import android.os.storage.StorageManager;
+import android.os.storage.VolumeInfo;
+import android.provider.Settings;
 import android.test.AndroidTestCase;
+import android.test.mock.MockContentResolver;
 
+import com.android.internal.util.test.FakeSettingsProvider;
 import com.android.server.storage.DiskStatsLoggingService.LogRunnable;
 
 import libcore.io.IoUtils;
@@ -46,14 +61,17 @@ import org.mockito.MockitoAnnotations;
 
 import java.io.File;
 import java.io.PrintStream;
+import java.lang.reflect.Field;
 import java.util.ArrayList;
 
 @RunWith(JUnit4.class)
 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 +84,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 +95,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 +119,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 +130,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,14 +164,57 @@ 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.
     }
 
+    @Test
+    public void testDontCrashOnRun() throws Exception {
+        DiskStatsLoggingService service = spy(new DiskStatsLoggingService());
+        BatteryManager batteryManager = mock(BatteryManager.class);
+        when(batteryManager.isCharging()).thenReturn(true);
+        doReturn(batteryManager).when(service).getSystemService(Context.BATTERY_SERVICE);
+        UserManager userManager = mock(UserManager.class);
+        when(userManager.getUsers()).thenReturn(new ArrayList<>());
+        doReturn(userManager).when(service).getSystemService(Context.USER_SERVICE);
+        doReturn(mSsm).when(service).getSystemService(Context.STORAGE_STATS_SERVICE);
+        doReturn(mock(StorageManager.class))
+                .when(service)
+                .getSystemService(Context.STORAGE_SERVICE);
+
+        MockContentResolver cr = new MockContentResolver();
+        cr.addProvider(Settings.AUTHORITY, new FakeSettingsProvider());
+        doReturn(cr).when(service).getContentResolver();
+
+        PackageManager pm = mock(PackageManager.class);
+        VolumeInfo volumeInfo =
+                new VolumeInfo(VolumeInfo.ID_PRIVATE_INTERNAL, VolumeInfo.TYPE_PRIVATE, null, null);
+        when(pm.getPrimaryStorageCurrentVolume()).thenReturn(volumeInfo);
+        doReturn(pm).when(service).getPackageManager();
+
+        doReturn(0).when(service).getUserId();
+
+        // UGGGGGHHHHHHH! jobFinished is a final method on JobService which crashes when called if
+        // the JobService isn't initialized for real. ServiceTestCase doesn't let us initialize a
+        // service which is built into the framework without crashing, though, so we can't make a
+        // real initialized service.
+        //
+        // And so, we use reflection to set the JobServiceEngine, which is used by the final method,
+        // to be something which won't crash when called.
+        final Field field = JobService.class.getDeclaredField("mEngine");
+        field.setAccessible(true);
+        field.set(service, mock(JobServiceEngine.class));
+
+        // Note: This won't clobber your on-device cache file because, technically,
+        // FrameworkServicesTests don't have write permission to actually overwrite the cache file.
+        // We log and fail on the write silently in this case.
+        service.onStartJob(null);
+    }
+
     private void writeDataToFile(File f, String data) throws Exception{
         PrintStream out = new PrintStream(f);
         out.print(data);