Prune obsolete UsageStats data on upgrade.

When a database upgrade is performed for UsageStats, prune all
UsageStats data that belongs to packages that have been uninstalled.
This ensures that all data in UsageStats in R belongs to packages that
are currently installed or to those packages whose DONT_DELETE_DATA flag
was set when uninstalling.

Also remove the clean-up mapping step on boot. That was added as a
safety measure to ensure the mappings file is always updated. However,
with the addition of the prune job on package uninstalls and this CL,
that step is now unnecessary.

Bug: 143889121
Test: atest UsageStatsDatabase
Change-Id: Ib3d24dead4cd0e23145c15e7b1f88e2e20aadcaa

2 files changed, 95 insertions, 38 deletions

diff --git a/services/usage/java/com/android/server/usage/UsageStatsDatabase.java b/services/usage/java/com/android/server/usage/UsageStatsDatabase.java
index d938baddd421..ce29527d19f2 100644
--- a/services/usage/java/com/android/server/usage/UsageStatsDatabase.java
+++ b/services/usage/java/com/android/server/usage/UsageStatsDatabase.java
@@ -52,6 +52,7 @@ import java.io.OutputStream;
 import java.nio.file.Files;
 import java.nio.file.StandardCopyOption;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 
 /**
@@ -122,6 +123,7 @@ public class UsageStatsDatabase {
     private int mCurrentVersion;
     private boolean mFirstUpdate;
     private boolean mNewUpdate;
+    private boolean mUpgradePerformed;
 
     // The obfuscated packages to tokens mappings file
     private final File mPackageMappingsFile;
@@ -325,6 +327,13 @@ public class UsageStatsDatabase {
         return mNewUpdate;
     }
 
+    /**
+     * Was an upgrade performed when this database was initialized?
+     */
+    boolean wasUpgradePerformed() {
+        return mUpgradePerformed;
+    }
+
     private void checkVersionAndBuildLocked() {
         int version;
         String buildFingerprint;
@@ -397,6 +406,8 @@ public class UsageStatsDatabase {
         if (mUpdateBreadcrumb.exists()) {
             // Files should be up to date with current version. Clear the version update breadcrumb
             mUpdateBreadcrumb.delete();
+            // update mUpgradePerformed after breadcrumb is deleted to indicate a successful upgrade
+            mUpgradePerformed = true;
         }
 
         if (mBackupsDir.exists() && !KEEP_BACKUP_DIR) {
@@ -594,6 +605,70 @@ public class UsageStatsDatabase {
         }
     }
 
+    /**
+     * Iterates through all the files on disk and prunes any data that belongs to packages that have
+     * been uninstalled (packages that are not in the given list).
+     * Note: this should only be called once, when there has been a database upgrade.
+     *
+     * @param installedPackages map of installed packages (package_name:package_install_time)
+     */
+    void prunePackagesDataOnUpgrade(HashMap<String, Long> installedPackages) {
+        if (installedPackages == null || installedPackages.isEmpty()) {
+            return;
+        }
+        synchronized (mLock) {
+            for (int i = 0; i < mIntervalDirs.length; i++) {
+                final File[] files = mIntervalDirs[i].listFiles();
+                if (files == null) {
+                    continue;
+                }
+                for (int j = 0; j < files.length; j++) {
+                    try {
+                        final IntervalStats stats = new IntervalStats();
+                        final AtomicFile atomicFile = new AtomicFile(files[j]);
+                        readLocked(atomicFile, stats, mCurrentVersion, mPackagesTokenData);
+                        if (!pruneStats(installedPackages, stats)) {
+                            continue; // no stats were pruned so no need to rewrite
+                        }
+                        writeLocked(atomicFile, stats, mCurrentVersion, mPackagesTokenData);
+                    } catch (Exception e) {
+                        Slog.e(TAG, "Failed to prune data from: " + files[j].toString());
+                    }
+                }
+            }
+        }
+    }
+
+    private boolean pruneStats(HashMap<String, Long> installedPackages, IntervalStats stats) {
+        boolean dataPruned = false;
+
+        // prune old package usage stats
+        for (int i = stats.packageStats.size() - 1; i >= 0; i--) {
+            final UsageStats usageStats = stats.packageStats.valueAt(i);
+            final Long timeInstalled = installedPackages.get(usageStats.mPackageName);
+            if (timeInstalled == null || timeInstalled > usageStats.mEndTimeStamp) {
+                stats.packageStats.removeAt(i);
+                dataPruned = true;
+            }
+        }
+        if (dataPruned) {
+            // ensure old stats don't linger around during the obfuscation step on write
+            stats.packageStatsObfuscated.clear();
+        }
+
+        // prune old events
+        for (int i = stats.events.size() - 1; i >= 0; i--) {
+            final UsageEvents.Event event = stats.events.get(i);
+            final Long timeInstalled = installedPackages.get(event.mPackage);
+            if (timeInstalled == null || timeInstalled > event.mTimeStamp) {
+                stats.events.remove(i);
+                dataPruned = true;
+            }
+        }
+
+        return dataPruned;
+    }
+
     public void onTimeChanged(long timeDiffMillis) {
         synchronized (mLock) {
             StringBuilder logBuilder = new StringBuilder();
diff --git a/services/usage/java/com/android/server/usage/UserUsageStatsService.java b/services/usage/java/com/android/server/usage/UserUsageStatsService.java
index 20c0bd6748a6..d30f2ad33bf2 100644
--- a/services/usage/java/com/android/server/usage/UserUsageStatsService.java
+++ b/services/usage/java/com/android/server/usage/UserUsageStatsService.java
@@ -34,10 +34,9 @@ import android.app.usage.UsageEvents.Event;
 import android.app.usage.UsageStats;
 import android.app.usage.UsageStatsManager;
 import android.content.Context;
-import android.content.pm.ApplicationInfo;
-import android.content.pm.PackageManagerInternal;
+import android.content.pm.PackageInfo;
+import android.content.pm.PackageManager;
 import android.content.res.Configuration;
-import android.os.Process;
 import android.os.SystemClock;
 import android.text.format.DateUtils;
 import android.util.ArrayMap;
@@ -47,7 +46,6 @@ import android.util.Slog;
 import android.util.SparseIntArray;
 
 import com.android.internal.util.IndentingPrintWriter;
-import com.android.server.LocalServices;
 import com.android.server.usage.UsageStatsDatabase.StatCombiner;
 
 import java.io.File;
@@ -55,7 +53,7 @@ import java.io.IOException;
 import java.text.SimpleDateFormat;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.HashSet;
+import java.util.HashMap;
 import java.util.List;
 
 /**
@@ -115,6 +113,9 @@ class UserUsageStatsService {
     void init(final long currentTimeMillis) {
         readPackageMappingsLocked();
         mDatabase.init(currentTimeMillis);
+        if (mDatabase.wasUpgradePerformed()) {
+            mDatabase.prunePackagesDataOnUpgrade(getInstalledPackages());
+        }
 
         int nullCount = 0;
         for (int i = 0; i < mCurrentStats.length; i++) {
@@ -184,7 +185,6 @@ class UserUsageStatsService {
     private void readPackageMappingsLocked() {
         mDatabase.readMappingsLocked();
         updatePackageMappingsLocked();
-        cleanUpPackageMappingsLocked();
     }
 
     /**
@@ -216,42 +216,24 @@ class UserUsageStatsService {
     }
 
     /**
-     * Queries Package Manager for a list of installed packages and removes those packages from
-     * mPackagesTokenData which are not installed any more.
-     * This will only happen once per device boot, when the user is unlocked for the first time.
+     * Fetches a map of package names to their install times. This includes all installed packages,
+     * including those packages which have been uninstalled with the DONT_DELETE_DATA flag.
+     * Note: this is supposed be a helper method which is only used on database upgrades - it should
+     * not be called otherwise since it's implementation performs a heavy query to package manager.
      */
-    private void cleanUpPackageMappingsLocked() {
-        final long timeNow = System.currentTimeMillis();
-        /*
-         Note (b/142501248): PackageManagerInternal#getInstalledApplications is not lightweight.
-         Once its implementation is updated, or it's replaced with a better alternative, update
-         the call here to use it. For now, using the heavy #getInstalledApplications is okay since
-         this clean-up is only performed once every boot.
-         */
-        final PackageManagerInternal packageManagerInternal =
-                LocalServices.getService(PackageManagerInternal.class);
-        if (packageManagerInternal == null) {
-            return;
+    private HashMap<String, Long> getInstalledPackages() {
+        final PackageManager packageManager = mContext.getPackageManager();
+        if (packageManager == null) {
+            return null;
         }
-        final List<ApplicationInfo> installedPackages =
-                packageManagerInternal.getInstalledApplications(0, mUserId, Process.SYSTEM_UID);
-        // convert the package list to a set for easy look-ups
-        final HashSet<String> packagesSet = new HashSet<>(installedPackages.size());
+        final List<PackageInfo> installedPackages = packageManager.getInstalledPackagesAsUser(
+                PackageManager.MATCH_UNINSTALLED_PACKAGES, mUserId);
+        final HashMap<String, Long> packagesMap = new HashMap<>();
         for (int i = installedPackages.size() - 1; i >= 0; i--) {
-            packagesSet.add(installedPackages.get(i).packageName);
-        }
-        final List<String> removedPackages = new ArrayList<>();
-        // populate list of packages that are found in the mappings but not in the installed list
-        for (int i = mDatabase.mPackagesTokenData.packagesToTokensMap.size() - 1; i >= 0; i--) {
-            if (!packagesSet.contains(mDatabase.mPackagesTokenData.packagesToTokensMap.keyAt(i))) {
-                removedPackages.add(mDatabase.mPackagesTokenData.packagesToTokensMap.keyAt(i));
-            }
-        }
-
-        // remove packages in the mappings that are no longer installed
-        for (int i = removedPackages.size() - 1; i >= 0; i--) {
-            mDatabase.mPackagesTokenData.removePackage(removedPackages.get(i), timeNow);
+            final PackageInfo packageInfo = installedPackages.get(i);
+            packagesMap.put(packageInfo.packageName, packageInfo.firstInstallTime);
         }
+        return packagesMap;
     }
 
     boolean pruneUninstalledPackagesData() {