Merge changes from topic "b143889121"

* changes:
  Prune obsolete UsageStats data on upgrade.
  Add a job to prune UsageStats data on package removals.

core/res/AndroidManifest.xml

@@ -5045,6 +5045,10 @@
                  android:permission="android.permission.BIND_JOB_SERVICE" >
         </service>
 
+        <service android:name="com.android.server.usage.UsageStatsIdleService"
+                 android:permission="android.permission.BIND_JOB_SERVICE" >
+        </service>
+
         <service android:name="com.android.server.net.watchlist.ReportWatchlistJobService"
                  android:permission="android.permission.BIND_JOB_SERVICE" >
         </service>

services/core/java/android/app/usage/UsageStatsManagerInternal.java

@@ -281,4 +281,13 @@ public abstract class UsageStatsManagerInternal {
             return mUsageRemaining;
         }
     }
+
+    /**
+     * Called by {@link com.android.server.usage.UsageStatsIdleService} when the device is idle to
+     * prune usage stats data for uninstalled packages.
+     *
+     * @param userId the user associated with the job
+     * @return {@code true} if the pruning was successful, {@code false} otherwise
+     */
+    public abstract boolean pruneUninstalledPackagesData(@UserIdInt int userId);
 }

services/tests/servicestests/src/com/android/server/usage/UsageStatsDatabaseTest.java

@@ -45,6 +45,7 @@ import java.io.FileOutputStream;
 import java.io.IOException;
 import java.util.List;
 import java.util.Locale;
+import java.util.Set;
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
@@ -93,6 +94,8 @@ public class UsageStatsDatabaseTest {
                 for (File f : usageFiles) {
                     f.delete();
                 }
+            } else {
+                intervalDir.delete();
             }
         }
     }
@@ -587,6 +590,7 @@ public class UsageStatsDatabaseTest {
         db.readMappingsLocked();
         db.init(1);
         db.putUsageStats(interval, mIntervalStats);
+        db.writeMappingsLocked();
 
         final String removedPackage = "fake.package.name0";
         // invoke handler call directly from test to remove package
@@ -594,19 +598,19 @@ public class UsageStatsDatabaseTest {
 
         List<IntervalStats> stats = db.queryUsageStats(interval, 0, mEndTime,
                 mIntervalStatsVerifier);
-        for (int i = 0; i < stats.size(); i++) {
-            final IntervalStats stat = stats.get(i);
-            if (stat.packageStats.containsKey(removedPackage)) {
-                fail("Found removed package " + removedPackage + " in package stats.");
+        assertEquals(1, stats.size(),
+                "Only one interval stats object should exist for the given time range.");
+        final IntervalStats stat = stats.get(0);
+        if (stat.packageStats.containsKey(removedPackage)) {
+            fail("Found removed package " + removedPackage + " in package stats.");
+            return;
+        }
+        for (int i = 0; i < stat.events.size(); i++) {
+            final Event event = stat.events.get(i);
+            if (removedPackage.equals(event.mPackage)) {
+                fail("Found an event from removed package " + removedPackage);
                 return;
             }
-            for (int j = 0; j < stat.events.size(); j++) {
-                final Event event = stat.events.get(j);
-                if (removedPackage.equals(event.mPackage)) {
-                    fail("Found an event from removed package " + removedPackage);
-                    return;
-                }
-            }
         }
     }
 
@@ -617,4 +621,90 @@ public class UsageStatsDatabaseTest {
         verifyPackageNotRetained(UsageStatsManager.INTERVAL_MONTHLY);
         verifyPackageNotRetained(UsageStatsManager.INTERVAL_YEARLY);
     }
+
+    private void verifyPackageDataIsRemoved(UsageStatsDatabase db, int interval,
+            String removedPackage) {
+        List<IntervalStats> stats = db.queryUsageStats(interval, 0, mEndTime,
+                mIntervalStatsVerifier);
+        assertEquals(1, stats.size(),
+                "Only one interval stats object should exist for the given time range.");
+        final IntervalStats stat = stats.get(0);
+        if (stat.packageStats.containsKey(removedPackage)) {
+            fail("Found removed package " + removedPackage + " in package stats.");
+            return;
+        }
+        for (int i = 0; i < stat.events.size(); i++) {
+            final Event event = stat.events.get(i);
+            if (removedPackage.equals(event.mPackage)) {
+                fail("Found an event from removed package " + removedPackage);
+                return;
+            }
+        }
+    }
+
+    private void verifyPackageDataIsNotRemoved(UsageStatsDatabase db, int interval,
+            Set<String> installedPackages) {
+        List<IntervalStats> stats = db.queryUsageStats(interval, 0, mEndTime,
+                mIntervalStatsVerifier);
+        assertEquals(1, stats.size(),
+                "Only one interval stats object should exist for the given time range.");
+        final IntervalStats stat = stats.get(0);
+        if (!stat.packageStats.containsAll(installedPackages)) {
+            fail("Could not find some installed packages in package stats.");
+            return;
+        }
+        // attempt to find an event from each installed package
+        for (String installedPackage : installedPackages) {
+            for (int i = 0; i < stat.events.size(); i++) {
+                if (installedPackage.equals(stat.events.get(i).mPackage)) {
+                    break;
+                }
+                if (i == stat.events.size() - 1) {
+                    fail("Could not find any event for: " + installedPackage);
+                    return;
+                }
+            }
+        }
+    }
+
+    @Test
+    public void testPackageDataIsRemoved() throws IOException {
+        UsageStatsDatabase db = new UsageStatsDatabase(mTestDir);
+        db.readMappingsLocked();
+        db.init(1);
+
+        // write stats to disk for each interval
+        db.putUsageStats(UsageStatsManager.INTERVAL_DAILY, mIntervalStats);
+        db.putUsageStats(UsageStatsManager.INTERVAL_WEEKLY, mIntervalStats);
+        db.putUsageStats(UsageStatsManager.INTERVAL_MONTHLY, mIntervalStats);
+        db.putUsageStats(UsageStatsManager.INTERVAL_YEARLY, mIntervalStats);
+        db.writeMappingsLocked();
+
+        final Set<String> installedPackages = mIntervalStats.packageStats.keySet();
+        final String removedPackage = installedPackages.iterator().next();
+        installedPackages.remove(removedPackage);
+
+        // mimic a package uninstall
+        db.onPackageRemoved(removedPackage, System.currentTimeMillis());
+
+        // mimic the idle prune job being triggered
+        db.pruneUninstalledPackagesData();
+
+        // read data from disk into a new db instance
+        UsageStatsDatabase newDB = new UsageStatsDatabase(mTestDir);
+        newDB.readMappingsLocked();
+        newDB.init(mEndTime);
+
+        // query data for each interval and ensure data for package doesn't exist
+        verifyPackageDataIsRemoved(newDB, UsageStatsManager.INTERVAL_DAILY, removedPackage);
+        verifyPackageDataIsRemoved(newDB, UsageStatsManager.INTERVAL_WEEKLY, removedPackage);
+        verifyPackageDataIsRemoved(newDB, UsageStatsManager.INTERVAL_MONTHLY, removedPackage);
+        verifyPackageDataIsRemoved(newDB, UsageStatsManager.INTERVAL_YEARLY, removedPackage);
+
+        // query data for each interval and ensure some data for installed packages exists
+        verifyPackageDataIsNotRemoved(newDB, UsageStatsManager.INTERVAL_DAILY, installedPackages);
+        verifyPackageDataIsNotRemoved(newDB, UsageStatsManager.INTERVAL_WEEKLY, installedPackages);
+        verifyPackageDataIsNotRemoved(newDB, UsageStatsManager.INTERVAL_MONTHLY, installedPackages);
+        verifyPackageDataIsNotRemoved(newDB, UsageStatsManager.INTERVAL_YEARLY, installedPackages);
+    }
 }

services/usage/java/com/android/server/usage/IntervalStats.java

@@ -448,8 +448,11 @@ public class IntervalStats {
     /**
      * Parses all of the tokens to strings in the obfuscated usage stats data. This includes
      * deobfuscating each of the package tokens and chooser actions and categories.
+     *
+     * @return {@code true} if any stats were omitted while deobfuscating, {@code false} otherwise.
      */
-    private void deobfuscateUsageStats(PackagesTokenData packagesTokenData) {
+    private boolean deobfuscateUsageStats(PackagesTokenData packagesTokenData) {
+        boolean dataOmitted = false;
         final int usageStatsSize = packageStatsObfuscated.size();
         for (int statsIndex = 0; statsIndex < usageStatsSize; statsIndex++) {
             final int packageToken = packageStatsObfuscated.keyAt(statsIndex);
@@ -457,6 +460,7 @@ public class IntervalStats {
             usageStats.mPackageName = packagesTokenData.getPackageString(packageToken);
             if (usageStats.mPackageName == null) {
                 Slog.e(TAG, "Unable to parse usage stats package " + packageToken);
+                dataOmitted = true;
                 continue;
             }
 
@@ -489,14 +493,18 @@ public class IntervalStats {
             }
             packageStats.put(usageStats.mPackageName, usageStats);
         }
+        return dataOmitted;
     }
 
     /**
      * Parses all of the tokens to strings in the obfuscated events data. This includes
      * deobfuscating the package token, along with any class, task root package/class tokens, and
      * shortcut or notification channel tokens.
+     *
+     * @return {@code true} if any events were omitted while deobfuscating, {@code false} otherwise.
      */
-    private void deobfuscateEvents(PackagesTokenData packagesTokenData) {
+    private boolean deobfuscateEvents(PackagesTokenData packagesTokenData) {
+        boolean dataOmitted = false;
         for (int i = this.events.size() - 1; i >= 0; i--) {
             final Event event = this.events.get(i);
             final int packageToken = event.mPackageToken;
@@ -504,6 +512,7 @@ public class IntervalStats {
             if (event.mPackage == null) {
                 Slog.e(TAG, "Unable to parse event package " + packageToken);
                 this.events.remove(i);
+                dataOmitted = true;
                 continue;
             }
 
@@ -543,6 +552,7 @@ public class IntervalStats {
                         Slog.e(TAG, "Unable to parse shortcut " + event.mShortcutIdToken
                                 + " for package " + packageToken);
                         this.events.remove(i);
+                        dataOmitted = true;
                         continue;
                     }
                     break;
@@ -554,21 +564,25 @@ public class IntervalStats {
                                 + event.mNotificationChannelIdToken + " for package "
                                 + packageToken);
                         this.events.remove(i);
+                        dataOmitted = true;
                         continue;
                     }
                     break;
             }
         }
+        return dataOmitted;
     }
 
     /**
      * Parses the obfuscated tokenized data held in this interval stats object.
      *
+     * @return {@code true} if any data was omitted while deobfuscating, {@code false} otherwise.
      * @hide
      */
-    public void deobfuscateData(PackagesTokenData packagesTokenData) {
-        deobfuscateUsageStats(packagesTokenData);
-        deobfuscateEvents(packagesTokenData);
+    public boolean deobfuscateData(PackagesTokenData packagesTokenData) {
+        final boolean statsOmitted = deobfuscateUsageStats(packagesTokenData);
+        final boolean eventsOmitted = deobfuscateEvents(packagesTokenData);
+        return statsOmitted || eventsOmitted;
     }
 
     /**

services/usage/java/com/android/server/usage/PackagesTokenData.java

@@ -162,15 +162,18 @@ public final class PackagesTokenData {
      *
      * @param packageName the package to be removed
      * @param timeRemoved the time stamp of when the package was removed
+     * @return the token mapped to the package removed or {@code PackagesTokenData.UNASSIGNED_TOKEN}
+     *         if not mapped
      */
-    public void removePackage(String packageName, long timeRemoved) {
+    public int removePackage(String packageName, long timeRemoved) {
         removedPackagesMap.put(packageName, timeRemoved);
 
         if (!packagesToTokensMap.containsKey(packageName)) {
-            return;
+            return UNASSIGNED_TOKEN;
         }
         final int packageToken = packagesToTokensMap.get(packageName).get(packageName);
         packagesToTokensMap.remove(packageName);
         tokensToPackagesMap.delete(packageToken);
+        return packageToken;
     }
 }

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) {
@@ -545,12 +556,119 @@ public class UsageStatsDatabase {
         }
     }
 
-    void onPackageRemoved(String packageName, long timeRemoved) {
+    /**
+     * Returns the token mapped to the package removed or {@code PackagesTokenData.UNASSIGNED_TOKEN}
+     * if not mapped.
+     */
+    int onPackageRemoved(String packageName, long timeRemoved) {
         synchronized (mLock) {
-            mPackagesTokenData.removePackage(packageName, timeRemoved);
+            return mPackagesTokenData.removePackage(packageName, timeRemoved);
         }
     }
 
+    /**
+     * Reads all the usage stats data on disk and rewrites it with any data related to uninstalled
+     * packages omitted.
+     */
+    boolean pruneUninstalledPackagesData() {
+        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]);
+                        if (!readLocked(atomicFile, stats, mCurrentVersion, mPackagesTokenData)) {
+                            continue; // no data was omitted when read so no need to rewrite
+                        }
+                        // Any data related to packages that have been removed would have failed
+                        // the deobfuscation step on read so the IntervalStats object here only
+                        // contains data for packages that are currently installed - all we need
+                        // to do here is write the data back to disk.
+                        writeLocked(atomicFile, stats, mCurrentVersion, mPackagesTokenData);
+                    } catch (Exception e) {
+                        Slog.e(TAG, "Failed to prune data from: " + files[j].toString());
+                        return false;
+                    }
+                }
+            }
+            try {
+                writeMappingsLocked();
+            } catch (IOException e) {
+                Slog.e(TAG, "Failed to write package mappings after pruning data.");
+                return false;
+            }
+            return true;
+        }
+    }
+
+    /**
+     * 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();
@@ -645,7 +763,6 @@ public class UsageStatsDatabase {
         }
 
         // filter out events
-        final int eventsSize = stats.events.size();
         for (int i = stats.events.size() - 1; i >= 0; i--) {
             final UsageEvents.Event event = stats.events.get(i);
             final Long timeRemoved = removedPackagesMap.get(event.mPackage);
@@ -942,13 +1059,17 @@ public class UsageStatsDatabase {
         readLocked(file, statsOut, mCurrentVersion, mPackagesTokenData);
     }
 
-    private static void readLocked(AtomicFile file, IntervalStats statsOut, int version,
+    /**
+     * Returns {@code true} if any stats were omitted while reading, {@code false} otherwise.
+     */
+    private static boolean readLocked(AtomicFile file, IntervalStats statsOut, int version,
             PackagesTokenData packagesTokenData) throws IOException {
+        boolean dataOmitted = false;
         try {
             FileInputStream in = file.openRead();
             try {
                 statsOut.beginTime = parseBeginTime(file);
-                readLocked(in, statsOut, version, packagesTokenData);
+                dataOmitted = readLocked(in, statsOut, version, packagesTokenData);
                 statsOut.lastTimeSaved = file.getLastModifiedTime();
             } finally {
                 try {
@@ -961,10 +1082,15 @@ public class UsageStatsDatabase {
             Slog.e(TAG, "UsageStatsDatabase", e);
             throw e;
         }
+        return dataOmitted;
     }
 
-    private static void readLocked(InputStream in, IntervalStats statsOut, int version,
+    /**
+     * Returns {@code true} if any stats were omitted while reading, {@code false} otherwise.
+     */
+    private static boolean readLocked(InputStream in, IntervalStats statsOut, int version,
             PackagesTokenData packagesTokenData) throws IOException {
+        boolean dataOmitted = false;
         switch (version) {
             case 1:
             case 2:
@@ -989,14 +1115,14 @@ public class UsageStatsDatabase {
                 } catch (IOException e) {
                     Slog.e(TAG, "Unable to read interval stats from proto.", e);
                 }
-                statsOut.deobfuscateData(packagesTokenData);
+                dataOmitted = statsOut.deobfuscateData(packagesTokenData);
                 break;
             default:
                 throw new RuntimeException(
                         "Unhandled UsageStatsDatabase version: " + Integer.toString(version)
                                 + " on read.");
         }
-
+        return dataOmitted;
     }
 
     /**

services/usage/java/com/android/server/usage/UsageStatsIdleService.java

@@ -0,0 +1,121 @@
+/*
+ * Copyright (C) 2019 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.app.job.JobInfo;
+import android.app.job.JobParameters;
+import android.app.job.JobScheduler;
+import android.app.job.JobService;
+import android.app.usage.UsageStatsManagerInternal;
+import android.content.ComponentName;
+import android.content.Context;
+import android.os.AsyncTask;
+import android.os.PersistableBundle;
+
+import com.android.server.LocalServices;
+
+/**
+ * JobService used to do any work for UsageStats while the device is idle.
+ */
+public class UsageStatsIdleService extends JobService {
+
+    /**
+     * Base job ID for the pruning job - must be unique within the system server uid.
+     */
+    private static final int PRUNE_JOB_ID = 546357475;
+
+    private static final String USER_ID_KEY = "user_id";
+    private static final String REMOVED_TOKENS_KEY = "removed_tokens";
+
+    static void scheduleJob(Context context, int userId, int token) {
+        final int userJobId = PRUNE_JOB_ID + userId; // unique job id per user
+        final ComponentName component = new ComponentName(context.getPackageName(),
+                UsageStatsIdleService.class.getName());
+        final PersistableBundle bundle = new PersistableBundle();
+        bundle.putInt(USER_ID_KEY, userId);
+        bundle.putIntArray(REMOVED_TOKENS_KEY,
+                getOrCreateRemovedTokens(fetchRemovedTokens(context, userId), token));
+        final JobInfo pruneJob = new JobInfo.Builder(userJobId, component)
+                .setRequiresDeviceIdle(true)
+                .setExtras(bundle)
+                .setPersisted(true)
+                .build();
+
+        final JobScheduler jobScheduler = context.getSystemService(JobScheduler.class);
+        jobScheduler.schedule(pruneJob);
+    }
+
+    static void cancelJob(Context context, int userId) {
+        final int userJobId = PRUNE_JOB_ID + userId; // unique job id per user
+        final JobScheduler jobScheduler = context.getSystemService(JobScheduler.class);
+        jobScheduler.cancel(userJobId);
+    }
+
+    /**
+     * Fetches an array of removed tokens from previous prune jobs, if any.
+     */
+    static int[] fetchRemovedTokens(Context context, int userId) {
+        final JobScheduler jobScheduler = context.getSystemService(JobScheduler.class);
+        final int userJobId = PRUNE_JOB_ID + userId; // unique job id per user
+        final JobInfo pendingJob = jobScheduler.getPendingJob(userJobId);
+        if (pendingJob != null) {
+            final PersistableBundle bundle = pendingJob.getExtras();
+            return bundle.getIntArray(REMOVED_TOKENS_KEY);
+        }
+        return null;
+    }
+
+    @Override
+    public boolean onStartJob(JobParameters params) {
+        final PersistableBundle bundle = params.getExtras();
+        final int userId = bundle.getInt(USER_ID_KEY, -1);
+        if (userId == -1) {
+            return false;
+        }
+
+        AsyncTask.execute(() -> {
+            final UsageStatsManagerInternal usageStatsManagerInternal = LocalServices.getService(
+                    UsageStatsManagerInternal.class);
+            final boolean pruned = usageStatsManagerInternal.pruneUninstalledPackagesData(userId);
+            jobFinished(params, !pruned); // reschedule if data was not pruned
+        });
+        return true;
+    }
+
+    @Override
+    public boolean onStopJob(JobParameters params) {
+        // Since the pruning job isn't a heavy job, we don't want to cancel it's execution midway.
+        return false;
+    }
+
+    /**
+     * Helper method to create a cloned array of removed tokens from previous jobs (if any) with
+     * the newly removed token at index 0. If there are no removed tokens from previous jobs, it
+     * simply returns a new array containing the current token.
+     */
+    private static int[] getOrCreateRemovedTokens(int[] previousRemovedTokens, int token) {
+        final int[] removedTokens;
+        if (previousRemovedTokens == null) {
+            removedTokens = new int[1];
+        } else {
+            removedTokens = new int[previousRemovedTokens.length + 1];
+            System.arraycopy(previousRemovedTokens, 0, removedTokens, 1,
+                    previousRemovedTokens.length);
+        }
+        removedTokens[0] = token;
+        return removedTokens;
+    }
+}

services/usage/java/com/android/server/usage/UsageStatsService.java

@@ -846,6 +846,8 @@ public class UsageStatsService extends SystemService implements
             mUserState.remove(userId);
             mAppStandby.onUserRemoved(userId);
             mAppTimeLimit.onUserRemoved(userId);
+            // Cancel any scheduled jobs for this user since the user is being removed.
+            UsageStatsIdleService.cancelJob(getContext(), userId);
         }
     }
 
@@ -855,7 +857,7 @@ public class UsageStatsService extends SystemService implements
     private void onPackageRemoved(int userId, String packageName) {
         synchronized (mLock) {
             final long timeRemoved = System.currentTimeMillis();
-            if (!mUserUnlockedStates.get(userId, false)) {
+            if (!mUserUnlockedStates.get(userId)) {
                 // If user is not unlocked and a package is removed for them, we will handle it
                 // when the user service is initialized and package manager is queried.
                 return;
@@ -869,6 +871,24 @@ public class UsageStatsService extends SystemService implements
         }
     }
 
+    /**
+     * Called by the Binder stub.
+     */
+    private boolean pruneUninstalledPackagesData(int userId) {
+        synchronized (mLock) {
+            if (!mUserUnlockedStates.get(userId)) {
+                return false; // user is no longer unlocked
+            }
+
+            final UserUsageStatsService userService = mUserState.get(userId);
+            if (userService == null) {
+                return false; // user was stopped or removed
+            }
+
+            return userService.pruneUninstalledPackagesData();
+        }
+    }
+
     /**
      * Called by the Binder stub.
      */
@@ -2124,6 +2144,11 @@ public class UsageStatsService extends SystemService implements
         public AppUsageLimitData getAppUsageLimit(String packageName, UserHandle user) {
             return mAppTimeLimit.getAppUsageLimit(packageName, user);
         }
+
+        @Override
+        public boolean pruneUninstalledPackagesData(int userId) {
+            return UsageStatsService.this.pruneUninstalledPackagesData(userId);
+        }
     }
 
     private class MyPackageMonitor extends PackageMonitor {

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++) {
@@ -175,53 +176,70 @@ class UserUsageStatsService {
     }
 
     void onPackageRemoved(String packageName, long timeRemoved) {
-        mDatabase.onPackageRemoved(packageName, timeRemoved);
+        final int token = mDatabase.onPackageRemoved(packageName, timeRemoved);
+        if (token != PackagesTokenData.UNASSIGNED_TOKEN) {
+            UsageStatsIdleService.scheduleJob(mContext, mUserId, token);
+        }
     }
 
     private void readPackageMappingsLocked() {
         mDatabase.readMappingsLocked();
-        cleanUpPackageMappingsLocked();
+        updatePackageMappingsLocked();
     }
 
     /**
-     * Queries Package Manager for a list of installed packages and removes those packages from
-     * mPackagesTokenData which are not installed any more.
+     * Queries Job Scheduler for any pending data prune jobs and if any exist, it updates the
+     * package mappings in memory by removing those tokens.
      * This will only happen once per device boot, when the user is unlocked for the first time.
      */
-    private void cleanUpPackageMappingsLocked() {
+    private void updatePackageMappingsLocked() {
         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) {
+        final int[] removedTokens = UsageStatsIdleService.fetchRemovedTokens(mContext, mUserId);
+        if (removedTokens == null) {
             return;
         }
-        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());
-        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);
+        for (int i = removedTokens.length - 1; i >= 0; i--) {
+            final String packageName =
+                    mDatabase.mPackagesTokenData.getPackageString(removedTokens[i]);
+            if (packageName == null) {
+                continue;
+            }
+            /*
+             Note: in most cases, packageName returned will be null since the package mappings file
+             on disk should have been updated when it was last persisted. This is to handle the rare
+             case of system crashing after a package was removed but the package mappings file
+             wasn't persisted to disk.
+             */
+            mDatabase.mPackagesTokenData.removePackage(packageName, timeNow);
         }
     }
 
+    /**
+     * 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 HashMap<String, Long> getInstalledPackages() {
+        final PackageManager packageManager = mContext.getPackageManager();
+        if (packageManager == null) {
+            return null;
+        }
+        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--) {
+            final PackageInfo packageInfo = installedPackages.get(i);
+            packagesMap.put(packageInfo.packageName, packageInfo.firstInstallTime);
+        }
+        return packagesMap;
+    }
+
+    boolean pruneUninstalledPackagesData() {
+        return mDatabase.pruneUninstalledPackagesData();
+    }
+
     private void onTimeChanged(long oldTime, long newTime) {
         persistActiveStats();
         mDatabase.onTimeChanged(newTime - oldTime);